Merge pull request #852 from jkatz/quote-literal

Add `QuoteLiteral` function.

Author: maddyblue

conn.go

@@ -1500,6 +1500,39 @@ func QuoteIdentifier(name string) string {
 	return `"` + strings.Replace(name, `"`, `""`, -1) + `"`
 }
 
+// QuoteLiteral quotes a 'literal' (e.g. a parameter, often used to pass literal
+// to DDL and other statements that do not accept parameters) to be used as part
+// of an SQL statement.  For example:
+//
+//    exp_date := pq.QuoteLiteral("2023-01-05 15:00:00Z")
+//    err := db.Exec(fmt.Sprintf("CREATE ROLE my_user VALID UNTIL %s", exp_date))
+//
+// Any single quotes in name will be escaped. Any backslashes (i.e. "\") will be
+// replaced by two backslashes (i.e. "\\") and the C-style escape identifier
+// that PostgreSQL provides ('E') will be prepended to the string.
+func QuoteLiteral(literal string) string {
+	// This follows the PostgreSQL internal algorithm for handling quoted literals
+	// from libpq, which can be found in the "PQEscapeStringInternal" function,
+	// which is found in the libpq/fe-exec.c source file:
+	// https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/interfaces/libpq/fe-exec.c
+	//
+	// substitute any single-quotes (') with two single-quotes ('')
+	literal = strings.Replace(literal, `'`, `''`, -1)
+	// determine if the string has any backslashes (\) in it.
+	// if it does, replace any backslashes (\) with two backslashes (\\)
+	// then, we need to wrap the entire string with a PostgreSQL
+	// C-style escape. Per how "PQEscapeStringInternal" handles this case, we
+	// also add a space before the "E"
+	if strings.Contains(literal, `\`) {
+		literal = strings.Replace(literal, `\`, `\\`, -1)
+		literal = ` E'` + literal + `'`
+	} else {
+		// otherwise, we can just wrap the literal with a pair of single quotes
+		literal = `'` + literal + `'`
+	}
+	return literal
+}
+
 func md5s(s string) string {
 	h := md5.New()
 	h.Write([]byte(s))

conn_test.go

@@ -1554,6 +1554,38 @@ func TestQuoteIdentifier(t *testing.T) {
 	}
 }
 
+func TestQuoteLiteral(t *testing.T) {
+	var cases = []struct {
+		input string
+		want  string
+	}{
+		{`foo`, `'foo'`},
+		{`foo bar baz`, `'foo bar baz'`},
+		{`foo'bar`, `'foo''bar'`},
+		{`foo\bar`, ` E'foo\\bar'`},
+		{`foo\ba'r`, ` E'foo\\ba''r'`},
+		{`foo"bar`, `'foo"bar'`},
+		{`foo\x00bar`, ` E'foo\\x00bar'`},
+		{`\x00foo`, ` E'\\x00foo'`},
+		{`'`, `''''`},
+		{`''`, `''''''`},
+		{`\`, ` E'\\'`},
+		{`'abc'; DROP TABLE users;`, `'''abc''; DROP TABLE users;'`},
+		{`\'`, ` E'\\'''`},
+		{`E'\''`, ` E'E''\\'''''`},
+		{`e'\''`, ` E'e''\\'''''`},
+		{`E'\'abc\'; DROP TABLE users;'`, ` E'E''\\''abc\\''; DROP TABLE users;'''`},
+		{`e'\'abc\'; DROP TABLE users;'`, ` E'e''\\''abc\\''; DROP TABLE users;'''`},
+	}
+
+	for _, test := range cases {
+		got := QuoteLiteral(test.input)
+		if got != test.want {
+			t.Errorf("QuoteLiteral(%q) = %v want %v", test.input, got, test.want)
+		}
+	}
+}
+
 func TestRowsResultTag(t *testing.T) {
 	type ResultTag interface {
 		Result() driver.Result