insert-returning (#552)

Summary:

- Support for `INSERT RETURNING`.
- Documented expected behaviour.
- Working for synchronously created objects **only**.
- Added robot test `Insert Returning Simple Projection`.

Author: general-kroll-4-life

.vscode/launch.json

@@ -175,6 +175,7 @@
         "select BackupId, BackupState from aws.cloudhsm.backups where region = 'rubbish-region' order by BackupId;",
         "select  rings.projectsId as project,  rings.locationsId as locale,  split_part(rings.name, '/', -1) as key_ring_name,  split_part(keys.name, '/', -1) as key_name,  json_extract(keys.\"versionTemplate\", '$.algorithm') as key_algorithm,  json_extract(keys.\"versionTemplate\", '$.protectionLevel') as key_protection_level  from  google.cloudkms.key_rings rings  inner join google.cloudkms.crypto_keys keys  on  keys.keyRingsId = split_part(rings.name, '/', -1)  and keys.projectsId = rings.projectsId  and keys.locationsId = rings.locationsId  where  rings.projectsId in ('testing-project', 'testing-project-two', 'testing-project-three')  and rings.locationsId in ('global', 'australia-southeast1', 'australia-southeast2')  order by project, locale, key_name ;",
         "delete from aws.cloud_control.resources  where  region = 'ap-southeast-1'  and data__TypeName = 'AWS::Logs::LogGroup'  and data__Identifier = 'LogGroupResourceExampleThird' ;",
+        "insert into google.storage.buckets( project, data__name) select  'testing-project', 'silly-bucket' returning projectNumber;",
       ],
       "default": "show providers;"
     },

docs/developer_guide.md

@@ -268,3 +268,13 @@ time ./stackql exec --cpuprofile=./select-disks-improved-05.profile --auth='{ "g
 ## AWS HTTP request signing
 
 https://docs.aws.amazon.com/sdk-for-go/api/aws/signer/v4/
