add: golang sql interface

chdb/driver/driver.go

@@ -0,0 +1,256 @@
+package chdb
+
+import (
+	"bytes"
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"fmt"
+	"reflect"
+	"time"
+
+	"github.com/apache/arrow/go/v14/arrow"
+	"github.com/apache/arrow/go/v14/arrow/array"
+	"github.com/apache/arrow/go/v14/arrow/decimal128"
+	"github.com/apache/arrow/go/v14/arrow/decimal256"
+	wrapper "github.com/chdb-io/chdb-go/chdb"
+	"github.com/chdb-io/chdb-go/chdbstable"
+
+	"github.com/apache/arrow/go/v14/arrow/ipc"
+)
+
+func init() {
+	sql.Register("chdb", Driver{})
+}
+
+type connector struct {
+}
+
+// Connect returns a connection to a database.
+func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
+	return &conn{}, nil
+}
+
+// Driver returns the underying Driver of the connector,
+// compatibility with the Driver method on sql.DB
+func (c *connector) Driver() driver.Driver { return Driver{} }
+
+type Driver struct{}
+
+// Open returns a new connection to the database.
+func (d Driver) Open(name string) (driver.Conn, error) {
+	return &conn{}, nil
+}
+
+// OpenConnector expects the same format as driver.Open
+func (d Driver) OpenConnector(dataSourceName string) (driver.Connector, error) {
+	return &connector{}, nil
+}
+
+type conn struct {
+}
+
+func (c *conn) Close() error {
+	return nil
+}
+
+func (c *conn) Query(query string, values []driver.Value) (driver.Rows, error) {
+	namedValues := make([]driver.NamedValue, len(values))
+	for i, value := range values {
+		namedValues[i] = driver.NamedValue{
+			// nb: Name field is optional
+			Ordinal: i,
+			Value:   value,
+		}
+	}
+	return c.QueryContext(context.Background(), query, namedValues)
+}
+
+func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	result := wrapper.Query(query, "Arrow")
+	buf := result.Buf()
+	if buf == nil {
+		return nil, fmt.Errorf("result is nil")
+	}
+	reader, err := ipc.NewFileReader(bytes.NewReader(buf))
+	if err != nil {
+		return nil, err
+	}
+	return &rows{localResult: result, reader: reader}, nil
+}
+
+func (c *conn) Begin() (driver.Tx, error) {
+	return nil, fmt.Errorf("does not support Transcation")
+}
+
+func (c *conn) Prepare(query string) (driver.Stmt, error) {
+	return c.PrepareContext(context.Background(), query)
+}
+
+func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	return nil, fmt.Errorf("does not support prepare statement")
+}
+
+// todo: func(c *conn) Prepare(query string)
+// todo: func(c *conn) PrepareContext(ctx context.Context, query string)
+// todo: prepared statment
+
+type rows struct {
+	localResult *chdbstable.LocalResult
+	reader      *ipc.FileReader
+	curRecord   arrow.Record
+	curRow      int64
+}
+
+func (r *rows) Columns() (out []string) {
+	sch := r.reader.Schema()
+	for i := 0; i < sch.NumFields(); i++ {
+		out = append(out, sch.Field(i).Name)
+	}
+	return
+}
+
+func (r *rows) Close() error {
+	if r.curRecord != nil {
+		r.curRecord = nil
+	}
+	// ignore reader close
+	_ = r.reader.Close()
+	r.reader = nil
+	r.localResult = nil
+	return nil
+}
+
+func (r *rows) Next(dest []driver.Value) error {
+	if r.curRecord != nil && r.curRow == r.curRecord.NumRows() {
+		r.curRecord = nil
+	}
+	for r.curRecord == nil {
+		record, err := r.reader.Read()
+		if err != nil {
+			return err
+		}
+		if record.NumRows() == 0 {
+			continue
+		}
+		r.curRecord = record
+		r.curRow = 0
+	}
+
+	for i, col := range r.curRecord.Columns() {
+		if col.IsNull(int(r.curRow)) {
+			dest[i] = nil
+			continue
+		}
+		switch col := col.(type) {
+		case *array.Boolean:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int8:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint8:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int16:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint16:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int32:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint32:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int64:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint64:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Float32:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Float64:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.String:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.LargeString:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Binary:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.LargeBinary:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Date32:
+			dest[i] = col.Value(int(r.curRow)).ToTime()
+		case *array.Date64:
+			dest[i] = col.Value(int(r.curRow)).ToTime()
+		case *array.Time32:
+			dest[i] = col.Value(int(r.curRow)).ToTime(col.DataType().(*arrow.Time32Type).Unit)
+		case *array.Time64:
+			dest[i] = col.Value(int(r.curRow)).ToTime(col.DataType().(*arrow.Time64Type).Unit)
+		case *array.Timestamp:
+			dest[i] = col.Value(int(r.curRow)).ToTime(col.DataType().(*arrow.TimestampType).Unit)
+		case *array.Decimal128:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Decimal256:
+			dest[i] = col.Value(int(r.curRow))
+		default:
+			return fmt.Errorf(
+				"not yet implemented populating from columns of type " + col.DataType().String(),
+			)
+		}
+	}
+
+	r.curRow++
+	return nil
+}
+
+func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
+	return r.reader.Schema().Field(index).Type.String()
+}
+
+func (r *rows) ColumnTypeNullable(index int) (nullable, ok bool) {
+	return r.reader.Schema().Field(index).Nullable, true
+}
+
+func (r *rows) ColumnTypePrecisionScale(index int) (precision, scale int64, ok bool) {
+	typ := r.reader.Schema().Field(index).Type
+	switch dt := typ.(type) {
+	case *arrow.Decimal128Type:
+		return int64(dt.Precision), int64(dt.Scale), true
+	case *arrow.Decimal256Type:
+		return int64(dt.Precision), int64(dt.Scale), true
+	}
+	return 0, 0, false
+}
+
+func (r *rows) ColumnTypeScanType(index int) reflect.Type {
+	switch r.reader.Schema().Field(index).Type.ID() {
+	case arrow.BOOL:
+		return reflect.TypeOf(false)
+	case arrow.INT8:
+		return reflect.TypeOf(int8(0))
+	case arrow.UINT8:
+		return reflect.TypeOf(uint8(0))
+	case arrow.INT16:
+		return reflect.TypeOf(int16(0))
+	case arrow.UINT16:
+		return reflect.TypeOf(uint16(0))
+	case arrow.INT32:
+		return reflect.TypeOf(int32(0))
+	case arrow.UINT32:
+		return reflect.TypeOf(uint32(0))
+	case arrow.INT64:
+		return reflect.TypeOf(int64(0))
+	case arrow.UINT64:
+		return reflect.TypeOf(uint64(0))
+	case arrow.FLOAT32:
+		return reflect.TypeOf(float32(0))
+	case arrow.FLOAT64:
+		return reflect.TypeOf(float64(0))
+	case arrow.DECIMAL128:
+		return reflect.TypeOf(decimal128.Num{})
+	case arrow.DECIMAL256:
+		return reflect.TypeOf(decimal256.Num{})
+	case arrow.BINARY:
+		return reflect.TypeOf([]byte{})
+	case arrow.STRING:
+		return reflect.TypeOf(string(""))
+	case arrow.TIME32, arrow.TIME64, arrow.DATE32, arrow.DATE64, arrow.TIMESTAMP:
+		return reflect.TypeOf(time.Time{})
+	}
+	return nil
+}

