fix: support release savepoint (#251)

pkg/executor/read_write_splitting.go

@@ -212,13 +212,26 @@ func (executor *ReadWriteSplittingExecutor) ExecutorComQuery(
 		if !ok {
 			return nil, 0, errors.New("there is no transaction")
 		}
-		defer executor.localTransactionMap.Delete(connectionID)
+		if stmt.SavepointName == "" {
+			defer executor.localTransactionMap.Delete(connectionID)
+		}
 		tx = txi.(proto.Tx)
 		// TODO add metrics
 		if result, err = tx.Rollback(spanCtx, stmt); err != nil {
 			return nil, 0, err
 		}
 		return result, 0, err
+	case *ast.ReleaseSavepointStmt:
+		txi, ok := executor.localTransactionMap.Load(connectionID)
+		if !ok {
+			return nil, 0, errors.New("there is no transaction")
+		}
+		defer executor.localTransactionMap.Delete(connectionID)
+		tx = txi.(proto.Tx)
+		if result, err = tx.ReleaseSavepoint(spanCtx, stmt.Name); err != nil {
+			return nil, 0, err
+		}
+		return result, 0, err
 	case *ast.InsertStmt, *ast.DeleteStmt, *ast.UpdateStmt:
 		txi, ok := executor.localTransactionMap.Load(connectionID)
 		if ok {
@@ -334,5 +347,5 @@ func (executor *ReadWriteSplittingExecutor) doPostFilter(ctx context.Context, re
 			return err
 		}
 	}
-	return nil
+	return err
 }

pkg/executor/sharding.go

@@ -291,5 +291,5 @@ func (executor *ShardingExecutor) doPostFilter(ctx context.Context, result proto
 			return err
 		}
 	}
-	return nil
+	return err
 }

pkg/executor/single_db.go

@@ -207,13 +207,26 @@ func (executor *SingleDBExecutor) ExecutorComQuery(
 		if !ok {
 			return nil, 0, errors.New("there is no transaction")
 		}
-		defer executor.localTransactionMap.Delete(connectionID)
+		if stmt.SavepointName == "" {
+			defer executor.localTransactionMap.Delete(connectionID)
+		}
 		tx = txi.(proto.Tx)
 		// TODO add metrics
 		if result, err = tx.Rollback(spanCtx, stmt); err != nil {
 			return nil, 0, err
 		}
 		return result, 0, err
+	case *ast.ReleaseSavepointStmt:
+		txi, ok := executor.localTransactionMap.Load(connectionID)
+		if !ok {
+			return nil, 0, errors.New("there is no transaction")
+		}
+		defer executor.localTransactionMap.Delete(connectionID)
+		tx = txi.(proto.Tx)
+		if result, err = tx.ReleaseSavepoint(spanCtx, stmt.Name); err != nil {
+			return nil, 0, err
+		}
+		return result, 0, err
 	default:
 		txi, ok := executor.localTransactionMap.Load(connectionID)
 		if ok {
@@ -285,5 +298,5 @@ func (executor *SingleDBExecutor) doPostFilter(ctx context.Context, result proto
 			return err
 		}
 	}
-	return nil
+	return err
 }

pkg/proto/interface.go

@@ -173,6 +173,7 @@ type (
 		ExecuteSqlDirectly(sql string, args ...interface{}) (Result, uint16, error)
 		Commit(ctx context.Context) (Result, error)
 		Rollback(ctx context.Context, stmt *ast.RollbackStmt) (Result, error)
+		ReleaseSavepoint(ctx context.Context, savepoint string) (result Result, err error)
 	}
 
 	DBManager interface {

pkg/sql/tx.go

@@ -165,9 +165,25 @@ func (tx *Tx) Rollback(ctx context.Context, stmt *ast.RollbackStmt) (result prot
 		result, err = tx.conn.Execute(ctx, fmt.Sprintf("ROLLBACK TO %s", stmt.SavepointName), false)
 	} else {
 		result, err = tx.conn.Execute(ctx, "ROLLBACK", false)
+		tx.db.pool.Put(tx.conn)
+		tx.Close()
 	}
+	return
+}
+
+func (tx *Tx) ReleaseSavepoint(ctx context.Context, savepoint string) (result proto.Result, err error) {
+	_, span := tracing.GetTraceSpan(ctx, tracing.TxReleaseSavePoint)
+	span.SetAttributes(attribute.KeyValue{Key: "db", Value: attribute.StringValue(tx.db.name)})
+	defer span.End()
+
+	if tx.closed.Load() {
+		return nil, nil
+	}
+	if tx.db == nil || tx.db.IsClosed() {
+		return nil, err2.ErrInvalidConn
+	}
+	result, err = tx.conn.Execute(ctx, fmt.Sprintf("RELEASE SAVEPOINT %s", savepoint), false)
 	tx.db.pool.Put(tx.conn)
-	tx.Close()
 	return
 }
 

pkg/tracing/constant.go

@@ -68,11 +68,12 @@ const (
 	DBLocalTransactionBegin = "db_local_transaction_begin"
 
 	// tx
-	TxQuery    = "tx_query"
-	TxExecSQL  = "tx_exec_sql"
-	TxExecStmt = "tx_exec_stmt"
-	TxCommit   = "db_local_transaction_commit"
-	TxRollback = "db_local_transaction_rollback"
+	TxQuery            = "tx_query"
+	TxExecSQL          = "tx_exec_sql"
+	TxExecStmt         = "tx_exec_stmt"
+	TxCommit           = "db_local_transaction_commit"
+	TxRollback         = "db_local_transaction_rollback"
+	TxReleaseSavePoint = "db_local_transaction_release_savepoint"
 
 	// conn
 	ConnQuery       = "conn_com_query"