Merge pull request #125 from tursodatabase/with_schemadb

Make it possible to signal to the driver that we're connecting to schema db and use nontransactional batches in such case

Author: haaawk

libsql/internal/hrana/stream_request.go

@@ -38,7 +38,7 @@ func ExecuteStoredStream(sqlId int32, params shared.Params, wantRows bool) (*Str
 	return &StreamRequest{Type: "execute", Stmt: stmt}, nil
 }
 
-func BatchStream(sqls []string, params []shared.Params, wantRows bool) (*StreamRequest, error) {
+func BatchStream(sqls []string, params []shared.Params, wantRows bool, transactional bool) (*StreamRequest, error) {
 	batch := &Batch{}
 	for idx, sql := range sqls {
 		s := sql
@@ -50,19 +50,23 @@ func BatchStream(sqls []string, params []shared.Params, wantRows bool) (*StreamR
 			return nil, err
 		}
 		var condition *BatchCondition
-		if idx > 0 {
-			prev_idx := int32(idx - 1)
-			condition = &BatchCondition{
-				Type: "ok",
-				Step: &prev_idx,
+		if transactional {
+			if idx > 0 {
+				prev_idx := int32(idx - 1)
+				condition = &BatchCondition{
+					Type: "ok",
+					Step: &prev_idx,
+				}
 			}
 		}
 		batch.Add(*stmt, condition)
 	}
-	rollback := "ROLLBACK"
-	last_idx := int32(len(sqls) - 1)
-	batch.Add(Stmt{Sql: &rollback, WantRows: false},
-		&BatchCondition{Type: "not", Cond: &BatchCondition{Type: "ok", Step: &last_idx}})
+	if transactional {
+		rollback := "ROLLBACK"
+		last_idx := int32(len(sqls) - 1)
+		batch.Add(Stmt{Sql: &rollback, WantRows: false},
+			&BatchCondition{Type: "not", Cond: &BatchCondition{Type: "ok", Step: &last_idx}})
+	}
 	return &StreamRequest{Type: "batch", Batch: batch}, nil
 }
 

libsql/internal/http/driver.go

@@ -6,6 +6,6 @@ import (
 	"github.com/tursodatabase/libsql-client-go/libsql/internal/http/hranaV2"
 )
 
-func Connect(url, jwt, host string) driver.Conn {
-	return hranaV2.Connect(url, jwt, host)
+func Connect(url, jwt, host string, schemaDb bool) driver.Conn {
+	return hranaV2.Connect(url, jwt, host, schemaDb)
 }

libsql/internal/http/hranaV2/hranaV2.go

@@ -36,8 +36,8 @@ func init() {
 	commitHash = "unknown"
 }
 
-func Connect(url, jwt, host string) driver.Conn {
-	return &hranaV2Conn{url, jwt, host, "", false, 0}
+func Connect(url, jwt, host string, schemaDb bool) driver.Conn {
+	return &hranaV2Conn{url, jwt, host, schemaDb, "", false, 0}
 }
 
 type hranaV2Stmt struct {
@@ -85,6 +85,7 @@ type hranaV2Conn struct {
 	url              string
 	jwt              string
 	host             string
+	schemaDb         bool
 	baton            string
 	streamClosed     bool
 	replicationIndex uint64
@@ -298,7 +299,7 @@ func (h *hranaV2Conn) executeStmt(ctx context.Context, query string, args []driv
 		}
 		msg.Add(*executeStream)
 	} else {
-		batchStream, err := hrana.BatchStream(stmts, params, wantRows)
+		batchStream, err := hrana.BatchStream(stmts, params, wantRows, !h.schemaDb)
 		if err != nil {
 			return nil, fmt.Errorf("failed to execute SQL: %s\n%w", query, err)
 		}
@@ -338,7 +339,11 @@ func (h *hranaV2Conn) ExecContext(ctx context.Context, query string, args []driv
 		}
 		lastInsertRowId := int64(0)
 		affectedRowCount := int64(0)
-		for idx := 0; idx < len(res.StepResults)-1; idx++ {
+		upperBound := len(res.StepResults)
+		if !h.schemaDb {
+			upperBound -= 1
+		}
+		for idx := 0; idx < upperBound; idx++ {
 			r := res.StepResults[idx]
 			rowId := r.GetLastInsertRowId()
 			if rowId > 0 {
@@ -458,8 +463,10 @@ func (h *hranaV2Conn) QueryContext(ctx context.Context, query string, args []dri
 		if err != nil {
 			return nil, err
 		}
-		res.StepResults = res.StepResults[:len(res.StepResults)-1]
-		res.StepErrors = res.StepErrors[:len(res.StepErrors)-1]
+		if !h.schemaDb {
+			res.StepResults = res.StepResults[:len(res.StepResults)-1]
+			res.StepErrors = res.StepErrors[:len(res.StepErrors)-1]
+		}
 		return shared.NewRows(&BatchResultRowsProvider{res}), nil
 	default:
 		return nil, fmt.Errorf("failed to execute SQL: %s\n%s", query, "unknown response type")

libsql/sql.go

@@ -17,6 +17,7 @@ type config struct {
 	authToken *string
 	tls       *bool
 	proxy     *string
+	schemaDb  *bool
 }
 
 type Option interface {
@@ -65,6 +66,16 @@ func WithProxy(proxy string) Option {
 	})
 }
 
+func WithSchemaDb(schemaDb bool) Option {
+	return option(func(o *config) error {
+		if o.tls != nil {
+			return fmt.Errorf("schemaDb already set")
+		}
+		o.schemaDb = &schemaDb
+		return nil
+	})
+}
+
 func (c config) connector(dbPath string) (driver.Connector, error) {
 	u, err := url.Parse(dbPath)
 	if err != nil {
@@ -144,11 +155,16 @@ func (c config) connector(dbPath string) (driver.Connector, error) {
 		}
 	}
 
+	schemaDb := false
+	if c.schemaDb != nil {
+		schemaDb = *c.schemaDb
+	}
+
 	if u.Scheme == "wss" || u.Scheme == "ws" {
 		return wsConnector{url: u.String(), authToken: authToken}, nil
 	}
 	if u.Scheme == "https" || u.Scheme == "http" {
-		return httpConnector{url: u.String(), authToken: authToken, host: host}, nil
+		return httpConnector{url: u.String(), authToken: authToken, host: host, schemaDb: schemaDb}, nil
 	}
 
 	return nil, fmt.Errorf("unsupported URL scheme: %s\nThis driver supports only URLs that start with libsql://, file://, https://, http://, wss:// and ws://", u.Scheme)
@@ -172,10 +188,11 @@ type httpConnector struct {
 	url       string
 	authToken string
 	host      string
+	schemaDb  bool
 }
 
 func (c httpConnector) Connect(_ctx context.Context) (driver.Conn, error) {
-	return http.Connect(c.url, c.authToken, c.host), nil
+	return http.Connect(c.url, c.authToken, c.host, c.schemaDb), nil
 }
 
 func (c httpConnector) Driver() driver.Driver {
@@ -324,7 +341,7 @@ func (d Driver) Open(dbUrl string) (driver.Conn, error) {
 		return ws.Connect(u.String(), jwt)
 	}
 	if u.Scheme == "https" || u.Scheme == "http" {
-		return http.Connect(u.String(), jwt, u.Host), nil
+		return http.Connect(u.String(), jwt, u.Host, false), nil
 	}
 
 	return nil, fmt.Errorf("unsupported URL scheme: %s\nThis driver supports only URLs that start with libsql://, file://, https://, http://, wss:// and ws://", u.Scheme)