Merge pull request #5 from mashiike/feature/ordinal-parameters

support parameter $ and ?

Author: mashiike

conn.go

@@ -66,7 +66,7 @@ func (conn *redshiftDataConn) QueryContext(ctx context.Context, query string, ar
 
 func (conn *redshiftDataConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
 	params := &redshiftdata.ExecuteStatementInput{
-		Sql:        nullif(query),
+		Sql:        nullif(rewriteQuery(query, len(args))),
 		Parameters: convertArgsToParameters(args),
 	}
 	_, output, err := conn.executeStatement(ctx, params)
@@ -76,6 +76,40 @@ func (conn *redshiftDataConn) ExecContext(ctx context.Context, query string, arg
 	return newResult(output), nil
 }
 
+func rewriteQuery(query string, paramsCount int) string {
+	if paramsCount == 0 {
+		return query
+	}
+	runes := make([]rune, 0, len(query))
+	stack := make([]rune, 0)
+	var exclamationCount int
+	for _, r := range query {
+		if len(stack) > 0 {
+			if r == stack[len(stack)-1] {
+				stack = stack[:len(stack)-1]
+				runes = append(runes, r)
+				continue
+			}
+		} else {
+			switch r {
+			case '?':
+				exclamationCount++
+				runes = append(runes, []rune(fmt.Sprintf(":%d", exclamationCount))...)
+				continue
+			case '$':
+				runes = append(runes, ':')
+				continue
+			}
+		}
+		switch r {
+		case '"', '\'':
+			stack = append(stack, r)
+		}
+		runes = append(runes, r)
+	}
+	return string(runes)
+}
+
 func convertArgsToParameters(args []driver.NamedValue) []types.SqlParameter {
 	if len(args) == 0 {
 		return nil

conn_test.go

@@ -0,0 +1,47 @@
+package redshiftdatasqldriver
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRewriteQuery(t *testing.T) {
+	cases := []struct {
+		casename    string
+		query       string
+		paramsCount int
+		expected    string
+	}{
+		{
+			casename:    "no params",
+			query:       `SELECT * FROM pg_user`,
+			paramsCount: 0,
+			expected:    `SELECT * FROM pg_user`,
+		},
+		{
+			casename:    "no change",
+			query:       `SELECT * FROM pg_user WHERE usename = :name`,
+			paramsCount: 1,
+			expected:    `SELECT * FROM pg_user WHERE usename = :name`,
+		},
+		{
+			casename:    "? rewrite",
+			query:       `SELECT 'hoge?' FROM pg_user WHERE usename = ? AND usesysid > ?`,
+			paramsCount: 1,
+			expected:    `SELECT 'hoge?' FROM pg_user WHERE usename = :1 AND usesysid > :2`,
+		},
+		{
+			casename:    "$ rewrite",
+			query:       `SELECT '3$1$' FROM table WHERE "$column" = $1 AND column1 > $2 AND column2 < $1`,
+			paramsCount: 1,
+			expected:    `SELECT '3$1$' FROM table WHERE "$column" = :1 AND column1 > :2 AND column2 < :1`,
+		},
+	}
+	for _, c := range cases {
+		t.Run(c.casename, func(t *testing.T) {
+			actual := rewriteQuery(c.query, c.paramsCount)
+			require.Equal(t, c.expected, actual)
+		})
+	}
+}