- More await pass through.

Author: general-kroll-4-life

internal/stackql/astanalysis/annotatedast/annotated_ast.go

@@ -27,6 +27,7 @@ func NewAnnotatedAst(parent AnnotatedAst, ast sqlparser.Statement) (AnnotatedAst
 
 type AnnotatedAst interface {
 	GetAST() sqlparser.Statement
+	IsAwait(sqlparser.SQLNode) bool
 	GetIndirect(sqlparser.SQLNode) (astindirect.Indirect, bool)
 	GetIndirects() map[string]astindirect.Indirect
 	GetMaterializedView(sqlparser.SQLNode) (astindirect.Indirect, bool)
@@ -65,6 +66,31 @@ type standardAnnotatedAst struct {
 	execIndirectCache    map[*sqlparser.Exec]astindirect.Indirect
 }
 
+func hasAwaitComment(comments sqlparser.Comments) bool {
+	if comments != nil {
+		commentDirectives := sqlparser.ExtractCommentDirectives(comments)
+		if commentDirectives.IsSet("AWAIT") {
+			return true
+		}
+	}
+	return false
+}
+
+func (aa *standardAnnotatedAst) IsAwait(node sqlparser.SQLNode) bool {
+	switch n := node.(type) {
+	case *sqlparser.Exec:
+		return hasAwaitComment(n.Comments)
+	case *sqlparser.Insert:
+		return hasAwaitComment(n.Comments)
+	case *sqlparser.Update:
+		return hasAwaitComment(n.Comments)
+	case *sqlparser.Delete:
+		return hasAwaitComment(n.Comments)
+	default:
+		return false
+	}
+}
+
 func (aa *standardAnnotatedAst) GetExecIndirect(selNode *sqlparser.Exec) (astindirect.Indirect, bool) {
 	rv, ok := aa.execIndirectCache[selNode]
 	return rv, ok

internal/stackql/astanalysis/routeanalysis/select_route_analysis.go

@@ -153,7 +153,7 @@ func (sp *standardSelectRoutePass) RoutePass() error {
 
 	// TODO: Do the proper SOLID treatment on router, etc.
 	// Might need to split into multiple passes.
-	v := router.NewTableRouteAstVisitor(pbi.GetHandlerCtx(), paramRouter)
+	v := router.NewTableRouteAstVisitor(pbi.GetHandlerCtx(), paramRouter, false)
 
 	err := v.Visit(node)
 

internal/stackql/dependencyplanner/dependencyplanner.go

@@ -497,7 +497,7 @@ func (dp *standardDependencyPlanner) orchestrate(
 	return idx, err
 }
 
-//nolint:gocognit // live with it
+//nolint:gocognit,funlen // live with it
 func (dp *standardDependencyPlanner) processAcquire(
 	sqlNode sqlparser.SQLNode,
 	annotationCtx taxonomy.AnnotationCtx,
@@ -571,7 +571,7 @@ func (dp *standardDependencyPlanner) processAcquire(
 		svc,
 		resource,
 		m,
-		false, // TODO: wire up isAwait
+		annotationCtx.IsAwait(),
 		[]util.AnnotatedTabulation{anTab},
 		dp.primitiveComposer.GetSQLEngine(),
 		prov.GetName(),

internal/stackql/parserutil/parser_util.go

@@ -17,6 +17,20 @@ const (
 	FloatBitSize int = 64
 )
 
+func hasAwaitComment(comments sqlparser.Comments) bool {
+	if comments != nil {
+		commentDirectives := sqlparser.ExtractCommentDirectives(comments)
+		if commentDirectives.IsSet("AWAIT") {
+			return true
+		}
+	}
+	return false
+}
+
+func IsAwait(comments sqlparser.Comments) bool {
+	return hasAwaitComment(comments)
+}
+
 //nolint:gocritic,staticcheck // TODO: clean up this and other tech debt
 func ExtractLHSTable(sqlFunc *sqlparser.FuncExpr) (*sqlparser.ColName, bool) {
 	if sqlFunc == nil {

internal/stackql/primitivegenerator/statement_analyzer.go

@@ -518,12 +518,22 @@ func extractVarDefFromExec(node *sqlparser.Exec, argName string) (*sqlparser.Exe
 
 func (pb *standardPrimitiveGenerator) parseComments(comments sqlparser.Comments) {
 	if comments != nil {
-		pb.PrimitiveComposer.SetCommentDirectives(sqlparser.ExtractCommentDirectives(comments))
-		// This pattern supports preset.
-		if pb.PrimitiveComposer.GetCommentDirectives().IsSet("AWAIT") {
-			pb.PrimitiveComposer.SetAwait(true)
+		commentDirectives := sqlparser.ExtractCommentDirectives(comments)
+		pb.PrimitiveComposer.SetCommentDirectives(commentDirectives)
+	}
+	if pb.hasAwaitComment(comments) {
+		pb.PrimitiveComposer.SetAwait(true)
+	}
+}
+
+func (pb *standardPrimitiveGenerator) hasAwaitComment(comments sqlparser.Comments) bool {
+	if comments != nil {
+		commentDirectives := sqlparser.ExtractCommentDirectives(comments)
+		if commentDirectives.IsSet("AWAIT") {
+			return true
 		}
 	}
+	return false
 }
 
 func (pb *standardPrimitiveGenerator) persistHerarchyToBuilder(
@@ -542,7 +552,8 @@ func (pb *standardPrimitiveGenerator) AnalyzeUnaryExec(
 	selectNode *sqlparser.Select,
 	cols []parserutil.ColumnHandle,
 ) (tablemetadata.ExtendedTableMetadata, error) {
-	err := pb.inferHeirarchyAndPersist(handlerCtx, node, nil, false)
+	isAwait := pb.hasAwaitComment(node.Comments)
+	err := pb.inferHeirarchyAndPersist(handlerCtx, node, nil, false, isAwait)
 	if err != nil {
 		return nil, err
 	}
@@ -1025,7 +1036,8 @@ func (pb *standardPrimitiveGenerator) AnalyzeInsert(pbi planbuilderinput.PlanBui
 	if !ok {
 		return fmt.Errorf("could not cast node of type '%T' to required Insert", pbi.GetStatement())
 	}
-	err := pb.inferHeirarchyAndPersist(handlerCtx, node, pbi.GetPlaceholderParams(), false)
+	isAwait := pb.hasAwaitComment(node.Comments)
+	err := pb.inferHeirarchyAndPersist(handlerCtx, node, pbi.GetPlaceholderParams(), false, isAwait)
 	if err != nil {
 		return err
 	}
@@ -1181,7 +1193,8 @@ func (pb *standardPrimitiveGenerator) AnalyzeUpdate(pbi planbuilderinput.PlanBui
 	if !ok {
 		return fmt.Errorf("could not cast node of type '%T' to required Update", pbi.GetStatement())
 	}
-	err := pb.inferHeirarchyAndPersist(handlerCtx, node, pbi.GetPlaceholderParams(), false)
+	isAwait := pb.hasAwaitComment(node.Comments)
+	err := pb.inferHeirarchyAndPersist(handlerCtx, node, pbi.GetPlaceholderParams(), false, isAwait)
 	if err != nil {
 		return err
 	}
@@ -1292,8 +1305,10 @@ func (pb *standardPrimitiveGenerator) inferHeirarchyAndPersist(
 	handlerCtx handler.HandlerContext,
 	node sqlparser.SQLNode,
 	parameters parserutil.ColumnKeyedDatastore,
-	viewPermissive bool) error {
-	heirarchy, err := taxonomy.GetHeirarchyFromStatement(handlerCtx, node, parameters, viewPermissive)
+	viewPermissive bool,
+	isAwait bool,
+) error {
+	heirarchy, err := taxonomy.GetHeirarchyFromStatement(handlerCtx, node, parameters, viewPermissive, isAwait)
 	if err != nil {
 		return err
 	}
@@ -1312,7 +1327,8 @@ func (pb *standardPrimitiveGenerator) analyzeDelete(
 	}
 	pb.parseComments(node.Comments)
 	paramMap, paramMapExists := pbi.GetAnnotatedAST().GetWhereParamMapsEntry(node.Where)
-	err := pb.inferHeirarchyAndPersist(handlerCtx, node, paramMap, false)
+	isAwait := pb.hasAwaitComment(node.Comments)
+	err := pb.inferHeirarchyAndPersist(handlerCtx, node, paramMap, false, isAwait)
 	if err != nil {
 		return err
 	}
@@ -1404,7 +1420,8 @@ func (pb *standardPrimitiveGenerator) analyzeDescribe(pbi planbuilderinput.PlanB
 		return fmt.Errorf("could not cast node of type '%T' to required Describe", pbi.GetStatement())
 	}
 	var err error
-	err = pb.inferHeirarchyAndPersist(handlerCtx, node, nil, true)
+	// DESCRIBE has no await or comments at this time
+	err = pb.inferHeirarchyAndPersist(handlerCtx, node, nil, true, false)
 	if err != nil {
 		return err
 	}
@@ -1515,12 +1532,14 @@ func (pb *standardPrimitiveGenerator) analyzeShow(
 	case "AUTH":
 		// TODO
 	case "INSERT":
-		err = pb.inferHeirarchyAndPersist(handlerCtx, node, nil, false)
+		// SHOW has no comments or concept of await
+		err = pb.inferHeirarchyAndPersist(handlerCtx, node, nil, false, false)
 		if err != nil {
 			return err
 		}
 	case "METHODS":
-		err = pb.inferHeirarchyAndPersist(handlerCtx, node, nil, false)
+		// SHOW has no comments or concept of await
+		err = pb.inferHeirarchyAndPersist(handlerCtx, node, nil, false, false)
 		if err != nil {
 			return err
 		}

internal/stackql/router/parameter_router.go

@@ -49,7 +49,7 @@ type ParameterRouter interface {
 	//   - Hierarchy.
 	//   - Columnar objects definitely assigned as HTTP method parameters.
 	//   - Error if applicable.
-	Route(tb sqlparser.TableExpr, handler handler.HandlerContext) (taxonomy.AnnotationCtx, error)
+	Route(tb sqlparser.TableExpr, handler handler.HandlerContext, isAwait bool) (taxonomy.AnnotationCtx, error)
 
 	// Detects:
 	//   - Dependency cycle.
@@ -511,14 +511,16 @@ func (pr *standardParameterRouter) isInvalidated(key string) bool {
 func (pr *standardParameterRouter) Route(
 	tb sqlparser.TableExpr,
 	handlerCtx handler.HandlerContext,
+	isAwait bool,
 ) (taxonomy.AnnotationCtx, error) {
-	return pr.route(tb, handlerCtx)
+	return pr.route(tb, handlerCtx, isAwait)
 }
 
 //nolint:funlen,gocognit,govet,gocyclo,cyclop // inherently complex functionality
 func (pr *standardParameterRouter) route(
 	tb sqlparser.TableExpr,
 	handlerCtx handler.HandlerContext,
+	isAwait bool,
 ) (taxonomy.AnnotationCtx, error) {
 	// TODO: Get rid of the dead set mess that is where paramters in preference.
 	for k, v := range pr.whereParamMap.GetMap() {
@@ -589,9 +591,9 @@ func (pr *standardParameterRouter) route(
 	priorParameters := runParamters.Clone()
 	// notOnStringParams := notOnParams.GetStringified()
 	// TODO: add parent params into the mix here.
-	hr, err := taxonomy.GetHeirarchyFromStatement(handlerCtx, tb, notOnParams, false)
+	hr, err := taxonomy.GetHeirarchyFromStatement(handlerCtx, tb, notOnParams, false, isAwait)
 	if err != nil {
-		hr, err = taxonomy.GetHeirarchyFromStatement(handlerCtx, tb, runParamters, false)
+		hr, err = taxonomy.GetHeirarchyFromStatement(handlerCtx, tb, runParamters, false, isAwait)
 	} else {
 		// If the where parameters are sufficient, then need to switch
 		// the Table - Paramater coupling object

internal/stackql/router/table_routing.go

@@ -30,14 +30,16 @@ type standardTableRouteAstVisitor struct {
 	router      ParameterRouter
 	tables      taxonomy.TblMap
 	annotations taxonomy.AnnotationCtxMap
+	isAwait     bool
 }
 
-func NewTableRouteAstVisitor(handlerCtx handler.HandlerContext, router ParameterRouter) TableRouteAstVisitor {
+func NewTableRouteAstVisitor(handlerCtx handler.HandlerContext, router ParameterRouter, isAwait bool) TableRouteAstVisitor {
 	return &standardTableRouteAstVisitor{
 		handlerCtx:  handlerCtx,
 		router:      router,
 		tables:      make(taxonomy.TblMap),
 		annotations: make(taxonomy.AnnotationCtxMap),
+		isAwait:     isAwait,
 	}
 }
 
@@ -55,13 +57,13 @@ func (v *standardTableRouteAstVisitor) analyzeAliasedTable(
 ) (taxonomy.AnnotationCtx, error) {
 	switch ex := tb.Expr.(type) {
 	case *sqlparser.Subquery:
-		return v.router.Route(tb, v.handlerCtx)
+		return v.router.Route(tb, v.handlerCtx, v.isAwait)
 	case sqlparser.TableName:
 		_, err := taxonomy.GetHeirarchyIDsFromParserNode(v.handlerCtx, ex)
 		if err != nil {
 			return nil, err
 		}
-		return v.router.Route(tb, v.handlerCtx)
+		return v.router.Route(tb, v.handlerCtx, v.isAwait)
 	default:
 		return nil, fmt.Errorf("table of type '%T' not currently supported", ex)
 	}
@@ -75,7 +77,7 @@ func (v *standardTableRouteAstVisitor) analyzeExecTableName(
 	tb *sqlparser.ExecSubquery,
 ) (taxonomy.AnnotationCtx, error) {
 	// cannot be recursive
-	return v.router.Route(tb, v.handlerCtx)
+	return v.router.Route(tb, v.handlerCtx, v.isAwait)
 }
 
 func (v *standardTableRouteAstVisitor) GetTableMap() taxonomy.TblMap {

internal/stackql/tablemetadata/hierarchy_objects.go

@@ -41,12 +41,14 @@ type HeirarchyObjects interface {
 	IsPGInternalObject() bool
 	SetIndirect(internaldto.RelationDTO)
 	GetIndirect() (internaldto.RelationDTO, bool)
+	IsAwait() bool
 }
 
-func NewHeirarchyObjects(hIDs internaldto.HeirarchyIdentifiers) HeirarchyObjects {
+func NewHeirarchyObjects(hIDs internaldto.HeirarchyIdentifiers, isAwait bool) HeirarchyObjects {
 	return &standardHeirarchyObjects{
 		heirarchyIDs: hIDs,
 		hr:           internaldto.NewHeirarchy(hIDs),
+		isAwait:      isAwait,
 	}
 }
 
@@ -56,6 +58,11 @@ type standardHeirarchyObjects struct {
 	prov          provider.IProvider
 	sqlDataSource sql_datasource.SQLDataSource
 	indirect      internaldto.RelationDTO
+	isAwait       bool
+}
+
+func (ho *standardHeirarchyObjects) IsAwait() bool {
+	return ho.isAwait
 }
 
 func (ho *standardHeirarchyObjects) GetIndirect() (internaldto.RelationDTO, bool) {

internal/stackql/taxonomy/hierarchy.go

@@ -184,6 +184,7 @@ func GetHeirarchyFromStatement(
 	node sqlparser.SQLNode,
 	parameters parserutil.ColumnKeyedDatastore,
 	viewPermissive bool,
+	isAwait bool,
 ) (tablemetadata.HeirarchyObjects, error) {
 	var hIDs internaldto.HeirarchyIdentifiers
 	getFirstAvailableMethod := false
@@ -205,10 +206,10 @@ func GetHeirarchyFromStatement(
 	case *sqlparser.AliasedTableExpr:
 		switch n.Expr.(type) { //nolint:gocritic // this is expressive enough
 		case *sqlparser.Subquery:
-			retVal := tablemetadata.NewHeirarchyObjects(hIDs)
+			retVal := tablemetadata.NewHeirarchyObjects(hIDs, isAwait)
 			return retVal, nil
 		}
-		return GetHeirarchyFromStatement(handlerCtx, n.Expr, parameters, viewPermissive)
+		return GetHeirarchyFromStatement(handlerCtx, n.Expr, parameters, viewPermissive, isAwait)
 	case *sqlparser.Show:
 		switch strings.ToUpper(n.Type) {
 		case "INSERT":
@@ -228,7 +229,7 @@ func GetHeirarchyFromStatement(
 	default:
 		return nil, fmt.Errorf("cannot resolve taxonomy")
 	}
-	retVal := tablemetadata.NewHeirarchyObjects(hIDs)
+	retVal := tablemetadata.NewHeirarchyObjects(hIDs, isAwait)
 	sqlDataSource, isSQLDataSource := handlerCtx.GetSQLDataSource(hIDs.GetProviderStr())
 	if isSQLDataSource {
 		retVal.SetSQLDataSource(sqlDataSource)