Initial commit.

Author: jbub

.drone.yml

@@ -0,0 +1,12 @@
+---
+kind: pipeline
+name: default
+
+clone:
+  depth: 50
+
+steps:
+  - name: build
+    image: golang:1.17
+    commands:
+      - go test -v ./...

.gitignore

@@ -0,0 +1,35 @@
+# os
+.DS_Store
+Thumbs.db
+Desktop.ini
+
+# Compiled Object files, Static and Dynamic libs (Shared Objects)
+*.o
+*.a
+*.so
+
+# Folders
+_obj
+_test
+
+# Architecture specific extensions/prefixes
+*.[568vq]
+[568vq].out
+
+*.cgo1.go
+*.cgo2.c
+_cgo_defun.c
+_cgo_gotypes.go
+_cgo_export.*
+
+_testmain.go
+
+*.exe
+*.test
+*.prof
+
+# coverage
+*.out
+
+# goland
+.idea/

CHANGELOG.md

@@ -0,0 +1,3 @@
+## v0.0.1
+
+- Initial release.

LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2021 Juraj Bubniak
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,60 @@
+# sqlcommenter
+[![GoDoc](https://img.shields.io/badge/go-documentation-blue.svg?style=flat-square)](https://godoc.org/github.com/jbub/sqlcommenter)
+[![Build Status](https://cloud.drone.io/api/badges/jbub/sqlcommenter/status.svg)](https://cloud.drone.io/jbub/sqlcommenter)
+[![Go Report Card](https://goreportcard.com/badge/github.com/jbub/sqlcommenter)](https://goreportcard.com/report/github.com/jbub/sqlcommenter)
+
+Go implementation of https://google.github.io/sqlcommenter/.
+
+## Usage with pgx stdlib driver
+
+```go
+package main
+
+import (
+    "context"
+    "database/sql"
+
+    "github.com/jackc/pgx/v4/stdlib"
+    "github.com/jbub/sqlcommenter"
+)
+
+type contextKey int
+
+const contextKeyUserID contextKey = 0
+
+func withUserID(ctx context.Context, key string) context.Context {
+    return context.WithValue(ctx, contextKeyUserID, key)
+}
+
+func userIDFromContext(ctx context.Context) string {
+    return ctx.Value(contextKeyUserID).(string)
+}
+
+func main() {
+    pgxDrv := stdlib.GetDefaultDriver()
+    drv := sqlcommenter.WrapDriver(pgxDrv,
+        sqlcommenter.WithAttrPairs("application", "hello-app"),
+        sqlcommenter.WithAttrFunc(func(ctx context.Context) sqlcommenter.Attrs {
+            return sqlcommenter.AttrPairs("user-id", userIDFromContext(ctx))
+        }),
+    )
+
+    sql.Register("pgx-sqlcommenter", drv)
+
+    db, err := sql.Open("pgx-sqlcommenter", "postgres://user@host:5432/db")
+    if err != nil {
+        // handle error
+    }
+    defer db.Close()
+    
+    ctx := context.Background()
+
+    rows, err := db.QueryContext(withUserID(ctx, "22"), "SELECT 1")
+    if err != nil {
+        // handle error
+    }
+    defer rows.Close()
+    
+    // will produce the following query: SELECT 1 /* application='hello-app',user-id='22' */
+}
+```

attrs.go

@@ -0,0 +1,59 @@
+package sqlcommenter
+
+import (
+	"net/url"
+	"sort"
+	"strings"
+)
+
+func AttrPairs(pairs ...string) Attrs {
+	if len(pairs)%2 == 1 {
+		panic("got odd number of pairs")
+	}
+	attrs := make(Attrs, len(pairs)/2)
+	for i := 0; i < len(pairs); i += 2 {
+		attrs[pairs[i]] = pairs[i+1]
+	}
+	return attrs
+}
+
+type Attr struct {
+	Key   string
+	Value string
+}
+
+type Attrs map[string]string
+
+func (a Attrs) Encode() string {
+	total := len(a)
+	keys := make([]string, 0, total)
+	for k := range a {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+
+	var b strings.Builder
+	for i, key := range keys {
+		b.WriteString(encodeKey(key))
+		b.WriteByte('=')
+		b.WriteString(encodeValue(a[key]))
+		if i < total-1 {
+			b.WriteByte(',')
+		}
+	}
+	return b.String()
+}
+
+func (a Attrs) Update(other Attrs) {
+	for k, v := range other {
+		a[k] = v
+	}
+}
+
+func encodeKey(k string) string {
+	return url.QueryEscape(k)
+}
+
+func encodeValue(v string) string {
+	return "'" + strings.ReplaceAll(url.PathEscape(v), "+", "%20") + "'"
+}

attrs_test.go

@@ -0,0 +1,42 @@
+package sqlcommenter
+
+import (
+	"testing"
+)
+
+func TestAttrsEncode(t *testing.T) {
+	cases := []struct {
+		name  string
+		attrs Attrs
+		want  string
+	}{
+		{
+			name: "no attrs",
+		},
+		{
+			name: "single attr",
+			attrs: map[string]string{
+				"key": "value",
+			},
+			want: "key='value'",
+		},
+		{
+			name: "multiple attrs",
+			attrs: map[string]string{
+				"key":  "value",
+				"2key": "value 33",
+				"key3": "44  value",
+			},
+			want: "2key='value%2033',key='value',key3='44%20%20value'",
+		},
+	}
+
+	for _, cs := range cases {
+		t.Run(cs.name, func(t *testing.T) {
+			got := cs.attrs.Encode()
+			if want := cs.want; want != got {
+				t.Errorf("got '%v', want '%v'", got, want)
+			}
+		})
+	}
+}

comment.go

@@ -0,0 +1,52 @@
+package sqlcommenter
+
+import (
+	"context"
+	"strings"
+)
+
+const (
+	commentStart = "/*"
+	commentEnd   = "*/"
+)
+
+func Comment(query string, opts ...Option) string {
+	if len(opts) == 0 {
+		return query
+	}
+	if strings.Contains(query, commentStart) {
+		return query
+	}
+	return newCommenter(opts...).comment(context.Background(), query)
+}
+
+func newCommenter(opts ...Option) *commenter {
+	cmt := &commenter{}
+	for _, opt := range opts {
+		opt(cmt)
+	}
+	return cmt
+}
+
+type commenter struct {
+	providers []AttrProvider
+}
+
+func (c *commenter) comment(ctx context.Context, query string) string {
+	attrs := c.attrs(ctx)
+	if len(attrs) == 0 {
+		return query
+	}
+	return query + " " + commentStart + " " + attrs.Encode() + " " + commentEnd
+}
+
+func (c *commenter) attrs(ctx context.Context) Attrs {
+	if len(c.providers) == 0 {
+		return nil
+	}
+	attrs := make(Attrs)
+	for _, prov := range c.providers {
+		attrs.Update(prov.GetAttrs(ctx))
+	}
+	return attrs
+}

comment_test.go

@@ -0,0 +1,54 @@
+package sqlcommenter
+
+import (
+	"testing"
+)
+
+func TestComment(t *testing.T) {
+	cases := []struct {
+		name  string
+		query string
+		opts  []Option
+		want  string
+	}{
+		{
+			name: "empty query",
+		},
+		{
+			name:  "empty query with whitespace",
+			query: "  ",
+			want:  "  ",
+		},
+		{
+			name:  "query with comment",
+			query: "SELECT 1  /* comment */",
+			want:  "SELECT 1  /* comment */",
+		},
+		{
+			name:  "query without attrs",
+			query: "SELECT 1",
+			want:  "SELECT 1",
+		},
+		{
+			name:  "query with single attr",
+			query: "SELECT 1",
+			opts:  []Option{WithAttrPairs("key", "value")},
+			want:  "SELECT 1 /* key='value' */",
+		},
+		{
+			name:  "query with multiple attrs",
+			query: "SELECT 1",
+			opts:  []Option{WithAttrPairs("key", "1value", "key2", "  value 2")},
+			want:  "SELECT 1 /* key='1value',key2='%20%20value%202' */",
+		},
+	}
+
+	for _, cs := range cases {
+		t.Run(cs.name, func(t *testing.T) {
+			got := Comment(cs.query, cs.opts...)
+			if want := cs.want; want != got {
+				t.Fatalf("got '%v', want '%v'", got, want)
+			}
+		})
+	}
+}