Repository: jackc/pgx
Pull request #960 (2021-03-13 10:16:30)

Add comment support when sanitizing SQL queries
There is a potential error when using comments in a SQL query because they are not processed, for example, if there are quotes, double quotes, or a dollar sign inside the comment.
Added support for comments and example tests.


Patch:
diff --git a/go.mod b/go.mod
index 5027afad..bfe6e3e5 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module github.com/jackc/pgx/v4
 go 1.12
 
 require (
+	github.com/Masterminds/semver/v3 v3.1.1 // indirect
 	github.com/cockroachdb/apd v1.1.0
 	github.com/go-kit/kit v0.10.0
 	github.com/gofrs/uuid v3.2.0+incompatible
diff --git a/go.sum b/go.sum
index 8a2c0bb2..e7636934 100644
--- a/go.sum
+++ b/go.sum
@@ -3,6 +3,8 @@ cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/Knetic/govaluate v3.0.1-0.20171022003610-9aa49832a739+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0=
+github.com/Masterminds/semver/v3 v3.1.1 h1:hLg3sBzpNErnxhQtUy/mmLR2I9foDujNK030IGemrRc=
+github.com/Masterminds/semver/v3 v3.1.1/go.mod h1:VPu/7SZ7ePZ3QOrcuXROw5FAcLl4a0cBrbBpGY/8hQs=
 github.com/Shopify/sarama v1.19.0/go.mod h1:FVkBWblsNy7DGZRfXLU0O9RCGt5g3g3yEuWXgklEdEo=
 github.com/Shopify/toxiproxy v2.1.4+incompatible/go.mod h1:OXgGpZ6Cli1/URJOF1DMxUHB2q5Ap20/P/eIdh4G0pI=
 github.com/VividCortex/gohistogram v1.0.0/go.mod h1:Pf5mBqqDxYaXu3hDrrU+w6nw50o/4+TcAqDqk/vUH7g=
diff --git a/internal/sanitize/sanitize.go b/internal/sanitize/sanitize.go
index 0e7ddf13..30a4a757 100644
--- a/internal/sanitize/sanitize.go
+++ b/internal/sanitize/sanitize.go
@@ -83,7 +83,7 @@ func NewQuery(sql string) (*Query, error) {
 }
 
 func QuoteString(str string) string {
-	return "'" + strings.Replace(str, "'", "''", -1) + "'"
+	return "'" + strings.ReplaceAll(str, "'", "''") + "'"
 }
 
 func QuoteBytes(buf []byte) string {
@@ -94,6 +94,7 @@ type sqlLexer struct {
 	src     string
 	start   int
 	pos     int
+	nested  int // multiline comment nesting level.
 	stateFn stateFn
 	parts   []Part
 }
@@ -125,6 +126,18 @@ func rawState(l *sqlLexer) stateFn {
 				l.start = l.pos
 				return placeholderState
 			}
+		case '-':
+			nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:])
+			if nextRune == '-' {
+				l.pos += width
+				return oneLineCommentState
+			}
+		case '/':
+			nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:])
+			if nextRune == '*' {
+				l.pos += width
+				return multilineCommentState
+			}
 		case utf8.RuneError:
 			if l.pos-l.start > 0 {
 				l.parts = append(l.parts, l.src[l.start:l.pos])
@@ -225,6 +238,61 @@ func escapeStringState(l *sqlLexer) stateFn {
 	}
 }
 
+func oneLineCommentState(l *sqlLexer) stateFn {
+	for {
+		r, width := utf8.DecodeRuneInString(l.src[l.pos:])
+		l.pos += width
+
+		switch r {
+		case '\\':
+			_, width = utf8.DecodeRuneInString(l.src[l.pos:])
+			l.pos += width
+		case '\n':
+			return rawState
+		case utf8.RuneError:
+			if l.pos-l.start > 0 {
+				l.parts = append(l.parts, l.src[l.start:l.pos])
+				l.start = l.pos
+			}
+			return nil
+		}
+	}
+}
+
+func multilineCommentState(l *sqlLexer) stateFn {
+	for {
+		r, width := utf8.DecodeRuneInString(l.src[l.pos:])
+		l.pos += width
+
+		switch r {
+		case '/':
+			nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:])
+			if nextRune == '*' {
+				l.pos += width
+				l.nested++
+			}
+		case '*':
+			nextRune, width := utf8.DecodeRuneInString(l.src[l.pos:])
+			if nextRune != '/' {
+				continue
+			}
+
+			l.pos += width
+			if l.nested == 0 {
+				return rawState
+			}
+			l.nested--
+
+		case utf8.RuneError:
+			if l.pos-l.start > 0 {
+				l.parts = append(l.parts, l.src[l.start:l.pos])
+				l.start = l.pos
+			}
+			return nil
+		}
+	}
+}
+
 // SanitizeSQL replaces placeholder values with args. It quotes and escapes args
 // as necessary. This function is only safe when standard_conforming_strings is
 // on.
