- Working localy for synchronous.

Author: general-kroll-4-life

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,23 @@ 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

@@ -10,7 +10,6 @@ import (
 	"github.com/stackql/any-sdk/anysdk"
 
 	"github.com/stackql/any-sdk/pkg/logging"
-	"github.com/stackql/any-sdk/pkg/streaming"
 	"github.com/stackql/stackql/internal/stackql/acid/txn_context"
 	"github.com/stackql/stackql/internal/stackql/astanalysis/routeanalysis"
 	"github.com/stackql/stackql/internal/stackql/handler"
@@ -926,6 +925,7 @@ func (pgb *standardPlanGraphBuilder) handleInsert(pbi planbuilderinput.PlanBuild
 			// 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)
@@ -938,13 +938,12 @@ func (pgb *standardPlanGraphBuilder) handleInsert(pbi planbuilderinput.PlanBuild
 			if rcErr != nil {
 				return rcErr
 			}
-			bldr = primitivebuilder.NewSingleSelectAcquire(
-				pgb.planGraphHolder,
-				handlerCtx,
-				rc,
+			bldrInput.SetTableInsertionContainer(rc)
+			bldr = primitivebuilder.NewSingleAcquireAndSelect(
+				bldrInput,
 				primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx(),
+				primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
 				nil,
-				streaming.NewNopMapStream(),
 			)
 		} else {
 			bldr = primitivebuilder.NewInsertOrUpdate(

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/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
 }

internal/stackql/primitivegenerator/select.go

@@ -10,6 +10,7 @@ import (
 	"github.com/stackql/stackql/internal/stackql/astindirect"
 	"github.com/stackql/stackql/internal/stackql/astvisit"
 	"github.com/stackql/stackql/internal/stackql/dependencyplanner"
+	"github.com/stackql/stackql/internal/stackql/internal_data_transfer/builder_input"
 	"github.com/stackql/stackql/internal/stackql/parserutil"
 	"github.com/stackql/stackql/internal/stackql/planbuilderinput"
 	"github.com/stackql/stackql/internal/stackql/primitivebuilder"
@@ -251,12 +252,16 @@ func (pb *standardPrimitiveGenerator) analyzeSelect(pbi planbuilderinput.PlanBui
 				return indirectErr
 			}
 			annotatedAST.SetSelectIndirect(node, selIndirect)
+			bldrInput := builder_input.NewBuilderInput(
+				pChild.GetPrimitiveComposer().GetGraphHolder(),
+				handlerCtx,
+				tbl,
+			)
+			bldrInput.SetTxnCtrlCtrs(pChild.GetPrimitiveComposer().GetTxnCtrlCtrs())
+			bldrInput.SetTableInsertionContainer(insertionContainer)
 			pChild.GetPrimitiveComposer().SetBuilder(
 				primitivebuilder.NewSingleAcquireAndSelect(
-					pChild.GetPrimitiveComposer().GetGraphHolder(),
-					pChild.GetPrimitiveComposer().GetTxnCtrlCtrs(),
-					handlerCtx,
-					insertionContainer,
+					bldrInput,
 					pChild.GetPrimitiveComposer().GetInsertPreparedStatementCtx(),
 					pChild.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
 					nil))

internal/stackql/primitivegenerator/statement_analyzer.go

@@ -704,14 +704,20 @@ func (pb *standardPrimitiveGenerator) analyzeExec(pbi planbuilderinput.PlanBuild
 		return indirectErr
 	}
 	annotatedAST.SetExecIndirect(node, selIndirect)
+	bldrInput := builder_input.NewBuilderInput(
+		pb.PrimitiveComposer.GetGraphHolder(),
+		handlerCtx,
+		tbl,
+	)
+	bldrInput.SetTxnCtrlCtrs(pb.PrimitiveComposer.GetTxnCtrlCtrs())
+	bldrInput.SetTableInsertionContainer(insertionContainer)
 	pb.PrimitiveComposer.SetBuilder(
 		primitivebuilder.NewSingleAcquireAndSelect(
-			pb.PrimitiveComposer.GetGraphHolder(),
-			pb.PrimitiveComposer.GetTxnCtrlCtrs(),
-			handlerCtx,
-			insertionContainer,
+			bldrInput,
 			pb.PrimitiveComposer.GetInsertPreparedStatementCtx(),
-			pb.PrimitiveComposer.GetSelectPreparedStatementCtx(), nil))
+			pb.PrimitiveComposer.GetSelectPreparedStatementCtx(),
+			nil,
+		))
 	return nil
 }
 
@@ -1127,13 +1133,20 @@ func (pb *standardPrimitiveGenerator) AnalyzeInsert(pbi planbuilderinput.PlanBui
 	if err != nil {
 		return err
 	}
+	columnHandles := []parserutil.ColumnHandle{}
+	if len(node.SelectExprs) > 0 {
+		columnHandles, err = parserutil.ExtractInsertReturningColumnNames(node, handlerCtx.GetASTFormatter())
+		if err != nil {
+			return err
+		}
+	}
 	err = pb.analyzeUnaryAction(
 		pbi,
 		handlerCtx,
 		node,
 		nil,
 		tbl,
-		[]parserutil.ColumnHandle{},
+		columnHandles,
 		methodAnalysisOutput,
 	)
 	if err != nil {