chdb/driver/driver_test.go

@@ -0,0 +1,46 @@
+package chdb
+
+import (
+	"database/sql"
+	"testing"
+)
+
+func TestDb(t *testing.T) {
+	db, err := sql.Open("chdb", "")
+	if err != nil {
+		t.Errorf("open db fail")
+	}
+	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")
+			}
+		}
+	}
+}

go.mod

@@ -2,13 +2,25 @@ module github.com/chdb-io/chdb-go
 
 go 1.21.5
 
-require github.com/c-bata/go-prompt v0.2.6
+require (
+	github.com/apache/arrow/go/v14 v14.0.2
+	github.com/c-bata/go-prompt v0.2.6
+)
 
 require (
-	github.com/mattn/go-colorable v0.1.7 // indirect
-	github.com/mattn/go-isatty v0.0.12 // indirect
+	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/google/flatbuffers v23.5.26+incompatible // indirect
+	github.com/klauspost/compress v1.16.7 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.5 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/mattn/go-runewidth v0.0.9 // indirect
 	github.com/mattn/go-tty v0.0.3 // indirect
+	github.com/pierrec/lz4/v4 v4.1.18 // indirect
 	github.com/pkg/term v1.2.0-beta.2 // indirect
-	golang.org/x/sys v0.0.0-20200918174421-af09f7315aff // indirect
+	github.com/zeebo/xxh3 v1.0.2 // indirect
+	golang.org/x/mod v0.13.0 // indirect
+	golang.org/x/sys v0.13.0 // indirect
+	golang.org/x/tools v0.14.0 // indirect
+	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 )