diff --git a/stdlib/sql.go b/stdlib/sql.go
index 10c104da..59d24c62 100644
--- a/stdlib/sql.go
+++ b/stdlib/sql.go
@@ -514,7 +514,7 @@ func (r *Rows) ColumnTypeScanType(index int) reflect.Type {
 
 func (r *Rows) Close() error {
 	r.rows.Close()
-	return nil
+	return r.rows.Err()
 }
 
 func (r *Rows) Next(dest []driver.Value) error {


Test patch:
diff --git a/internal/sanitize/sanitize_test.go b/internal/sanitize/sanitize_test.go
index fff14896..344c46b0 100644
--- a/internal/sanitize/sanitize_test.go
+++ b/internal/sanitize/sanitize_test.go
@@ -60,6 +60,30 @@ func TestNewQuery(t *testing.T) {
 			sql:      `select e'escape string\' $42', $1`,
 			expected: sanitize.Query{Parts: []sanitize.Part{`select e'escape string\' $42', `, 1}},
 		},
+		{
+			sql:      `select /* a baby's toy */ 'barbie', $1`,
+			expected: sanitize.Query{Parts: []sanitize.Part{`select /* a baby's toy */ 'barbie', `, 1}},
+		},
+		{
+			sql:      `select /* *_* */ $1`,
+			expected: sanitize.Query{Parts: []sanitize.Part{`select /* *_* */ `, 1}},
+		},
+		{
+			sql:      `select 42 /* /* /* 42 */ */ */, $1`,
+			expected: sanitize.Query{Parts: []sanitize.Part{`select 42 /* /* /* 42 */ */ */, `, 1}},
+		},
+		{
+			sql:      "select -- a baby's toy\n'barbie', $1",
+			expected: sanitize.Query{Parts: []sanitize.Part{"select -- a baby's toy\n'barbie', ", 1}},
+		},
+		{
+			sql:      "select 42 -- is a Thinker's favorite number",
+			expected: sanitize.Query{Parts: []sanitize.Part{"select 42 -- is a Thinker's favorite number"}},
+		},
+		{
+			sql:      "select 42, -- \\nis a Thinker's favorite number\n$1",
+			expected: sanitize.Query{Parts: []sanitize.Part{"select 42, -- \\nis a Thinker's favorite number\n", 1}},
+		},
 	}
 
 	for i, tt := range successTests {
diff --git a/stdlib/sql_test.go b/stdlib/sql_test.go
index a053c6cc..7fb4640a 100644
--- a/stdlib/sql_test.go
+++ b/stdlib/sql_test.go
@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Masterminds/semver/v3"
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgx/v4"
 	"github.com/jackc/pgx/v4/stdlib"
@@ -45,6 +46,37 @@ func skipCockroachDB(t testing.TB, db *sql.DB, msg string) {
 	require.NoError(t, err)
 }
 
+func skipPostgreSQLVersion(t testing.TB, db *sql.DB, constraintStr, msg string) {
+	conn, err := db.Conn(context.Background())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	err = conn.Raw(func(driverConn interface{}) error {
+		conn := driverConn.(*stdlib.Conn).Conn()
+		serverVersionStr := conn.PgConn().ParameterStatus("server_version")
+		// if not PostgreSQL do nothing
+		if serverVersionStr == "" {
+			return nil
+		}
+
+		serverVersion, err := semver.NewVersion(serverVersionStr)
+		if err != nil {
+			return err
+		}
+
+		c, err := semver.NewConstraint(constraintStr)
+		if err != nil {
+			return err
+		}
+
+		if c.Check(serverVersion) {
+			t.Skip(msg)
+		}
+		return nil
+	})
+	require.NoError(t, err)
+}
+
 func testWithAndWithoutPreferSimpleProtocol(t *testing.T, f func(t *testing.T, db *sql.DB)) {
 	t.Run("SimpleProto",
 		func(t *testing.T) {
@@ -1047,3 +1079,43 @@ func TestRegisterConnConfig(t *testing.T) {
 	assert.Equal(t, "Query", l.msg)
 	assert.Equal(t, "select 1", l.data["sql"])
 }
+
+// https://github.com/jackc/pgx/issues/958
+func TestConnQueryRowConstraintErrors(t *testing.T) {
+	testWithAndWithoutPreferSimpleProtocol(t, func(t *testing.T, db *sql.DB) {
+		skipPostgreSQLVersion(t, db, "< 11", "Test requires PG 11+")
+		skipCockroachDB(t, db, "Server does not support deferred constraint (https://github.com/cockroachdb/cockroach/issues/31632)")
+
+		_, err := db.Exec(`create temporary table defer_test (
+			id text primary key,
+			n int not null, unique (n),
+			unique (n) deferrable initially deferred )`)
+		require.NoError(t, err)
+
+		_, err = db.Exec(`drop function if exists test_trigger cascade`)
+		require.NoError(t, err)
+
+		_, err = db.Exec(`create function test_trigger() returns trigger language plpgsql as $$
+		begin
+		if new.n = 4 then
+			raise exception 'n cant be 4!';
+		end if;
+		return new;
+	end$$`)
+		require.NoError(t, err)
+
+		_, err = db.Exec(`create constraint trigger test
+			after insert or update on defer_test
+			deferrable initially deferred
+			for each row
+			execute function test_trigger()`)
+		require.NoError(t, err)
+
+		_, err = db.Exec(`insert into defer_test (id, n) values ('a', 1), ('b', 2), ('c', 3)`)
+		require.NoError(t, err)
+
+		var id string
+		err = db.QueryRow(`insert into defer_test (id, n) values ('e', 4) returning id`).Scan(&id)
+		assert.Error(t, err)
+	})
+}