+
+## Selection from non-select DML
+
+`INSERT RETURNING` can function in two mechanisms:
+
+- Synchronous responses, such as [`google.storage.buckets`](https://cloud.google.com/storage/docs/json_api/v1/buckets/insert).  The returning clause is a projection on the immediately available reponse body.
+- Asynchronous responses, such as [`google.compute.instances`](https://cloud.google.com/compute/docs/reference/rest/v1/instances/insert).  The returning clause is a projection on the reponse body **after** the await flow has concluded.
+
+Future use cases for `UPDATE RETURNING`, `REPLACE RETURNING` and `DELETE RETURNING` will function the same observable fashion.
+

go.mod

@@ -22,7 +22,7 @@ require (
 	github.com/stackql/any-sdk v0.1.3-beta02
 	github.com/stackql/go-suffix-map v0.0.1-alpha01
 	github.com/stackql/psql-wire v0.1.1-beta23
-	github.com/stackql/stackql-parser v0.0.14-alpha05
+	github.com/stackql/stackql-parser v0.0.15-alpha06
 	github.com/stretchr/testify v1.10.0
 	golang.org/x/sync v0.10.0
 	gonum.org/v1/gonum v0.11.0

go.sum

@@ -494,8 +494,8 @@ github.com/stackql/readline v0.0.2-alpha05 h1:ID4QzGdplFBsrSnTuz8pvKzWw96JbrJg8f
 github.com/stackql/readline v0.0.2-alpha05/go.mod h1:OFAYOdXk/X4+5GYiDXFfaGrk+bCN6Qv0SYY5HNzD2E0=
 github.com/stackql/stackql-go-sqlite3 v1.0.3-stackql h1:j0yt6T5thZuz5+HIr81PXz2AClAtCko0vzr5tm8C1g8=
 github.com/stackql/stackql-go-sqlite3 v1.0.3-stackql/go.mod h1:HemqCrcMK2xyhMMMt6oZ7ERDtoSmyyDsw5LBBcTZ+Rk=
-github.com/stackql/stackql-parser v0.0.14-alpha05 h1:DLcsaeTypH5p1T+g/9NqaGYdC9uIQ+7pZetnkjHi0G0=
-github.com/stackql/stackql-parser v0.0.14-alpha05/go.mod h1:iyB47SvRS+Fvpn7joF7mHAkeiWSq83TbUhglRmLzPLQ=
+github.com/stackql/stackql-parser v0.0.15-alpha06 h1:bdaudybbEmrR9m88CO1HpX/eKda2+gjiy4Kf6ZjwPoQ=
+github.com/stackql/stackql-parser v0.0.15-alpha06/go.mod h1:iyB47SvRS+Fvpn7joF7mHAkeiWSq83TbUhglRmLzPLQ=
 github.com/stackql/stackql-provider-registry v0.0.1-rc06 h1:MgroWOr0bSqjSTDGnXB0UoZGFXpW3SRtN0EFkzB8Rpo=
 github.com/stackql/stackql-provider-registry v0.0.1-rc06/go.mod h1:87rVxnS2aRASK20lBQgoYA0o7FSJTZBGGRaWFR7IDm4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

internal/stackql/internal_data_transfer/builder_input/builder_input.go

@@ -9,6 +9,7 @@ import (
 	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/internaldto"
 	"github.com/stackql/stackql/internal/stackql/primitivegraph"
 	"github.com/stackql/stackql/internal/stackql/provider"
+	"github.com/stackql/stackql/internal/stackql/tableinsertioncontainer"
 	"github.com/stackql/stackql/internal/stackql/tablemetadata"
 )
 
@@ -51,27 +52,30 @@ type BuilderInput interface {
 	SetIsTargetPhysicalTable(isPhysical bool)
 	SetTxnCtrlCtrs(internaldto.TxnControlCounters)
 	GetTxnCtrlCtrs() (internaldto.TxnControlCounters, bool)
+	GetTableInsertionContainer() (tableinsertioncontainer.TableInsertionContainer, bool)
+	SetTableInsertionContainer(tableinsertioncontainer.TableInsertionContainer)
 }
 
 type builderInput struct {
-	graphHolder       primitivegraph.PrimitiveGraphHolder
-	handlerCtx        handler.HandlerContext
-	paramMap          map[int]map[string]interface{}
-	tbl               tablemetadata.ExtendedTableMetadata
-	dependencyNode    primitivegraph.PrimitiveNode
-	commentDirectives sqlparser.CommentDirectives
-	isAwait           bool
-	verb              string
-	inputAlias        string
-	isUndo            bool
-	node              sqlparser.SQLNode
-	paramMapStream    streaming.MapStream
-	httpPrepStream    anysdk.HttpPreparatorStream
-	op                anysdk.OperationStore
-	prov              provider.IProvider
-	annotatedAst      annotatedast.AnnotatedAst
-	isTargetPhysical  bool
-	txnCtrlCtrs       internaldto.TxnControlCounters
+	graphHolder             primitivegraph.PrimitiveGraphHolder
+	handlerCtx              handler.HandlerContext
+	paramMap                map[int]map[string]interface{}
+	tbl                     tablemetadata.ExtendedTableMetadata
+	dependencyNode          primitivegraph.PrimitiveNode
+	commentDirectives       sqlparser.CommentDirectives
+	isAwait                 bool
+	verb                    string
+	inputAlias              string
+	isUndo                  bool
+	node                    sqlparser.SQLNode
+	paramMapStream          streaming.MapStream
+	httpPrepStream          anysdk.HttpPreparatorStream
+	op                      anysdk.OperationStore
+	prov                    provider.IProvider
+	annotatedAst            annotatedast.AnnotatedAst
+	isTargetPhysical        bool
+	txnCtrlCtrs             internaldto.TxnControlCounters
+	tableInsertionContainer tableinsertioncontainer.TableInsertionContainer
 }
 
 func NewBuilderInput(
@@ -88,6 +92,14 @@ func NewBuilderInput(
 	}
 }
 
+func (bi *builderInput) GetTableInsertionContainer() (tableinsertioncontainer.TableInsertionContainer, bool) {
+	return bi.tableInsertionContainer, bi.tableInsertionContainer != nil
+}
+
+func (bi *builderInput) SetTableInsertionContainer(ti tableinsertioncontainer.TableInsertionContainer) {
+	bi.tableInsertionContainer = ti
+}
+
 func (bi *builderInput) SetTxnCtrlCtrs(tcc internaldto.TxnControlCounters) {
 	bi.txnCtrlCtrs = tcc
 }

internal/stackql/parserutil/parser_util.go

@@ -88,6 +88,26 @@ func ExtractSelectColumnNames(selStmt *sqlparser.Select, formatter sqlparser.Nod
 	return colNames, err
 }
 
+func ExtractInsertReturningColumnNames(
+	insertStmt *sqlparser.Insert,
+	formatter sqlparser.NodeFormatter,
+) ([]ColumnHandle, error) {
+	var colNames []ColumnHandle
+	var err error
+	for _, node := range insertStmt.SelectExprs {
+		switch node := node.(type) {
+		case *sqlparser.AliasedExpr:
+			cn, cErr := inferColNameFromExpr(node.Expr, formatter, node.As.GetRawVal())
+			if cErr != nil {
+				return nil, cErr
+			}
+			colNames = append(colNames, cn)
+		case *sqlparser.StarExpr:
+		}
+	}
+	return colNames, err
+}
+
 func ExtractInsertColumnNames(insertStmt *sqlparser.Insert) ([]string, error) {
 	var colNames []string
 	var err error

internal/stackql/planbuilder/plan_builder.go

@@ -24,6 +24,7 @@ import (
 	"github.com/stackql/stackql/internal/stackql/primitivebuilder"
 	"github.com/stackql/stackql/internal/stackql/primitivegenerator"
 	"github.com/stackql/stackql/internal/stackql/primitivegraph"
+	"github.com/stackql/stackql/internal/stackql/tableinsertioncontainer"
 	"github.com/stackql/stackql/internal/stackql/tablemetadata"
 	"github.com/stackql/stackql/internal/stackql/util"
 
@@ -919,9 +920,36 @@ func (pgb *standardPlanGraphBuilder) handleInsert(pbi planbuilderinput.PlanBuild
 		if isPhysicalTable {
 			bldrInput.SetIsTargetPhysicalTable(true)
 		}
-		bldr := primitivebuilder.NewInsertOrUpdate(
-			bldrInput,
-		)
+		var bldr primitivebuilder.Builder
+		if len(node.SelectExprs) > 0 {
+			// Two cases:
+			//   1. Synchronous.  Equivalent to select.
+			//   2. Asynchronous.  Whole other story.
+			//   Synchronous only for now...
+			tableMeta, tableMetaExists := bldrInput.GetTableMetadata()
+			if !tableMetaExists {
+				return fmt.Errorf("could not obtain table metadata for node '%s'", node.Action)
+			}
+			rc, rcErr := tableinsertioncontainer.NewTableInsertionContainer(
+				tableMeta,
+				handlerCtx.GetSQLEngine(),
+				handlerCtx.GetTxnCounterMgr(),
+			)
+			if rcErr != nil {
+				return rcErr
+			}
+			bldrInput.SetTableInsertionContainer(rc)
+			bldr = primitivebuilder.NewSingleAcquireAndSelect(
+				bldrInput,
+				primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx(),
+				primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
+				nil,
+			)
+		} else {
+			bldr = primitivebuilder.NewInsertOrUpdate(
+				bldrInput,
+			)
+		}
 		err = bldr.Build()
 		if err != nil {
 			return err

internal/stackql/primitivebuilder/exec.go

@@ -6,6 +6,7 @@ import (
 	"github.com/stackql/stackql/internal/stackql/drm"
 	"github.com/stackql/stackql/internal/stackql/execution"
 	"github.com/stackql/stackql/internal/stackql/handler"
+	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/builder_input"
 	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/internaldto"
 	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/primitive_context"
 	"github.com/stackql/stackql/internal/stackql/primitive"
@@ -92,11 +93,14 @@ func (ss *Exec) Build() error {
 			return analysisErr
 		}
 		methodAnalysisOutput.GetInsertTabulation()
-		deFactoSelectBuilder := NewSingleAcquireAndSelect(
+		bldrInput := builder_input.NewBuilderInput(
 			ss.graph,
-			ss.tcc,
-			ss.handlerCtx.Clone(),
-			nil,
+			handlerCtx.Clone(),
+			tbl,
+		)
+		bldrInput.SetTxnCtrlCtrs(ss.tcc)
+		deFactoSelectBuilder := NewSingleAcquireAndSelect(
+			bldrInput,
 			nil,
 			nil,
 			nil,

internal/stackql/primitivebuilder/insert_or_update.go

@@ -8,28 +8,28 @@ import (
 	"github.com/stackql/stackql/internal/stackql/primitivegraph"
 )
 
-type InsertOrUpdate struct {
+type insertOrUpdate struct {
 	bldrInput builder_input.BuilderInput
 	root      primitivegraph.PrimitiveNode
 }
 
 func NewInsertOrUpdate(
 	bldrInput builder_input.BuilderInput,
 ) Builder {
-	return &InsertOrUpdate{
+	return &insertOrUpdate{
 		bldrInput: bldrInput,
 	}
 }
 
-func (ss *InsertOrUpdate) GetRoot() primitivegraph.PrimitiveNode {
+func (ss *insertOrUpdate) GetRoot() primitivegraph.PrimitiveNode {
 	return ss.root
 }
 
-func (ss *InsertOrUpdate) GetTail() primitivegraph.PrimitiveNode {
+func (ss *insertOrUpdate) GetTail() primitivegraph.PrimitiveNode {
 	return ss.root
 }
 
-func (ss *InsertOrUpdate) Build() error {
+func (ss *insertOrUpdate) Build() error {
 	node, nodeExists := ss.bldrInput.GetParserNode()
 	if !nodeExists {
 		return fmt.Errorf("mutation executor: node does not exist")

internal/stackql/primitivebuilder/single_acquire_and_select.go

@@ -3,8 +3,7 @@ package primitivebuilder
 import (
 	"github.com/stackql/any-sdk/pkg/streaming"
 	"github.com/stackql/stackql/internal/stackql/drm"
-	"github.com/stackql/stackql/internal/stackql/handler"
-	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/internaldto"
+	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/builder_input"
 	"github.com/stackql/stackql/internal/stackql/primitivegraph"
 	"github.com/stackql/stackql/internal/stackql/tableinsertioncontainer"
 )
@@ -13,17 +12,24 @@ type SingleAcquireAndSelect struct {
 	graph          primitivegraph.PrimitiveGraphHolder
 	acquireBuilder Builder
 	selectBuilder  Builder
+	bldrInput      builder_input.BuilderInput
+	root           primitivegraph.PrimitiveNode
 }
 
 func NewSingleAcquireAndSelect(
-	graph primitivegraph.PrimitiveGraphHolder,
-	txnControlCounters internaldto.TxnControlCounters, //nolint:revive // future proofing
-	handlerCtx handler.HandlerContext,
-	insertContainer tableinsertioncontainer.TableInsertionContainer,
+	// graph primitivegraph.PrimitiveGraphHolder,
+	// txnControlCounters internaldto.TxnControlCounters, //nolint:revive // future proofing
+	// handlerCtx handler.HandlerContext,
+	// insertContainer tableinsertioncontainer.TableInsertionContainer,
+	bldrInput builder_input.BuilderInput,
 	insertCtx drm.PreparedStatementCtx,
 	selectCtx drm.PreparedStatementCtx,
 	rowSort func(map[string]map[string]interface{}) []string,
 ) Builder {
+	graph, _ := bldrInput.GetGraphHolder()
+	// txnControlCounters, _ := bldrInput.GetTxnCtrlCtrs()
+	handlerCtx, _ := bldrInput.GetHandlerContext()
+	insertContainer, _ := bldrInput.GetTableInsertionContainer()
 	return &SingleAcquireAndSelect{
 		graph: graph,
 		acquireBuilder: NewSingleSelectAcquire(
@@ -38,6 +44,7 @@ func NewSingleAcquireAndSelect(
 			[]tableinsertioncontainer.TableInsertionContainer{insertContainer},
 			rowSort,
 			streaming.NewNopMapStream()),
+		bldrInput: bldrInput,
 	}
 }
 
@@ -60,5 +67,14 @@ func (ss *SingleAcquireAndSelect) Build() error {
 	}
 	graph := ss.graph
 	graph.NewDependency(ss.acquireBuilder.GetTail(), ss.selectBuilder.GetRoot(), 1.0)
+	rootNode := ss.acquireBuilder.GetRoot()
+	ss.root = rootNode
+	dependencyNode, dependencyNodeExists := ss.bldrInput.GetDependencyNode()
+	if dependencyNodeExists {
+		//nolint:errcheck // TODO: fix this
+		rootNode.SetInputAlias("", dependencyNode.ID())
+		ss.graph.NewDependency(dependencyNode, rootNode, 1.0)
+		// ss.root = dependencyNode // dont think this is needed
+	}
 	return nil
 }