From 830b6ef3500a933d2061e9760cd4b349358c490f Mon Sep 17 00:00:00 2001
From: Dmitriy Seredenko <dmitriy.seredenko@kredito.de>
Date: Thu, 8 Jun 2017 09:18:30 +0200
Subject: [PATCH] * Support $1, $n queries like in PostgresSQL * Added tests
 for it * Small formatting update

---
 conn.go          | 18 +++++++++++++++---
 response.go      |  3 ++-
 response_test.go | 24 +++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/conn.go b/conn.go
index ddf9aa2..41bad8d 100644
--- a/conn.go
+++ b/conn.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"database/sql/driver"
 	"errors"
+	"log"
+	"regexp"
 	"strings"
 	"sync"
 )
@@ -66,9 +68,19 @@ func (c *FakeConn) Prepare(query string) (driver.Stmt, error) {
 // context is for the preparation of the statement,
 // it must not store the context within the statement itself.
 func (c *FakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
-	var firstStmt = &FakeStmt{q: query, connection: c}          // Create statement
-	firstStmt.placeholders = len(strings.Split(query, "?")) - 1 // Checking how many placeholders do we have
-	queryParts := strings.Split(query, " ")                     // By First statement define the query type
+	var firstStmt = &FakeStmt{q: query, connection: c}
+	// Checking how many placeholders do we have
+	if strings.Contains(query, "$1") {
+		r, err := regexp.Compile(`[$]\d+`)
+		if err != nil {
+			log.Fatalf(`Cant't compile regexp with err [%v]`, err)
+		}
+		firstStmt.placeholders = len(strings.Split(r.ReplaceAllString(query, `$$$`), "$$")) - 1 // Postgres notation
+	} else {
+		firstStmt.placeholders = len(strings.Split(query, "?")) - 1 // Postgres notation
+	}
+
+	queryParts := strings.Split(query, " ") // By First statement define the query type
 	firstStmt.command = strings.ToUpper(queryParts[0])
 	return firstStmt, nil
 }
diff --git a/response.go b/response.go
index c19f300..a420679 100644
--- a/response.go
+++ b/response.go
@@ -1,13 +1,14 @@
 package go_mocket
 
 import (
+	"database/sql"
 	"database/sql/driver"
 	"fmt"
 	"log"
 	"reflect"
 	"strings"
-	"database/sql"
 )
+
 const (
 	DRIVER_NAME = "MOCK_FAKE_DRIVER"
 )
diff --git a/response_test.go b/response_test.go
index bd86ed7..4fba5ea 100644
--- a/response_test.go
+++ b/response_test.go
@@ -43,7 +43,7 @@ func CreateUsersWithError(db *sql.DB) error {
 	return err
 }
 
-func InsertRecord(db *sql.DB) int64  {
+func InsertRecord(db *sql.DB) int64 {
 	res, err := db.Exec(`INSERT INTO foo VALUES("bar", ?))`, "value")
 	if err != nil {
 		return 0
@@ -143,4 +143,26 @@ func TestResponses(t *testing.T) {
 			t.Fatalf("Last insert id not returned. Expected: [%v] , Got: [%v]", mockedId, returnedId)
 		}
 	})
+
+	t.Run(`Recognise both ? and $1 Postgres placeholders for raw query`, func(t *testing.T) {
+		t.Run("Question mark", func(t *testing.T) {
+			testFunc := func(db *sql.DB) string {
+				var name string
+				err := db.QueryRow(`SELECT * FROM foo a = $1 AND b = $2 AND c = $3`, "value", "value2", "value3").Scan(&name)
+				if err != nil {
+					t.Fatalf("Test function failed [%v]", err)
+					return ""
+				}
+				return name
+			}
+
+			Catcher.Reset().NewMock().WithQuery("SELECT * FROM foo ").WithReply([]map[string]interface{}{{"name": "full_name"}})
+			returnedName := testFunc(DB)
+
+			if returnedName != "full_name" {
+				t.Fatalf("Returned name mismatches. Expected: [%v] , Got: [%v]", "full_name", returnedName)
+			}
+
+		})
+	})
 }