Escape Parameters (#9)

* escape parameters

* fix test

* rename const

Author: Heqing Ya

connection.go

@@ -65,7 +65,10 @@ func (c *Conn) QueryContext(ctx context.Context, q string, args []driver.NamedVa
 	}
 
 	tmpl := template(q)
-	stmt := statement(tmpl, args)
+	stmt, err := statement(tmpl, args)
+	if err != nil {
+		return nil, err
+	}
 	return query(ctx, session, stmt)
 }
 
@@ -77,7 +80,10 @@ func (c *Conn) ExecContext(ctx context.Context, q string, args []driver.NamedVal
 	}
 
 	tmpl := template(q)
-	stmt := statement(tmpl, args)
+	stmt, err := statement(tmpl, args)
+	if err != nil {
+		return nil, err
+	}
 	return exec(ctx, session, stmt)
 }
 

escaper.go

@@ -0,0 +1,29 @@
+package dbsql
+
+import (
+	"database/sql/driver"
+	"fmt"
+	"strings"
+	"time"
+)
+
+const (
+	TimeFmt = "2006-01-02T15:04:05.999-07:00"
+)
+
+func EscapeArg(arg driver.NamedValue) (string, error) {
+	switch v := arg.Value.(type) {
+	case int64:
+		return fmt.Sprintf("%v", v), nil
+	case float64:
+		return fmt.Sprintf("%v", v), nil
+	case bool:
+		return fmt.Sprintf("%v", v), nil
+	case string:
+		return fmt.Sprintf("'%v'", strings.ReplaceAll(v, "'", "''")), nil
+	case time.Time:
+		return "'" + v.Format(TimeFmt) + "'", nil
+	default:
+		return "", fmt.Errorf("unsupported parameter type %T for value %v", v, v)
+	}
+}

escaper_test.go

@@ -0,0 +1,31 @@
+package dbsql
+
+import (
+	"database/sql/driver"
+	"testing"
+	"time"
+)
+
+func TestEscaper(t *testing.T) {
+	testcases := []struct {
+		Value          driver.Value
+		ExpectedOutput string
+	}{
+		{Value: "a'b'c", ExpectedOutput: `'a''b''c'`},
+		{Value: int64(1024), ExpectedOutput: `1024`},
+		{Value: float64(1024.5), ExpectedOutput: `1024.5`},
+		{Value: true, ExpectedOutput: "true"},
+		{Value: time.Date(2020, time.April, 11, 21, 34, 01, 0, time.UTC), ExpectedOutput: "'2020-04-11T21:34:01+00:00'"},
+	}
+
+	for _, test := range testcases {
+		actual, err := EscapeArg(driver.NamedValue{Value: test.Value})
+		if err != nil {
+			t.Error(err)
+		}
+		if actual != test.ExpectedOutput {
+			t.Errorf("expecting %v, actual value: %v", test.ExpectedOutput, actual)
+		}
+
+	}
+}

statement.go

@@ -64,7 +64,10 @@ func (s *Stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driv
 	if err != nil {
 		return nil, err
 	}
-	stmt := statement(s.stmt, args)
+	stmt, err := statement(s.stmt, args)
+	if err != nil {
+		return nil, err
+	}
 	return query(ctx, session, stmt)
 }
 
@@ -74,7 +77,10 @@ func (s *Stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (drive
 	if err != nil {
 		return nil, err
 	}
-	stmt := statement(s.stmt, args)
+	stmt, err := statement(s.stmt, args)
+	if err != nil {
+		return nil, err
+	}
 	return exec(ctx, session, stmt)
 }
 
@@ -92,7 +98,7 @@ func template(query string) string {
 	return query
 }
 
-func statement(tmpl string, args []driver.NamedValue) string {
+func statement(tmpl string, args []driver.NamedValue) (string, error) {
 	stmt := tmpl
 	for _, arg := range args {
 		var re *regexp.Regexp
@@ -101,10 +107,14 @@ func statement(tmpl string, args []driver.NamedValue) string {
 		} else {
 			re = regexp.MustCompile(fmt.Sprintf("@p%d%s", arg.Ordinal, `\b`))
 		}
-		val := fmt.Sprintf("%v", arg.Value)
+		escaped, err := EscapeArg(arg)
+		if err != nil {
+			return "", err
+		}
+		val := fmt.Sprintf("%v", escaped)
 		stmt = re.ReplaceAllString(stmt, val)
 	}
-	return stmt
+	return stmt, nil
 }
 
 func query(ctx context.Context, session *hive.Session, stmt string) (driver.Rows, error) {

statement_test.go

@@ -16,7 +16,7 @@ func TestStatement(t *testing.T) {
 			args: []driver.NamedValue{
 				{Ordinal: 1, Value: "val_1"},
 			},
-			target: "val_1 p1",
+			target: "'val_1' p1",
 		},
 		{
 			stmt: "@p1 @p10 @p11 @named @named1 @p1",
@@ -25,12 +25,15 @@ func TestStatement(t *testing.T) {
 				{Ordinal: 10, Name: "named", Value: "val_named"},
 				{Ordinal: 11, Value: "val_11"},
 			},
-			target: "val_1 @p10 val_11 val_named @named1 val_1",
+			target: "'val_1' @p10 'val_11' 'val_named' @named1 'val_1'",
 		},
 	}
 
 	for _, tt := range tests {
-		result := statement(tt.stmt, tt.args)
+		result, err := statement(tt.stmt, tt.args)
+		if err != nil {
+			t.Error(err)
+		}
 
 		if result != tt.target {
 			t.Fatalf("mismatch for statement: %q\n\ttarget: %q\n\tresult: %q", tt.stmt, tt.target, result)