Skip to content

Commit 200f70c

Browse files
kapv89kapv89
authored
Merge pull request #192 from sjs994/commenter-sql
Update go sql to as a driver
2 parents f04f24a + 46920a4 commit 200f70c

File tree

8 files changed

+503
-233
lines changed

8 files changed

+503
-233
lines changed

‎go/database/sql/connection.go‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
// Copyright 2022 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package sql
16+
17+
import (
18+
"context"
19+
"database/sql/driver"
20+
"fmt"
21+
"runtime/debug"
22+
"strings"
23+
24+
"github.com/google/sqlcommenter/go/core"
25+
)
26+
27+
var attemptedToAutosetApplication = false
28+
29+
type sqlCommenterConn struct {
30+
driver.Conn
31+
options core.CommenterOptions
32+
}
33+
34+
func newSQLCommenterConn(conn driver.Conn, options core.CommenterOptions) *sqlCommenterConn {
35+
return &sqlCommenterConn{
36+
Conn: conn,
37+
options: options,
38+
}
39+
}
40+
41+
func (s *sqlCommenterConn) Query(query string, args []driver.Value) (driver.Rows, error) {
42+
queryer, ok := s.Conn.(driver.Queryer)
43+
if !ok {
44+
return nil, driver.ErrSkip
45+
}
46+
ctx := context.Background()
47+
commentedQuery := s.withComment(ctx, query)
48+
return queryer.Query(commentedQuery, args)
49+
}
50+
51+
func (s *sqlCommenterConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
52+
queryer, ok := s.Conn.(driver.QueryerContext)
53+
if !ok {
54+
return nil, driver.ErrSkip
55+
}
56+
commentedQuery := s.withComment(ctx, query)
57+
return queryer.QueryContext(ctx, commentedQuery, args)
58+
}
59+
60+
func (s *sqlCommenterConn) Exec(query string, args []driver.Value) (driver.Result, error) {
61+
execor, ok := s.Conn.(driver.Execer)
62+
if !ok {
63+
return nil, driver.ErrSkip
64+
}
65+
ctx := context.Background()
66+
commentedQuery := s.withComment(ctx, query)
67+
return execor.Exec(commentedQuery, args)
68+
}
69+
70+
func (s *sqlCommenterConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
71+
execor, ok := s.Conn.(driver.ExecerContext)
72+
if !ok {
73+
return nil, driver.ErrSkip
74+
}
75+
commentedQuery := s.withComment(ctx, query)
76+
return execor.ExecContext(ctx, commentedQuery, args)
77+
}
78+
79+
func (s *sqlCommenterConn) PrepareContext(ctx context.Context, query string) (stmt driver.Stmt, err error) {
80+
preparer, ok := s.Conn.(driver.ConnPrepareContext)
81+
if !ok {
82+
return nil, driver.ErrSkip
83+
}
84+
commentedQuery := s.withComment(ctx, query)
85+
return preparer.PrepareContext(ctx, commentedQuery)
86+
}
87+
88+
func (s *sqlCommenterConn) Raw() driver.Conn {
89+
return s.Conn
90+
}
91+
92+
// ***** Commenter Functions *****
93+
94+
func (conn *sqlCommenterConn) withComment(ctx context.Context, query string) string {
95+
var commentsMap = map[string]string{}
96+
query = strings.TrimSpace(query)
97+
config := conn.options.Config
98+
99+
// Sorted alphabetically
100+
if config.EnableAction && (ctx.Value(core.Action) != nil) {
101+
commentsMap[core.Action] = ctx.Value(core.Action).(string)
102+
}
103+
104+
// `driver` information should not be coming from framework.
105+
// So, explicitly adding that here.
106+
if config.EnableDBDriver {
107+
commentsMap[core.Driver] = fmt.Sprintf("database/sql:%s", conn.options.Tags.DriverName)
108+
}
109+
110+
if config.EnableFramework && (ctx.Value(core.Framework) != nil) {
111+
commentsMap[core.Framework] = ctx.Value(core.Framework).(string)
112+
}
113+
114+
if config.EnableRoute && (ctx.Value(core.Route) != nil) {
115+
commentsMap[core.Route] = ctx.Value(core.Route).(string)
116+
}
117+
118+
if config.EnableTraceparent {
119+
carrier := core.ExtractTraceparent(ctx)
120+
if val, ok := carrier["traceparent"]; ok {
121+
commentsMap[core.Traceparent] = val
122+
}
123+
}
124+
125+
if config.EnableApplication {
126+
if !attemptedToAutosetApplication && conn.options.Tags.Application == "" {
127+
attemptedToAutosetApplication = true
128+
bi, ok := debug.ReadBuildInfo()
129+
if ok {
130+
conn.options.Tags.Application = bi.Path
131+
}
132+
}
133+
if conn.options.Tags.Application != "" {
134+
commentsMap[core.Application] = conn.options.Tags.Application
135+
}
136+
}
137+
138+
var commentsString string = ""
139+
if len(commentsMap) > 0 { // Converts comments map to string and appends it to query
140+
commentsString = fmt.Sprintf("/*%s*/", core.ConvertMapToComment(commentsMap))
141+
}
142+
143+
// A semicolon at the end of the SQL statement means the query ends there.
144+
// We need to insert the comment before that to be considered as part of the SQL statemtent.
145+
if query[len(query)-1:] == ";" {
146+
return fmt.Sprintf("%s%s;", strings.TrimSuffix(query, ";"), commentsString)
147+
}
148+
return fmt.Sprintf("%s%s", query, commentsString)
149+
}
150+
151+
// ***** Commenter Functions *****

