diff --git a/interpolate.go b/interpolate.go index 7a8933e..323f4da 100644 --- a/interpolate.go +++ b/interpolate.go @@ -170,6 +170,19 @@ func postgresqlInterpolate(query string, args ...interface{}) (string, error) { continue } + var tokenStart rune + for pos := offset - sz; pos > 0; { + previous, size := utf8.DecodeLastRuneInString(query[:pos]) + if !unicode.IsLetter(previous) && previous != '_' && previous != '$' && (previous < '0' || previous > '9') { + break + } + tokenStart = previous + pos -= size + } + if tokenStart == '_' || unicode.IsLetter(tokenStart) { + continue + } + oldSz := sz pos := offset r, sz = utf8.DecodeRuneInString(target) @@ -212,7 +225,11 @@ func postgresqlInterpolate(query string, args ...interface{}) (string, error) { } // Try to find the beginning of dollar quote. - for ; sz != 0 && r != '$' && unicode.IsLetter(r); r, sz = utf8.DecodeRuneInString(target) { + if r != '$' && r != '_' && !unicode.IsLetter(r) { + continue + } + + for ; sz != 0 && (unicode.IsLetter(r) || r == '_' || ('0' <= r && r <= '9')); r, sz = utf8.DecodeRuneInString(target) { pos += sz target = query[pos:] } @@ -221,7 +238,7 @@ func postgresqlInterpolate(query string, args ...interface{}) (string, error) { break } - if !unicode.IsLetter(r) && r != '$' { + if r != '$' { continue } diff --git a/interpolate_test.go b/interpolate_test.go index 1d01ad0..4701bfe 100644 --- a/interpolate_test.go +++ b/interpolate_test.go @@ -391,6 +391,45 @@ func TestFlavorInterpolate(t *testing.T) { } } +func TestPostgreSQLInterpolateDollarQuoteTags(t *testing.T) { + for _, tag := range []string{"", "tag", "_", "tag_name", "tag1", "_tag_1", "标签1"} { + t.Run("tag="+tag, func(t *testing.T) { + delimiter := "$" + tag + "$" + quoted := delimiter + "literal $2" + delimiter + query, err := PostgreSQL.Interpolate("SELECT "+quoted+", $1", []interface{}{42}) + if err != nil { + t.Fatal(err) + } + if want := "SELECT " + quoted + ", 42"; query != want { + t.Fatalf("unexpected query: got %q, want %q", query, want) + } + }) + } + + // $0$ is invalid SQL; this only checks that a digit cannot start a quote tag. + query, err := PostgreSQL.Interpolate("SELECT $0$, $1", []interface{}{42}) + if err != nil { + t.Fatal(err) + } + if want := "SELECT $0$, 42"; query != want { + t.Fatalf("unexpected query: got %q, want %q", query, want) + } +} + +func TestPostgreSQLInterpolateDollarQuoteIdentifiers(t *testing.T) { + for _, identifier := range []string{"foo$_bar$", "foo$bar$", "foo$$", "foo1$_bar$", "foo$2$_bar$", "标签$_bar$"} { + t.Run(identifier, func(t *testing.T) { + query, err := PostgreSQL.Interpolate("SELECT "+identifier+", $1", []interface{}{42}) + if err != nil { + t.Fatal(err) + } + if want := "SELECT " + identifier + ", 42"; query != want { + t.Fatalf("unexpected query: got %q, want %q", query, want) + } + }) + } +} + func TestANSIInterpolationEscapesQuotes(t *testing.T) { payload := "x'; DROP TABLE users; --" cases := []struct {