Skip to content

Commit cab2314

Browse files
general-kroll-4-lifegeneral-kroll-4-life
authored
update-replace-returning (#555)
Summary: - Support for `UPDATE RETURNING`. - Support for `REPLACE RETURNING`. - Added robot test `Update Returning Simple Projection`. - Added robot test `Update Async Returning Simple Projection`. - Added robot test `Replace Returning Simple Projection`. - Added robot test `Replace Async Returning Simple Projection`.
1 parent 414a233 commit cab2314

File tree

16 files changed

+510
-103
lines changed

16 files changed

+510
-103
lines changed

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ require (
1919
github.com/spf13/cobra v1.4.0
2020
github.com/spf13/pflag v1.0.5
2121
github.com/spf13/viper v1.10.1
22-
github.com/stackql/any-sdk v0.1.4-alpha11
22+
github.com/stackql/any-sdk v0.1.4-alpha12
2323
github.com/stackql/go-suffix-map v0.0.1-alpha01
2424
github.com/stackql/psql-wire v0.1.1-beta23
2525
github.com/stackql/stackql-parser v0.0.15-alpha06

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,8 +484,8 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
484484
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
485485
github.com/spf13/viper v1.10.1 h1:nuJZuYpG7gTj/XqiUwg8bA0cp1+M2mC3J4g5luUYBKk=
486486
github.com/spf13/viper v1.10.1/go.mod h1:IGlFPqhNAPKRxohIzWpI5QEy4kuI7tcl5WvR+8qy1rU=
487-
github.com/stackql/any-sdk v0.1.4-alpha11 h1:WeFvLcml37sq+HbVk3hiXuImQ578SE9LM/K1MzqFovo=
488-
github.com/stackql/any-sdk v0.1.4-alpha11/go.mod h1:AKS/g28y7m4SWL/YW8veE9MCNy8XJgaicVibemVE9e8=
487+
github.com/stackql/any-sdk v0.1.4-alpha12 h1:SozoNb4+7+y+0bVGmO8ozFtTTpYhEWBHVDuN9SmkFuA=
488+
github.com/stackql/any-sdk v0.1.4-alpha12/go.mod h1:AKS/g28y7m4SWL/YW8veE9MCNy8XJgaicVibemVE9e8=
489489
github.com/stackql/go-suffix-map v0.0.1-alpha01 h1:TDUDS8bySu41Oo9p0eniUeCm43mnRM6zFEd6j6VUaz8=
490490
github.com/stackql/go-suffix-map v0.0.1-alpha01/go.mod h1:QAi+SKukOyf4dBtWy8UMy+hsXXV+yyEE4vmBkji2V7g=
491491
github.com/stackql/psql-wire v0.1.1-beta23 h1:1ayYMjZArfDcIMyEOKnm+Bp1zRCISw8pguvTFuUhhVQ=

‎internal/stackql/parserutil/parser_util.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,26 @@ func ExtractInsertReturningColumnNames(
108108
return colNames, err
109109
}
110110

111+
func ExtractUpdateReturningColumnNames(
112+
updateStmt *sqlparser.Update,
113+
formatter sqlparser.NodeFormatter,
114+
) ([]ColumnHandle, error) {
115+
var colNames []ColumnHandle
116+
var err error
117+
for _, node := range updateStmt.SelectExprs {
118+
switch node := node.(type) {
119+
case *sqlparser.AliasedExpr:
120+
cn, cErr := inferColNameFromExpr(node.Expr, formatter, node.As.GetRawVal())
121+
if cErr != nil {
122+
return nil, cErr
123+
}
124+
colNames = append(colNames, cn)
125+
case *sqlparser.StarExpr:
126+
}
127+
}
128+
return colNames, err
129+
}
130+
111131
func ExtractInsertColumnNames(insertStmt *sqlparser.Insert) ([]string, error) {
112132
var colNames []string
113133
var err error

‎internal/stackql/planbuilder/plan_builder.go‎

Lines changed: 107 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -949,75 +949,110 @@ func (pgb *standardPlanGraphBuilder) handleInsert(pbi planbuilderinput.PlanBuild
949949
if isPhysicalTable {
950950
bldrInput.SetIsTargetPhysicalTable(true)
951951
}
952-
var bldr primitivebuilder.Builder
953-
if len(node.SelectExprs) > 0 {
954-
// Two cases:
955-
// 1. Synchronous. Equivalent to select.
956-
// 2. Asynchronous. Whole other story.
957-
tableMeta, tableMetaExists := bldrInput.GetTableMetadata()
958-
if !tableMetaExists {
959-
return fmt.Errorf("could not obtain table metadata for node '%s'", node.Action)
960-
}
961-
rc, rcErr := tableinsertioncontainer.NewTableInsertionContainer(
962-
tableMeta,
963-
handlerCtx.GetSQLEngine(),
964-
handlerCtx.GetTxnCounterMgr(),
952+
return pgb.handleMutationOperation(
953+
handlerCtx,
954+
pbi,
955+
primitiveGenerator,
956+
node,
957+
tbl,
958+
selectPrimitiveNode,
959+
bldrInput,
960+
isAwait,
961+
)
962+
}
963+
pr := primitive.NewGenericPrimitive(nil, nil, nil, primitive_context.NewPrimitiveContext())
964+
pgb.planGraphHolder.CreatePrimitiveNode(pr)
965+
return nil
966+
}
967+
968+
func (pgb *standardPlanGraphBuilder) handleMutationOperation(
969+
handlerCtx handler.HandlerContext,
970+
pbi planbuilderinput.PlanBuilderInput,
971+
primitiveGenerator primitivegenerator.PrimitiveGenerator,
972+
node sqlparser.SQLNode,
973+
tbl tablemetadata.ExtendedTableMetadata,
974+
selectPrimitiveNode primitivegraph.PrimitiveNode,
975+
bldrInput builder_input.BuilderInput,
976+
isAwait bool,
977+
) error {
978+
var returningExpressions sqlparser.SelectExprs
979+
var inputAction string
980+
var bldr primitivebuilder.Builder
981+
switch n := node.(type) {
982+
case *sqlparser.Insert:
983+
returningExpressions = n.SelectExprs
984+
inputAction = n.Action
985+
case *sqlparser.Update:
986+
returningExpressions = n.SelectExprs
987+
inputAction = n.Action
988+
default:
989+
return fmt.Errorf("unsupported mutation operation of type '%T'", node)
990+
}
991+
//nolint:nestif // acceptable complexity
992+
if len(returningExpressions) > 0 {
993+
// Two cases:
994+
// 1. Synchronous. Equivalent to select.
995+
// 2. Asynchronous. Whole other story.
996+
tableMeta, tableMetaExists := bldrInput.GetTableMetadata()
997+
if !tableMetaExists {
998+
return fmt.Errorf("could not obtain table metadata for node '%s'", inputAction)
999+
}
1000+
rc, rcErr := tableinsertioncontainer.NewTableInsertionContainer(
1001+
tableMeta,
1002+
handlerCtx.GetSQLEngine(),
1003+
handlerCtx.GetTxnCounterMgr(),
1004+
)
1005+
if rcErr != nil {
1006+
return rcErr
1007+
}
1008+
bldrInput.SetTableInsertionContainer(rc)
1009+
bldrInput.SetIsReturning(true)
1010+
if !isAwait {
1011+
bldr = primitivebuilder.NewSingleAcquireAndSelect(
1012+
bldrInput,
1013+
primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx(),
1014+
primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
1015+
nil,
9651016
)
966-
if rcErr != nil {
967-
return rcErr
968-
}
969-
bldrInput.SetTableInsertionContainer(rc)
970-
bldrInput.SetIsReturning(true)
971-
if !isAwait {
972-
bldr = primitivebuilder.NewSingleAcquireAndSelect(
973-
bldrInput,
974-
primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx(),
975-
primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
976-
nil,
977-
)
978-
} else {
979-
bldrInput.SetIsAwait(true)
980-
bldrInput.SetIsReturning(true)
981-
bldrInput.SetInsertCtx(primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx())
982-
lhsBldr := primitivebuilder.NewInsertOrUpdate(
983-
bldrInput,
984-
)
985-
newBldrInput := builder_input.NewBuilderInput(
986-
pgb.planGraphHolder,
987-
handlerCtx,
988-
tbl,
989-
)
990-
newBldrInput.SetParserNode(node)
991-
newBldrInput.SetAnnotatedAST(pbi.GetAnnotatedAST())
992-
newBldrInput.SetTxnCtrlCtrs(pbi.GetTxnCtrlCtrs())
993-
newBldrInput.SetTableInsertionContainer(rc)
994-
newBldrInput.SetDependencyNode(selectPrimitiveNode)
995-
newBldrInput.SetIsAwait(isAwait)
996-
rhsBldr := primitivebuilder.NewSingleSelect(
997-
pgb.planGraphHolder, handlerCtx, primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
998-
[]tableinsertioncontainer.TableInsertionContainer{rc},
999-
nil,
1000-
streaming.NewNopMapStream(),
1001-
)
1002-
bldr = primitivebuilder.NewDependencySubDAGBuilder(
1003-
pgb.planGraphHolder,
1004-
[]primitivebuilder.Builder{lhsBldr},
1005-
rhsBldr,
1006-
)
1007-
}
10081017
} else {
1009-
bldr = primitivebuilder.NewInsertOrUpdate(
1018+
bldrInput.SetIsAwait(true)
1019+
bldrInput.SetIsReturning(true)
1020+
bldrInput.SetInsertCtx(primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx())
1021+
lhsBldr := primitivebuilder.NewInsertOrUpdate(
10101022
bldrInput,
10111023
)
1024+
newBldrInput := builder_input.NewBuilderInput(
1025+
pgb.planGraphHolder,
1026+
handlerCtx,
1027+
tbl,
1028+
)
1029+
newBldrInput.SetParserNode(node)
1030+
newBldrInput.SetAnnotatedAST(pbi.GetAnnotatedAST())
1031+
newBldrInput.SetTxnCtrlCtrs(pbi.GetTxnCtrlCtrs())
1032+
newBldrInput.SetTableInsertionContainer(rc)
1033+
newBldrInput.SetDependencyNode(selectPrimitiveNode)
1034+
newBldrInput.SetIsAwait(isAwait)
1035+
rhsBldr := primitivebuilder.NewSingleSelect(
1036+
pgb.planGraphHolder, handlerCtx, primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
1037+
[]tableinsertioncontainer.TableInsertionContainer{rc},
1038+
nil,
1039+
streaming.NewNopMapStream(),
1040+
)
1041+
bldr = primitivebuilder.NewDependencySubDAGBuilder(
1042+
pgb.planGraphHolder,
1043+
[]primitivebuilder.Builder{lhsBldr},
1044+
rhsBldr,
1045+
)
10121046
}
1013-
err = bldr.Build()
1014-
if err != nil {
1015-
return err
1016-
}
1017-
return nil
1047+
} else {
1048+
bldr = primitivebuilder.NewInsertOrUpdate(
1049+
bldrInput,
1050+
)
1051+
}
1052+
err := bldr.Build()
1053+
if err != nil {
1054+
return err
10181055
}
1019-
pr := primitive.NewGenericPrimitive(nil, nil, nil, primitive_context.NewPrimitiveContext())
1020-
pgb.planGraphHolder.CreatePrimitiveNode(pr)
10211056
return nil
10221057
}
10231058

@@ -1072,14 +1107,17 @@ func (pgb *standardPlanGraphBuilder) handleUpdate(pbi planbuilderinput.PlanBuild
10721107
bldrInput.SetTxnCtrlCtrs(pbi.GetTxnCtrlCtrs())
10731108
bldrInput.SetIsTargetPhysicalTable(true)
10741109
}
1075-
bldr := primitivebuilder.NewInsertOrUpdate(
1110+
isAwait := primitiveGenerator.GetPrimitiveComposer().IsAwait()
1111+
return pgb.handleMutationOperation(
1112+
handlerCtx,
1113+
pbi,
1114+
primitiveGenerator,
1115+
node,
1116+
tbl,
1117+
selectPrimitiveNode,
10761118
bldrInput,
1119+
isAwait,
10771120
)
1078-
err = bldr.Build()
1079-
if err != nil {
1080-
return err
1081-
}
1082-
return nil
10831121
}
10841122
pr := primitive.NewGenericPrimitive(nil, nil, nil, primitive_context.NewPrimitiveContext())
10851123
pgb.planGraphHolder.CreatePrimitiveNode(pr)

‎internal/stackql/primitivebuilder/exec.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ func (ss *Exec) Build() error {
8787
svc,
8888
true,
8989
[]anysdk.ColumnDescriptor{},
90+
false,
9091
)
9192
analyser := anysdk.NewMethodAnalyzer()
9293
methodAnalysisOutput, analysisErr := analyser.AnalyzeUnaryAction(analysisInput)

‎internal/stackql/primitivebuilder/insert_or_update.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,9 @@ func (ss *insertOrUpdate) Build() error {
4444
}
4545
case *sqlparser.Update:
4646
mutableInput.SetVerb("update")
47+
if len(node.SelectExprs) > 0 {
48+
mutableInput.SetIsReturning(true)
49+
}
4750
default:
4851
return fmt.Errorf("mutation executor: cannnot accomodate node of type '%T'", node)
4952
}

0 commit comments

Comments
 (0)