‎go/database/sql/connection_test.go‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
// Copyright 2022 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package sql
16+
17+
import (
18+
"context"
19+
"testing"
20+
21+
"github.com/google/sqlcommenter/go/core"
22+
)
23+
24+
func TestWithComment_NoContext(t *testing.T) {
25+
testBasicConn := &mockConn{}
26+
testCases := []struct {
27+
desc string
28+
commenterOptions core.CommenterOptions
29+
query string
30+
wantQuery string
31+
}{
32+
{
33+
desc: "empty commenter options",
34+
commenterOptions: core.CommenterOptions{},
35+
query: "SELECT 1;",
36+
wantQuery: "SELECT 1;",
37+
},
38+
{
39+
desc: "only enable DBDriver",
40+
commenterOptions: core.CommenterOptions{
41+
Config: core.CommenterConfig{EnableDBDriver: true},
42+
},
43+
query: "SELECT 1;",
44+
wantQuery: "SELECT 1/*db_driver=database%2Fsql%3A*/;",
45+
},
46+
{
47+
desc: "enable DBDriver and pass static tag driver name",
48+
commenterOptions: core.CommenterOptions{
49+
Config: core.CommenterConfig{EnableDBDriver: true},
50+
Tags: core.StaticTags{DriverName: "postgres"},
51+
},
52+
query: "SELECT 1;",
53+
wantQuery: "SELECT 1/*db_driver=database%2Fsql%3Apostgres*/;",
54+
},
55+
{
56+
desc: "enable DBDriver and pass all static tags",
57+
commenterOptions: core.CommenterOptions{
58+
Config: core.CommenterConfig{EnableDBDriver: true},
59+
Tags: core.StaticTags{DriverName: "postgres", Application: "app-1"},
60+
},
61+
query: "SELECT 1;",
62+
wantQuery: "SELECT 1/*db_driver=database%2Fsql%3Apostgres*/;",
63+
},
64+
{
65+
desc: "enable other tags and pass all static tags",
66+
commenterOptions: core.CommenterOptions{
67+
Config: core.CommenterConfig{EnableDBDriver: true, EnableApplication: true, EnableFramework: true},
68+
Tags: core.StaticTags{DriverName: "postgres", Application: "app-1"},
69+
},
70+
query: "SELECT 1;",
71+
wantQuery: "SELECT 1/*application=app-1,db_driver=database%2Fsql%3Apostgres*/;",
72+
},
73+
}
74+
for _, tc := range testCases {
75+
testConn := newSQLCommenterConn(testBasicConn, tc.commenterOptions)
76+
ctx := context.Background()
77+
if got, want := testConn.withComment(ctx, tc.query), tc.wantQuery; got != want {
78+
t.Errorf("testConn.withComment(ctx, %q) = %q, want = %q", tc.query, got, want)
79+
}
80+
}
81+
}
82+
83+
func TestWithComment_WithContext(t *testing.T) {
84+
testBasicConn := &mockConn{}
85+
testCases := []struct {
86+
desc string
87+
commenterOptions core.CommenterOptions
88+
ctx context.Context
89+
query string
90+
wantQuery string
91+
}{
92+
{
93+
desc: "empty commenter options",
94+
commenterOptions: core.CommenterOptions{},
95+
ctx: getContextWithKeyValue(
96+
map[string]string{
97+
"route": "listData",
98+
"framework": "custom-golang",
99+
},
100+
),
101+
query: "SELECT 1;",
102+
wantQuery: "SELECT 1;",
103+
},
104+
{
105+
desc: "only all options but context has few tags",
106+
commenterOptions: core.CommenterOptions{
107+
Config: core.CommenterConfig{
108+
EnableDBDriver: true,
109+
EnableRoute: true,
110+
EnableFramework: true,
111+
EnableController: true,
112+
EnableAction: true,
113+
EnableTraceparent: true,
114+
EnableApplication: true,
115+
},
116+
Tags: core.StaticTags{DriverName: "postgres", Application: "app-1"},
117+
},
118+
ctx: getContextWithKeyValue(
119+
map[string]string{
120+
"route": "listData",
121+
"framework": "custom-golang",
122+
},
123+
),
124+
query: "SELECT 1;",
125+
wantQuery: "SELECT 1/*application=app-1,db_driver=database%2Fsql%3Apostgres,framework=custom-golang,route=listData*/;",
126+
},
127+
{
128+
desc: "only all options but context contains all tags",
129+
commenterOptions: core.CommenterOptions{
130+
Config: core.CommenterConfig{
131+
EnableDBDriver: true,
132+
EnableRoute: true,
133+
EnableFramework: true,
134+
EnableController: true,
135+
EnableAction: true,
136+
EnableTraceparent: true,
137+
EnableApplication: true,
138+
},
139+
Tags: core.StaticTags{DriverName: "postgres", Application: "app-1"},
140+
},
141+
ctx: getContextWithKeyValue(
142+
map[string]string{
143+
"route": "listData",
144+
"framework": "custom-golang",
145+
"controller": "custom-controller",
146+
"action": "any action",
147+
},
148+
),
149+
query: "SELECT 1;",
150+
wantQuery: "SELECT 1/*action=any+action,application=app-1,db_driver=database%2Fsql%3Apostgres,framework=custom-golang,route=listData*/;",
151+
},
152+
}
153+
for _, tc := range testCases {
154+
testConn := newSQLCommenterConn(testBasicConn, tc.commenterOptions)
155+
if got, want := testConn.withComment(tc.ctx, tc.query), tc.wantQuery; got != want {
156+
t.Errorf("testConn.withComment(ctx, %q) = %q, want = %q", tc.query, got, want)
157+
}
158+
}
159+
}
160+
161+
func getContextWithKeyValue(vals map[string]string) context.Context {
162+
ctx := context.Background()
163+
for k, v := range vals {
164+
ctx = context.WithValue(ctx, k, v)
165+
}
166+
return ctx
167+
}

0 commit comments

Comments
 (0)