From e0b2d8d7ad0ba5e4d253f3e5bf3168d5bf4867db Mon Sep 17 00:00:00 2001
From: Andrei Goncear <agoncear@malwarebytes.com>
Date: Fri, 19 Jul 2024 08:59:22 +0000
Subject: [PATCH] add arguments compilation trough sqlbuilder

---
 chdb/driver/driver.go      | 26 ++++++++++++++++++++++++--
 chdb/driver/driver_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 go.mod                     |  2 ++
 go.sum                     |  5 +++++
 4 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/chdb/driver/driver.go b/chdb/driver/driver.go
index 9c2061e..9977582 100644
--- a/chdb/driver/driver.go
+++ b/chdb/driver/driver.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/chdb-io/chdb-go/chdb"
 	"github.com/chdb-io/chdb-go/chdbstable"
+	"github.com/huandu/go-sqlbuilder"
 	"github.com/parquet-go/parquet-go"
 
 	"github.com/apache/arrow/go/v15/arrow/ipc"
@@ -196,9 +197,30 @@ func (c *conn) Query(query string, values []driver.Value) (driver.Rows, error) {
 	return c.QueryContext(context.Background(), query, namedValues)
 }
 
-func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+func (c *conn) compileArguments(query string, args []driver.NamedValue) (string, error) {
+	var compiledQuery string
+	if len(args) > 0 {
+		compiledArgs := make([]interface{}, len(args))
+		for idx := range args {
+			compiledArgs[idx] = args[idx].Value
+		}
+		compiled, err := sqlbuilder.ClickHouse.Interpolate(query, compiledArgs)
+		if err != nil {
+			return "", err
+		}
+		compiledQuery = compiled
+	} else {
+		compiledQuery = query
+	}
+	return compiledQuery, nil
+}
 
-	result, err := c.QueryFun(query, c.driverType.String(), c.udfPath)
+func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	compiledQuery, err := c.compileArguments(query, args)
+	if err != nil {
+		return nil, err
+	}
+	result, err := c.QueryFun(compiledQuery, c.driverType.String(), c.udfPath)
 	if err != nil {
 		return nil, err
 	}
diff --git a/chdb/driver/driver_test.go b/chdb/driver/driver_test.go
index 12fe356..fb6b4ae 100644
--- a/chdb/driver/driver_test.go
+++ b/chdb/driver/driver_test.go
@@ -47,6 +47,44 @@ func TestDb(t *testing.T) {
 	}
 }
 
+func TestDbWithCompiledArgs(t *testing.T) {
+	db, err := sql.Open("chdb", "")
+	if err != nil {
+		t.Errorf("open db fail, err:%s", err)
+	}
+	if db.Ping() != nil {
+		t.Errorf("ping db fail")
+	}
+	rows, err := db.Query(`SELECT ?, ?`, 1, "abc")
+	if err != nil {
+		t.Errorf("run Query fail, err:%s", err)
+	}
+	cols, err := rows.Columns()
+	if err != nil {
+		t.Errorf("get result columns fail, err: %s", err)
+	}
+	if len(cols) != 2 {
+		t.Errorf("select result columns length should be 2")
+	}
+	var (
+		bar int
+		foo string
+	)
+	defer rows.Close()
+	for rows.Next() {
+		err := rows.Scan(&bar, &foo)
+		if err != nil {
+			t.Errorf("scan fail, err: %s", err)
+		}
+		if bar != 1 {
+			t.Errorf("expected error")
+		}
+		if foo != "abc" {
+			t.Errorf("expected error")
+		}
+	}
+}
+
 func TestDbWithOpt(t *testing.T) {
 	for _, kv := range []struct {
 		opt       string
diff --git a/go.mod b/go.mod
index 783aeb2..c97fee5 100644
--- a/go.mod
+++ b/go.mod
@@ -13,6 +13,8 @@ require (
 	github.com/goccy/go-json v0.10.3 // indirect
 	github.com/google/flatbuffers v24.3.25+incompatible // indirect
 	github.com/google/uuid v1.6.0 // indirect
+	github.com/huandu/go-sqlbuilder v1.27.3 // indirect
+	github.com/huandu/xstrings v1.4.0 // indirect
 	github.com/klauspost/compress v1.17.9 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.8 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
diff --git a/go.sum b/go.sum
index ca90107..444020a 100644
--- a/go.sum
+++ b/go.sum
@@ -14,6 +14,11 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
 github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
+github.com/huandu/go-assert v1.1.6/go.mod h1:JuIfbmYG9ykwvuxoJ3V8TB5QP+3+ajIA54Y44TmkMxs=
+github.com/huandu/go-sqlbuilder v1.27.3 h1:cNVF9vQP4i7rTk6XXJIEeMbGkZbxfjcITeJzobJK44k=
+github.com/huandu/go-sqlbuilder v1.27.3/go.mod h1:mS0GAtrtW+XL6nM2/gXHRJax2RwSW1TraavWDFAc1JA=
+github.com/huandu/xstrings v1.4.0 h1:D17IlohoQq4UcpqD7fDk80P7l+lwAmlFaBHgOipl2FU=
+github.com/huandu/xstrings v1.4.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
 github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
 github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/klauspost/cpuid/v2 v2.2.8 h1:+StwCXwm9PdpiEkPyzBXIy+M9KUb4ODm0Zarf1kS5BM=