- PAtched unit testing garbage.

Author: general-kroll-4-life

internal/stackql/execution/mono_valent_execution.go

@@ -1478,6 +1478,7 @@ func GetMonitorExecutor(
 	precursor primitive.IPrimitive,
 	initialCtx primitive.IPrimitiveCtx,
 	comments sqlparser.CommentDirectives,
+	isReturning bool,
 ) (primitive.IPrimitive, error) {
 	m := op
 	// tableName, err := mv.tableMeta.GetTableName()
@@ -1528,7 +1529,7 @@ func GetMonitorExecutor(
 		//nolint:nestif // acceptable for now
 		if endTimeOk && endTime != "" {
 			targetLink, targetLinkOK := body["targetLink"]
-			if targetLinkOK {
+			if targetLinkOK && isReturning {
 				authCtx, authErr := pc.GetAuthContext(prStr)
 				if authErr != nil {
 					return internaldto.NewExecutorOutput(nil, nil, nil, nil, authErr)
@@ -1568,6 +1569,7 @@ func GetMonitorExecutor(
 				}
 				return prepareResultSet(&asyncPrim, pc, target, operationDescriptor)
 			}
+			return prepareResultSet(&asyncPrim, pc, body, operationDescriptor)
 		}
 		url, ok := body["selfLink"]
 		if !ok {

internal/stackql/internal_data_transfer/builder_input/builder_input.go

@@ -31,6 +31,8 @@ type BuilderInput interface {
 	GetOperationStore() (anysdk.OperationStore, bool)
 	SetOperationStore(op anysdk.OperationStore)
 	IsAwait() bool
+	IsReturning() bool
+	SetIsReturning(bool)
 	GetVerb() string
 	GetInputAlias() string
 	IsUndo() bool
@@ -64,6 +66,7 @@ type builderInput struct {
 	dependencyNode          primitivegraph.PrimitiveNode
 	commentDirectives       sqlparser.CommentDirectives
 	isAwait                 bool
+	isReturning             bool
 	verb                    string
 	inputAlias              string
 	isUndo                  bool
@@ -92,6 +95,14 @@ func NewBuilderInput(
 	}
 }
 
+func (bi *builderInput) IsReturning() bool {
+	return bi.isReturning
+}
+
+func (bi *builderInput) SetIsReturning(isReturning bool) {
+	bi.isReturning = isReturning
+}
+
 func (bi *builderInput) GetTableInsertionContainer() (tableinsertioncontainer.TableInsertionContainer, bool) {
 	return bi.tableInsertionContainer, bi.tableInsertionContainer != nil
 }

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"
@@ -954,13 +953,11 @@ func (pgb *standardPlanGraphBuilder) handleInsert(pbi planbuilderinput.PlanBuild
 			if !isAwait {
 				bldr = returningBldr
 			} else {
-				rhsBldr := primitivebuilder.NewSingleSelect(
-					pgb.planGraphHolder,
-					handlerCtx,
+				rhsBldr := primitivebuilder.NewSingleAcquireAndSelect(
+					bldrInput,
+					primitiveGenerator.GetPrimitiveComposer().GetInsertPreparedStatementCtx(),
 					primitiveGenerator.GetPrimitiveComposer().GetSelectPreparedStatementCtx(),
-					[]tableinsertioncontainer.TableInsertionContainer{rc},
 					nil,
-					streaming.NewNopMapStream(),
 				)
 				bldr = primitivebuilder.NewDependencySubDAGBuilder(
 					pgb.planGraphHolder,

internal/stackql/primitivebuilder/async_compose.go

@@ -15,8 +15,9 @@ func composeAsyncMonitor(
 	prov provider.IProvider,
 	method anysdk.OperationStore,
 	commentDirectives sqlparser.CommentDirectives,
+	isReturning bool,
 ) (primitive.IPrimitive, error) {
-	asm, err := NewAsyncMonitor(handlerCtx, prov, method)
+	asm, err := NewAsyncMonitor(handlerCtx, prov, method, isReturning)
 	if err != nil {
 		return nil, err
 	}
@@ -31,7 +32,7 @@ func composeAsyncMonitor(
 		handlerCtx.GetOutfile(),
 		handlerCtx.GetOutErrFile(),
 	)
-	primitive, err := asm.GetMonitorPrimitive(prov, method, precursor, pl, commentDirectives)
+	primitive, err := asm.GetMonitorPrimitive(prov, method, precursor, pl, commentDirectives, isReturning)
 	if err != nil {
 		return nil, err
 	}

internal/stackql/primitivebuilder/asyncmonitor.go

@@ -22,6 +22,7 @@ type IAsyncMonitor interface {
 		precursor primitive.IPrimitive,
 		initialCtx primitive.IPrimitiveCtx,
 		comments sqlparser.CommentDirectives,
+		isReturning bool,
 	) (primitive.IPrimitive, error)
 }
 
@@ -120,11 +121,12 @@ func NewAsyncMonitor(
 	handlerCtx handler.HandlerContext,
 	prov provider.IProvider,
 	op anysdk.OperationStore,
+	isReturning bool,
 ) (IAsyncMonitor, error) {
 	//nolint:gocritic //TODO: refactor
 	switch prov.GetProviderString() {
 	case "google":
-		return newGoogleAsyncMonitor(handlerCtx, prov, op, prov.GetVersion())
+		return newGoogleAsyncMonitor(handlerCtx, prov, op, prov.GetVersion(), isReturning)
 	}
 	return nil, fmt.Errorf(
 		"async operation monitor for provider = '%s', api version = '%s' currently not supported",
@@ -136,6 +138,7 @@ func newGoogleAsyncMonitor(
 	prov provider.IProvider,
 	op anysdk.OperationStore,
 	version string, //nolint:unparam // TODO: refactor
+	isReturning bool,
 ) (IAsyncMonitor, error) {
 	//nolint:gocritic //TODO: refactor
 	switch version {
@@ -160,11 +163,12 @@ func (gm *DefaultGoogleAsyncMonitor) GetMonitorPrimitive(
 	precursor primitive.IPrimitive,
 	initialCtx primitive.IPrimitiveCtx,
 	comments sqlparser.CommentDirectives,
+	isReturning bool,
 ) (primitive.IPrimitive, error) {
 	//nolint:gocritic,staticcheck //TODO: refactor
 	switch strings.ToLower(prov.GetVersion()) {
 	default:
-		return gm.getV1Monitor(prov, op, precursor, initialCtx, comments)
+		return gm.getV1Monitor(prov, op, precursor, initialCtx, comments, isReturning)
 	}
 }
 
@@ -174,6 +178,7 @@ func (gm *DefaultGoogleAsyncMonitor) getV1Monitor(
 	precursor primitive.IPrimitive,
 	initialCtx primitive.IPrimitiveCtx,
 	comments sqlparser.CommentDirectives,
+	isReturning bool,
 ) (primitive.IPrimitive, error) {
 	provider, providerErr := prov.GetProvider()
 	if providerErr != nil {
@@ -186,6 +191,7 @@ func (gm *DefaultGoogleAsyncMonitor) getV1Monitor(
 		precursor,
 		initialCtx,
 		comments,
+		isReturning,
 	)
 	if exPrepErr != nil {
 		return nil, exPrepErr

internal/stackql/primitivebuilder/delete.go

@@ -93,7 +93,7 @@ func (ss *Delete) Build() error {
 		primitive_context.NewPrimitiveContext(),
 	)
 	if ss.isAwait {
-		deletePrimitive, err = composeAsyncMonitor(handlerCtx, deletePrimitive, prov, method, nil)
+		deletePrimitive, err = composeAsyncMonitor(handlerCtx, deletePrimitive, prov, method, nil, false)
 	}
 	if err != nil {
 		return err

internal/stackql/primitivebuilder/exec.go

@@ -174,7 +174,7 @@ func (ss *Exec) Build() error {
 		ss.graph.CreatePrimitiveNode(execPrimitive)
 		return nil
 	}
-	pr, err := composeAsyncMonitor(handlerCtx, execPrimitive, prov, m, nil)
+	pr, err := composeAsyncMonitor(handlerCtx, execPrimitive, prov, m, nil, false) // returning hardcoded to false for now
 	if err != nil {
 		return err
 	}

internal/stackql/primitivebuilder/generic_http_reversal.go

@@ -25,6 +25,7 @@ type genericHTTPReversal struct {
 	op                anysdk.OperationStore
 	commentDirectives sqlparser.CommentDirectives
 	isAwait           bool
+	isReturning       bool
 	verb              string // may be "insert" or "update"
 	inputAlias        string
 	isUndo            bool
@@ -206,7 +207,8 @@ func (gh *genericHTTPReversal) Build() error {
 				if err != nil {
 					return internaldto.NewErroneousExecutorOutput(err)
 				}
-				execPrim, execErr := composeAsyncMonitor(handlerCtx, dependentInsertPrimitive, prov, m, commentDirectives)
+				execPrim, execErr := composeAsyncMonitor(
+					handlerCtx, dependentInsertPrimitive, prov, m, commentDirectives, gh.isReturning)
 				if execErr != nil {
 					return internaldto.NewErroneousExecutorOutput(execErr)
 				}

internal/stackql/primitivebuilder/generic_http_stream_input.go

@@ -32,6 +32,7 @@ type genericHTTPStreamInput struct {
 	dependencyNode    primitivegraph.PrimitiveNode
 	parserNode        sqlparser.SQLNode
 	isAwait           bool
+	isReturning       bool
 	verb              string // may be "insert" or "update"
 	inputAlias        string
 	isUndo            bool
@@ -70,6 +71,7 @@ func newGenericHTTPStreamInput(
 		commentDirectives: commentDirectives,
 		dependencyNode:    dependencyNode,
 		isAwait:           builderInput.IsAwait(),
+		isReturning:       builderInput.IsReturning(),
 		verb:              builderInput.GetVerb(),
 		inputAlias:        builderInput.GetInputAlias(),
 		isUndo:            builderInput.IsUndo(),
@@ -357,7 +359,8 @@ func (gh *genericHTTPStreamInput) Build() error {
 				if err != nil {
 					return internaldto.NewErroneousExecutorOutput(err)
 				}
-				execPrim, execErr := composeAsyncMonitor(handlerCtx, dependentInsertPrimitive, prov, m, commentDirectives)
+				execPrim, execErr := composeAsyncMonitor(
+					handlerCtx, dependentInsertPrimitive, prov, m, commentDirectives, gh.isReturning)
 				if execErr != nil {
 					return internaldto.NewErroneousExecutorOutput(execErr)
 				}

internal/stackql/primitivebuilder/insert_or_update.go

@@ -38,6 +38,9 @@ func (ss *insertOrUpdate) Build() error {
 	switch node := node.(type) {
 	case *sqlparser.Insert:
 		mutableInput.SetVerb("insert")
+		if len(node.SelectExprs) > 0 {
+			mutableInput.SetIsReturning(true)
+		}
 	case *sqlparser.Update:
 		mutableInput.SetVerb("update")
 	default: