feat: support transaction (#102)

* feat: support transaction

* add begin, commit ,rollback

* fix serverinfo

* fix unit tests

* fix

* f

* settest suit

* fix tests

* add test for commit, rollback

* fix

* add tests

* fix

* fix

* fix commit

Author: hantmac

connection.go

@@ -25,14 +25,15 @@ const (
 )
 
 type DatabendConn struct {
-	ctx    context.Context
-	cfg    *Config
-	cancel context.CancelFunc
-	closed int32
-	stmts  []*databendStmt
-	logger *log.Logger
-	rest   *APIClient
-	commit func() error
+	ctx         context.Context
+	cfg         *Config
+	cancel      context.CancelFunc
+	closed      int32
+	stmts       []*databendStmt
+	logger      *log.Logger
+	rest        *APIClient
+	batchMode   bool
+	batchInsert func() error
 }
 
 func (dc *DatabendConn) exec(ctx context.Context, query string, args ...driver.Value) (driver.Result, error) {
@@ -95,11 +96,23 @@ func (dc *DatabendConn) query(ctx context.Context, query string, args ...driver.
 	return newNextRows(ctx, dc, r0)
 }
 
-//func (dc *DatabendConn) Begin() (driver.Tx, error) {
-//	return dc.BeginTx(dc.ctx, driver.TxOptions{})
-//}
+func (dc *DatabendConn) Begin() (driver.Tx, error) {
 
-func (dc *DatabendConn) Begin() (driver.Tx, error) { return dc, nil }
+	return dc.BeginTx(dc.ctx, driver.TxOptions{})
+}
+
+func (dc *DatabendConn) BeginTx(
+	ctx context.Context,
+	opts driver.TxOptions) (
+	driver.Tx, error) {
+	if dc.rest == nil {
+		return nil, driver.ErrBadConn
+	}
+	if _, err := dc.exec(ctx, "BEGIN"); err != nil {
+		return nil, err
+	}
+	return &databendTx{dc}, nil
+}
 
 func (dc *DatabendConn) cleanup() {
 	// must flush log buffer while the process is running.
@@ -128,7 +141,8 @@ func (dc *DatabendConn) prepare(ctx context.Context, query string) (*databendStm
 	if err != nil {
 		return nil, err
 	}
-	dc.commit = batch.BatchInsert
+	dc.batchInsert = batch.BatchInsert
+	dc.batchMode = true
 	stmt := &databendStmt{
 		dc:    dc,
 		query: query,
@@ -193,22 +207,15 @@ func (dc *DatabendConn) QueryContext(ctx context.Context, query string, args []d
 	return dc.query(ctx, query, values...)
 }
 
-// Commit applies prepared statement if it exists
-func (dc *DatabendConn) Commit() (err error) {
-	if dc.commit == nil {
+// ExecuteBatch applies batch prepared statement if it exists
+func (dc *DatabendConn) ExecuteBatch() (err error) {
+	if dc.batchInsert == nil {
 		return nil
 	}
 	defer func() {
-		dc.commit = nil
+		dc.batchInsert = nil
 	}()
-	return dc.commit()
-}
-
-// Rollback cleans prepared statement
-func (dc *DatabendConn) Rollback() error {
-	dc.commit = nil
-	dc.Close()
-	return nil
+	return dc.batchInsert()
 }
 
 // checkQueryID checks if query_id exists in context, if not, generate a new one

query.go

@@ -89,10 +89,20 @@ type SessionState struct {
 	// KeepServerSessionSecs uint64            `json:"keep_server_session_secs,omitempty"`
 
 	Settings map[string]string `json:"settings,omitempty"`
+
+	// txn
+	TxnState       string     `json:"txn_state,omitempty"`
+	LastServerInfo ServerInfo `json:"last_server_info,omitempty"`
+	LastQueryIds   []string   `json:"last_query_ids,omitempty"`
 }
 
 type StageAttachmentConfig struct {
 	Location          string            `json:"location"`
 	FileFormatOptions map[string]string `json:"file_format_options,omitempty"`
 	CopyOptions       map[string]string `json:"copy_options,omitempty"`
 }
+
+type ServerInfo struct {
+	Id        string `json:"id"`
+	StartTime string `json:"start_time"`
+}

query_test.go

@@ -17,7 +17,7 @@ func Test_SessionState(t *testing.T) {
 	}
 	buf, err := json.Marshal(ss)
 	require.NoError(t, err)
-	assert.Equal(t, `{"database":"db1"}`, string(buf))
+	assert.Equal(t, "{\"database\":\"db1\",\"last_server_info\":{\"id\":\"\",\"start_time\":\"\"}}", string(buf))
 
 	buf = []byte(`{"database":"db1", "secondary_roles": []}`)
 	err = json.Unmarshal(buf, ss)

stmt.go

@@ -1,32 +1,25 @@
 package godatabend
 
 import (
-	"context"
 	"database/sql/driver"
-	"regexp"
 
 	"github.com/pkg/errors"
 )
 
-var (
-	splitInsertRe = regexp.MustCompile(`(?si)(.+\s*VALUES)\s*(\(.+\))`)
-)
-
 type databendStmt struct {
-	dc        *DatabendConn
-	closed    int32
-	prefix    string
-	pattern   string
-	index     []int
-	batchMode bool
-	args      [][]driver.Value
-	query     string
-	batch     Batch
+	dc      *DatabendConn
+	closed  int32
+	prefix  string
+	pattern string
+	index   []int
+	args    [][]driver.Value
+	query   string
+	batch   Batch
 }
 
 func (stmt *databendStmt) Close() error {
 	logger.WithContext(stmt.dc.ctx).Infoln("Stmt.Close")
-	return nil
+	return stmt.dc.Close()
 }
 
 func (stmt *databendStmt) NumInput() int {
@@ -41,12 +34,6 @@ func (stmt *databendStmt) Exec(args []driver.Value) (driver.Result, error) {
 		return nil, err
 	}
 
-	//2. /v1/upload_to_stage csv file
-
-	// 3. copy into db.table from @~/csv
-
-	// 4. delete the file ?
-
 	return driver.RowsAffected(0), nil
 }
 
@@ -62,8 +49,3 @@ func (stmt *databendStmt) Query(args []driver.Value) (driver.Rows, error) {
 	logger.WithContext(stmt.dc.ctx).Infoln("Stmt.Query")
 	return nil, errors.New("only Exec method supported in batch mode")
 }
-
-func (stmt *databendStmt) commit(ctx context.Context) error {
-	logger.WithContext(stmt.dc.ctx).Infoln("Stmt Commit")
-	return stmt.batch.BatchInsert()
-}

tests/docker-compose.yml

@@ -7,7 +7,7 @@ services:
     volumes:
       - ./data:/data
   databend:
-    image: docker.io/datafuselabs/databend
+    image: datafuselabs/databend:v1.2.360-nightly
     environment:
       - QUERY_DEFAULT_USER=databend
       - QUERY_DEFAULT_PASSWORD=databend

tests/main_test.go

@@ -29,6 +29,15 @@ const (
 	createTable2 = `create table %s (a string);`
 )
 
+var (
+	dsn = "http://databend:databend@localhost:8000?presigned_url_disabled=true"
+)
+
+func init() {
+	dsn = os.Getenv("TEST_DATABEND_DSN")
+	//dsn = "http://databend:databend@localhost:8000?presigned_url_disabled=true"
+}
+
 func TestDatabendSuite(t *testing.T) {
 	suite.Run(t, new(DatabendTestSuite))
 }
@@ -44,7 +53,6 @@ type DatabendTestSuite struct {
 func (s *DatabendTestSuite) SetupSuite() {
 	var err error
 
-	dsn := os.Getenv("TEST_DATABEND_DSN")
 	s.NotEmpty(dsn)
 	s.db, err = sql.Open("databend", dsn)
 	s.Nil(err)
@@ -80,6 +88,7 @@ func (s *DatabendTestSuite) SetupTest() {
 
 func (s *DatabendTestSuite) TearDownTest() {
 	// t := s.T()
+	s.SetupSuite()
 
 	// t.Logf("teardown test with table %s", s.table)
 	_, err := s.db.Exec(fmt.Sprintf("DROP TABLE %s", s.table))
@@ -251,6 +260,52 @@ func (s *DatabendTestSuite) TestQueryNull() {
 	s.r.NoError(rows.Close())
 }
 
+func (s *DatabendTestSuite) TestTransactionCommit() {
+	tx, err := s.db.Begin()
+	s.r.Nil(err)
+
+	_, err = tx.Exec(fmt.Sprintf("INSERT INTO %s (i64) VALUES (?)", s.table), int64(1))
+	s.r.Nil(err)
+
+	err = tx.Commit()
+	s.r.Nil(err)
+
+	rows, err := s.db.Query(fmt.Sprintf("SELECT * FROM %s", s.table))
+	s.r.Nil(err)
+
+	result, err := scanValues(rows)
+	s.r.Nil(err)
+	s.r.Equal([][]interface{}{[]interface{}{"1", "NULL", "NULL", "NULL", "NULL", "NULL", "NULL", "NULL", "NULL"}}, result)
+
+	s.r.NoError(rows.Close())
+}
+
+func (s *DatabendTestSuite) TestTransactionRollback() {
+	tx, err := s.db.Begin()
+	s.r.Nil(err)
+
+	_, err = tx.Exec(fmt.Sprintf("INSERT INTO %s (i64) VALUES (?)", s.table), int64(1))
+	s.r.Nil(err)
+	rows, err := s.db.Query(fmt.Sprintf("SELECT * FROM %s", s.table))
+	s.r.Nil(err)
+
+	result, err := scanValues(rows)
+	s.r.Nil(err)
+	s.r.Equal([][]interface{}(nil), result)
+
+	err = tx.Rollback()
+	s.r.Nil(err)
+
+	rows, err = s.db.Query(fmt.Sprintf("SELECT * FROM %s", s.table))
+	s.r.Nil(err)
+
+	result, err = scanValues(rows)
+	s.r.Nil(err)
+	s.r.Empty(result)
+
+	s.r.NoError(rows.Close())
+}
+
 func scanValues(rows *sql.Rows) (interface{}, error) {
 	var err error
 	var result [][]interface{}

tests/session_test.go

@@ -9,6 +9,7 @@ import (
 )
 
 func (s *DatabendTestSuite) TestChangeDatabase() {
+	s.SetupSuite()
 	r := require.New(s.T())
 	var result string
 

transaction.go

@@ -0,0 +1,43 @@
+package godatabend
+
+import (
+	"database/sql/driver"
+)
+
+type databendTx struct {
+	dc *DatabendConn
+}
+
+func (tx *databendTx) Commit() (err error) {
+	if tx.dc == nil || tx.dc.rest == nil {
+		return driver.ErrBadConn
+	}
+	defer func() {
+		tx.dc.batchInsert = nil
+	}()
+	if tx.dc.batchMode && tx.dc.batchInsert != nil {
+		err = tx.dc.batchInsert()
+		if err != nil {
+			return
+		}
+	}
+	// complicity with old server version
+	if tx.dc.rest.sessionState.TxnState != "" {
+		_, err = tx.dc.exec(tx.dc.ctx, "COMMIT")
+		if err != nil {
+			return
+		}
+	}
+	return
+}
+
+func (tx *databendTx) Rollback() (err error) {
+	if tx.dc == nil || tx.dc.rest == nil {
+		return driver.ErrBadConn
+	}
+	_, err = tx.dc.exec(tx.dc.ctx, "ROLLBACK")
+	if err != nil {
+		return
+	}
+	return
+}