Merge branch 'main' into customNetListener

seanlafferty-ibm · seanlafferty-ibm · commit 2058262c7b24 · 2024-02-12T09:29:37.000-05:00
diff --git a/enginetest/queries/procedure_queries.go b/enginetest/queries/procedure_queries.go
@@ -1236,6 +1236,53 @@ END;`,
 			},
 		},
 	},
+	{
+		Name: "SQLEXCEPTION declare handler",
+		SetUpScript: []string{
+			`DROP TABLE IF EXISTS t1;`,
+			`CREATE TABLE t1 (pk BIGINT PRIMARY KEY);`,
+			`CREATE PROCEDURE eof()
+BEGIN
+	DECLARE a, b INT DEFAULT 1;
+    DECLARE cur1 CURSOR FOR SELECT * FROM t1;
+    OPEN cur1;
+    BEGIN
+		DECLARE EXIT HANDLER FOR SQLEXCEPTION SET a = 7;
+		tloop: LOOP
+			FETCH cur1 INTO b;
+            IF a > 1000 THEN
+				LEAVE tloop;
+            END IF;
+		END LOOP;
+    END;
+    CLOSE cur1;
+    SELECT a;
+END;`,
+			`CREATE PROCEDURE duplicate_key()
+BEGIN
+	DECLARE a, b INT DEFAULT 1;
+    BEGIN
+		DECLARE EXIT HANDLER FOR SQLEXCEPTION SET a = 7;
+		INSERT INTO t1 values (0);
+    END;
+    SELECT a;
+END;`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "CALL eof();",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "CALL duplicate_key();",
+				Expected: []sql.Row{{1}},
+			},
+			{
+				Query:    "CALL duplicate_key();",
+				Expected: []sql.Row{{7}},
+			},
+		},
+	},
 	{
 		Name: "DECLARE HANDLERs exit according to the block they were declared in",
 		SetUpScript: []string{
diff --git a/sql/expression/procedurereference.go b/sql/expression/procedurereference.go
@@ -15,6 +15,7 @@
 package expression
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 
@@ -47,8 +48,9 @@ type procedureCursorReferenceValue struct {
 type procedureHandlerReferenceValue struct {
 	Stmt        sql.Node
 	IsExit      bool
+	Action      DeclareHandlerAction
+	Cond        HandlerCondition
 	ScopeHeight int
-	//TODO: support more than just NOT FOUND
 }
 
 // ProcedureReferencable indicates that a sql.Node takes a *ProcedureReference returns a new copy with the reference set.
@@ -83,10 +85,11 @@ func (ppr *ProcedureReference) InitializeCursor(name string, selectStmt sql.Node
 }
 
 // InitializeHandler sets the given handler's statement.
-func (ppr *ProcedureReference) InitializeHandler(stmt sql.Node, returnsExitError bool) {
+func (ppr *ProcedureReference) InitializeHandler(stmt sql.Node, action DeclareHandlerAction, cond HandlerCondition) {
 	ppr.InnermostScope.Handlers = append(ppr.InnermostScope.Handlers, &procedureHandlerReferenceValue{
 		Stmt:        stmt,
-		IsExit:      returnsExitError,
+		Cond:        cond,
+		Action:      action,
 		ScopeHeight: ppr.height,
 	})
 }
@@ -393,12 +396,35 @@ func (upp *UnresolvedProcedureParam) WithChildren(children ...sql.Expression) (s
 	return upp, nil
 }
 
-// ProcedureBlockExitError contains the scope height that should exit.
-type ProcedureBlockExitError int
+// FetchEOF is a special EOF error that lets the loop implementation
+// differentiate between this io.EOF
+var FetchEOF = errors.New("exhausted fetch iterator")
 
-var _ error = ProcedureBlockExitError(0)
+type HandlerConditionType uint8
 
-// Error implements the error interface.
-func (b ProcedureBlockExitError) Error() string {
-	return "Block that EXIT handler was declared in could somehow not be found"
+const (
+	HandlerConditionUnknown HandlerConditionType = iota
+	HandlerConditionNotFound
+	HandlerConditionSqlException
+)
+
+type HandlerCondition struct {
+	SqlStatePrefix string
+	Type           HandlerConditionType
+}
+
+type DeclareHandlerAction byte
+
+const (
+	DeclareHandlerAction_Continue DeclareHandlerAction = iota
+	DeclareHandlerAction_Exit
+	DeclareHandlerAction_Undo
+)
+
+func (c *HandlerCondition) Matches(err error) bool {
+	if errors.Is(err, FetchEOF) {
+		return c.Type == HandlerConditionNotFound
+	} else {
+		return c.Type == HandlerConditionSqlException
+	}
 }
diff --git a/sql/plan/declare_handler.go b/sql/plan/declare_handler.go
@@ -21,20 +21,12 @@ import (
 	"github.com/dolthub/go-mysql-server/sql/expression"
 )
 
-type DeclareHandlerAction byte
-
-const (
-	DeclareHandlerAction_Continue DeclareHandlerAction = iota
-	DeclareHandlerAction_Exit
-	DeclareHandlerAction_Undo
-)
-
 // DeclareHandler represents the DECLARE ... HANDLER statement.
 type DeclareHandler struct {
-	Action    DeclareHandlerAction
+	Action    expression.DeclareHandlerAction
 	Statement sql.Node
 	Pref      *expression.ProcedureReference
-	//TODO: implement other conditions besides NOT FOUND
+	Condition expression.HandlerCondition
 }
 
 var _ sql.Node = (*DeclareHandler)(nil)
@@ -43,13 +35,14 @@ var _ sql.DebugStringer = (*DeclareHandler)(nil)
 var _ expression.ProcedureReferencable = (*DeclareHandler)(nil)
 
 // NewDeclareHandler returns a new *DeclareHandler node.
-func NewDeclareHandler(action DeclareHandlerAction, statement sql.Node) (*DeclareHandler, error) {
-	if action == DeclareHandlerAction_Undo {
+func NewDeclareHandler(action expression.DeclareHandlerAction, statement sql.Node, cond expression.HandlerCondition) (*DeclareHandler, error) {
+	if action == expression.DeclareHandlerAction_Undo {
 		return nil, sql.ErrDeclareHandlerUndo.New()
 	}
 	return &DeclareHandler{
 		Action:    action,
 		Statement: statement,
+		Condition: cond,
 	}, nil
 }
 
@@ -66,11 +59,11 @@ func (d *DeclareHandler) IsReadOnly() bool {
 func (d *DeclareHandler) String() string {
 	var action string
 	switch d.Action {
-	case DeclareHandlerAction_Continue:
+	case expression.DeclareHandlerAction_Continue:
 		action = "CONTINUE"
-	case DeclareHandlerAction_Exit:
+	case expression.DeclareHandlerAction_Exit:
 		action = "EXIT"
-	case DeclareHandlerAction_Undo:
+	case expression.DeclareHandlerAction_Undo:
 		action = "UNDO"
 	}
 	return fmt.Sprintf("DECLARE %s HANDLER FOR NOT FOUND %s", action, d.Statement.String())
@@ -80,11 +73,11 @@ func (d *DeclareHandler) String() string {
 func (d *DeclareHandler) DebugString() string {
 	var action string
 	switch d.Action {
-	case DeclareHandlerAction_Continue:
+	case expression.DeclareHandlerAction_Continue:
 		action = "CONTINUE"
-	case DeclareHandlerAction_Exit:
+	case expression.DeclareHandlerAction_Exit:
 		action = "EXIT"
-	case DeclareHandlerAction_Undo:
+	case expression.DeclareHandlerAction_Undo:
 		action = "UNDO"
 	}
 	return fmt.Sprintf("DECLARE %s HANDLER FOR NOT FOUND %s", action, sql.DebugString(d.Statement))
diff --git a/sql/planbuilder/ddl.go b/sql/planbuilder/ddl.go
@@ -997,6 +997,7 @@ func (b *Builder) buildExternalCreateIndex(inScope *scope, ddl *ast.DDL) (outSco
 		config,
 	)
 	createIndex.Catalog = b.cat
+	outScope.node = createIndex
 	return
 }
 
diff --git a/sql/planbuilder/proc.go b/sql/planbuilder/proc.go
@@ -40,7 +40,7 @@ const (
 
 type procCtx struct {
 	s          *scope
-	handler    *plan.DeclareHandler
+	handlers   []*plan.DeclareHandler
 	conditions map[string]*plan.DeclareCondition
 	vars       map[string]scopeColumn
 	cursors    map[string]struct{}
@@ -119,15 +119,11 @@ func (p *procCtx) HasCursor(name string) bool {
 
 func (p *procCtx) AddHandler(h *plan.DeclareHandler) {
 	p.NewState(dsHandler)
-	if p.handler != nil {
-		err := sql.ErrDeclareHandlerDuplicate.New()
-		p.s.b.handleErr(err)
-	}
-	p.handler = h
+	p.handlers = append(p.handlers, h)
 }
 
 func (p *procCtx) HasHandler(name string) bool {
-	return p.handler != nil
+	return p.handlers != nil
 }
 
 func (p *procCtx) AddLabel(label string, isLoop bool) {
@@ -347,33 +343,46 @@ func (b *Builder) buildDeclareCursor(inScope *scope, d *ast.Declare) (outScope *
 func (b *Builder) buildDeclareHandler(inScope *scope, d *ast.Declare, query string) (outScope *scope) {
 	outScope = inScope.push()
 	dHandler := d.Handler
-	//TODO: support other condition values besides NOT FOUND
-	if len(dHandler.ConditionValues) != 1 || dHandler.ConditionValues[0].ValueType != ast.DeclareHandlerCondition_NotFound {
+	if len(dHandler.ConditionValues) != 1 {
 		err := sql.ErrUnsupportedSyntax.New(ast.String(d))
 		b.handleErr(err)
 	}
+
+	var cond expression.HandlerCondition
+
+	switch dHandler.ConditionValues[0].ValueType {
+	case ast.DeclareHandlerCondition_NotFound:
+		cond = expression.HandlerCondition{Type: expression.HandlerConditionNotFound}
+	case ast.DeclareHandlerCondition_SqlException:
+		cond = expression.HandlerCondition{Type: expression.HandlerConditionSqlException}
+	default:
+		err := sql.ErrUnsupportedSyntax.New(ast.String(d))
+		b.handleErr(err)
+	}
+
 	stmtScope := b.build(inScope, dHandler.Statement, query)
 
-	var action plan.DeclareHandlerAction
+	var action expression.DeclareHandlerAction
 	switch dHandler.Action {
 	case ast.DeclareHandlerAction_Continue:
-		action = plan.DeclareHandlerAction_Continue
+		action = expression.DeclareHandlerAction_Continue
 	case ast.DeclareHandlerAction_Exit:
-		action = plan.DeclareHandlerAction_Exit
+		action = expression.DeclareHandlerAction_Exit
 	case ast.DeclareHandlerAction_Undo:
-		action = plan.DeclareHandlerAction_Undo
+		action = expression.DeclareHandlerAction_Undo
 	default:
 		err := fmt.Errorf("unknown DECLARE ... HANDLER action: %v", dHandler.Action)
 		b.handleErr(err)
 	}
-	if action == plan.DeclareHandlerAction_Undo {
+	if action == expression.DeclareHandlerAction_Undo {
 		err := sql.ErrDeclareHandlerUndo.New()
 		b.handleErr(err)
 	}
 
 	handler := &plan.DeclareHandler{
 		Action:    action,
 		Statement: stmtScope.node,
+		Condition: cond,
 	}
 
 	inScope.proc.AddHandler(handler)
diff --git a/sql/rowexec/other.go b/sql/rowexec/other.go
@@ -71,39 +71,7 @@ func (b *BaseBuilder) buildDeallocateQuery(ctx *sql.Context, n *plan.DeallocateQ
 func (b *BaseBuilder) buildFetch(ctx *sql.Context, n *plan.Fetch, row sql.Row) (sql.RowIter, error) {
 	row, sch, err := n.Pref.FetchCursor(ctx, n.Name)
 	if err == io.EOF {
-		scope := n.Pref.InnermostScope
-		for scope != nil {
-			for i := len(scope.Handlers) - 1; i >= 0; i-- {
-				//TODO: handle more than NOT FOUND handlers, handlers should check if the error applies to them first
-				originalScope := n.Pref.InnermostScope
-				defer func() {
-					n.Pref.InnermostScope = originalScope
-				}()
-				n.Pref.InnermostScope = scope
-				handlerRefVal := scope.Handlers[i]
-
-				handlerRowIter, err := b.buildNodeExec(ctx, handlerRefVal.Stmt, nil)
-				if err != nil {
-					return sql.RowsToRowIter(), err
-				}
-				defer handlerRowIter.Close(ctx)
-
-				for {
-					_, err := handlerRowIter.Next(ctx)
-					if err == io.EOF {
-						break
-					} else if err != nil {
-						return sql.RowsToRowIter(), err
-					}
-				}
-				if handlerRefVal.IsExit {
-					return sql.RowsToRowIter(), expression.ProcedureBlockExitError(handlerRefVal.ScopeHeight)
-				}
-				return sql.RowsToRowIter(), io.EOF
-			}
-			scope = scope.Parent
-		}
-		return sql.RowsToRowIter(), err
+		return sql.RowsToRowIter(), expression.FetchEOF
 	} else if err != nil {
 		return nil, err
 	}
diff --git a/sql/rowexec/proc.go b/sql/rowexec/proc.go
diff --git a/sql/rowexec/proc_iters.go b/sql/rowexec/proc_iters.go
diff --git a/sql/rowexec/rel_iters.go b/sql/rowexec/rel_iters.go

Original file line number	Diff line number	Diff line change
`@@ -997,6 +997,7 @@ func (b Builder) buildExternalCreateIndex(inScope scope, ddl *ast.DDL) (outSco`
`997`	`997`	`config,`
`998`	`998`	`)`
`999`	`999`	`createIndex.Catalog = b.cat`
	`1000`	`+ outScope.node = createIndex`
`1000`	`1001`	`return`
`1001`	`1002`	`}`
`1002`	`1003`