From 45ab7520399b4ee875ef65b00a2f165fe3ede9e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Dec 2025 09:27:27 +0000 Subject: [PATCH 1/5] feat: add CTE support and robot tests for window functions - Implement CTE (Common Table Expression) support in execution layer - Extract CTEs from WITH clause during AST expansion - Convert CTE references to inline subqueries for processing - Store CTE registry in visitor for reference resolution - Add robot tests for window functions - ROW_NUMBER() OVER - SUM() OVER (with and without ORDER BY) - RANK() OVER - Add robot test for simple CTE usage - Test queries using WITH clause against accelerator types data --- .../astanalysis/earlyanalysis/ast_expand.go | 36 +++++++++++ .../cte/simple-cte.txt | 11 ++++ .../window-functions/rank-by-cards.txt | 11 ++++ .../window-functions/row-number-by-name.txt | 11 ++++ .../window-functions/running-sum-by-name.txt | 11 ++++ .../window-functions/sum-over-all.txt | 11 ++++ .../stackql_test_tooling/stackql_context.py | 28 +++++++++ .../stackql_mocked_from_cmd_line.robot | 60 +++++++++++++++++++ 8 files changed, 179 insertions(+) create mode 100644 test/assets/expected/simple-select/compute-accelerator-type/cte/simple-cte.txt create mode 100644 test/assets/expected/simple-select/compute-accelerator-type/window-functions/rank-by-cards.txt create mode 100644 test/assets/expected/simple-select/compute-accelerator-type/window-functions/row-number-by-name.txt create mode 100644 test/assets/expected/simple-select/compute-accelerator-type/window-functions/running-sum-by-name.txt create mode 100644 test/assets/expected/simple-select/compute-accelerator-type/window-functions/sum-over-all.txt diff --git a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go index 4e03cccce..2be8e080c 100644 --- a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go +++ b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go @@ -49,6 +49,7 @@ type indirectExpandAstVisitor struct { selectCount int mutateCount int createBuilder []primitivebuilder.Builder + cteRegistry map[string]*sqlparser.Subquery // CTE name -> subquery definition } func newIndirectExpandAstVisitor( @@ -75,6 +76,7 @@ func newIndirectExpandAstVisitor( tcc: tcc, whereParams: whereParams, indirectionDepth: indirectionDepth, + cteRegistry: make(map[string]*sqlparser.Subquery), } return rv, nil } @@ -214,6 +216,14 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error { addIf(node.StraightJoinHint, sqlparser.StraightJoinHint) addIf(node.SQLCalcFoundRows, sqlparser.SQLCalcFoundRowsStr) + // Extract CTEs from WITH clause and store in registry + if node.With != nil { + for _, cte := range node.With.CTEs { + cteName := cte.Name.GetRawVal() + v.cteRegistry[cteName] = cte.Subquery + } + } + if node.Comments != nil { node.Comments.Accept(v) //nolint:errcheck // future proof } @@ -785,6 +795,32 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error { return nil //nolint:nilerr //TODO: investigate } return nil + case sqlparser.TableName: + // Check if this table name is a CTE reference + tableName := n.GetRawVal() + if cteSubquery, isCTE := v.cteRegistry[tableName]; isCTE { + // Process CTE reference as a subquery + // Use the CTE name as the effective alias if no explicit alias is set + effectiveAlias := node.As.GetRawVal() + if effectiveAlias == "" { + effectiveAlias = tableName + } + // Create a synthetic AliasedTableExpr with the CTE subquery + syntheticExpr := &sqlparser.AliasedTableExpr{ + Expr: cteSubquery, + As: sqlparser.NewTableIdent(effectiveAlias), + } + sq := internaldto.NewSubqueryDTO(syntheticExpr, cteSubquery) + indirect, err := astindirect.NewSubqueryIndirect(sq) + if err != nil { + return nil //nolint:nilerr //TODO: investigate + } + err = v.processIndirect(syntheticExpr, indirect) + if err != nil { + return nil //nolint:nilerr //TODO: investigate + } + return nil + } } err := node.Expr.Accept(v) if err != nil { diff --git a/test/assets/expected/simple-select/compute-accelerator-type/cte/simple-cte.txt b/test/assets/expected/simple-select/compute-accelerator-type/cte/simple-cte.txt new file mode 100644 index 000000000..ac262bfda --- /dev/null +++ b/test/assets/expected/simple-select/compute-accelerator-type/cte/simple-cte.txt @@ -0,0 +1,11 @@ +|---------------------|-------------------------| +| name | maximumCardsPerInstance | +|---------------------|-------------------------| +| nvidia-tesla-p4 | 4 | +|---------------------|-------------------------| +| nvidia-tesla-p4-vws | 4 | +|---------------------|-------------------------| +| nvidia-tesla-t4 | 4 | +|---------------------|-------------------------| +| nvidia-tesla-t4-vws | 4 | +|---------------------|-------------------------| diff --git a/test/assets/expected/simple-select/compute-accelerator-type/window-functions/rank-by-cards.txt b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/rank-by-cards.txt new file mode 100644 index 000000000..2fa81ac89 --- /dev/null +++ b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/rank-by-cards.txt @@ -0,0 +1,11 @@ +|---------------------|------| +| name | rank | +|---------------------|------| +| nvidia-tesla-p4 | 1 | +|---------------------|------| +| nvidia-tesla-p4-vws | 1 | +|---------------------|------| +| nvidia-tesla-t4 | 1 | +|---------------------|------| +| nvidia-tesla-t4-vws | 1 | +|---------------------|------| diff --git a/test/assets/expected/simple-select/compute-accelerator-type/window-functions/row-number-by-name.txt b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/row-number-by-name.txt new file mode 100644 index 000000000..368933bb2 --- /dev/null +++ b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/row-number-by-name.txt @@ -0,0 +1,11 @@ +|---------------------|------------| +| name | row_number | +|---------------------|------------| +| nvidia-tesla-p4 | 1 | +|---------------------|------------| +| nvidia-tesla-p4-vws | 2 | +|---------------------|------------| +| nvidia-tesla-t4 | 3 | +|---------------------|------------| +| nvidia-tesla-t4-vws | 4 | +|---------------------|------------| diff --git a/test/assets/expected/simple-select/compute-accelerator-type/window-functions/running-sum-by-name.txt b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/running-sum-by-name.txt new file mode 100644 index 000000000..9df074a51 --- /dev/null +++ b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/running-sum-by-name.txt @@ -0,0 +1,11 @@ +|---------------------|-------------| +| name | running_sum | +|---------------------|-------------| +| nvidia-tesla-p4 | 4 | +|---------------------|-------------| +| nvidia-tesla-p4-vws | 8 | +|---------------------|-------------| +| nvidia-tesla-t4 | 12 | +|---------------------|-------------| +| nvidia-tesla-t4-vws | 16 | +|---------------------|-------------| diff --git a/test/assets/expected/simple-select/compute-accelerator-type/window-functions/sum-over-all.txt b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/sum-over-all.txt new file mode 100644 index 000000000..2dfacf37b --- /dev/null +++ b/test/assets/expected/simple-select/compute-accelerator-type/window-functions/sum-over-all.txt @@ -0,0 +1,11 @@ +|---------------------|-----------| +| name | total_sum | +|---------------------|-----------| +| nvidia-tesla-p4 | 16 | +|---------------------|-----------| +| nvidia-tesla-p4-vws | 16 | +|---------------------|-----------| +| nvidia-tesla-t4 | 16 | +|---------------------|-----------| +| nvidia-tesla-t4-vws | 16 | +|---------------------|-----------| diff --git a/test/python/stackql_test_tooling/stackql_context.py b/test/python/stackql_test_tooling/stackql_context.py index 2c881608f..fef1e2a8d 100644 --- a/test/python/stackql_test_tooling/stackql_context.py +++ b/test/python/stackql_test_tooling/stackql_context.py @@ -644,6 +644,15 @@ def generate_password() -> str: SELECT_MACHINE_TYPES_DESC = "select name from google.compute.machineTypes where project = 'testing-project' and zone = 'australia-southeast1-a' order by name desc;" SELECT_GOOGLE_COMPUTE_INSTANCE_IAM_POLICY = "SELECT etag FROM google.compute.instances_iam_policies WHERE project = 'testing-project' AND zone = 'australia-southeast1-a' AND resource = '000000001';" + # Window function tests using accelerator types + SELECT_ACCELERATOR_TYPES_ROW_NUMBER = "SELECT name, ROW_NUMBER() OVER (ORDER BY name ASC) as row_number FROM google.compute.acceleratorTypes WHERE project = 'testing-project' AND zone = 'australia-southeast1-a' ORDER BY name ASC;" + SELECT_ACCELERATOR_TYPES_SUM_OVER = "SELECT name, SUM(maximumCardsPerInstance) OVER () as total_sum FROM google.compute.acceleratorTypes WHERE project = 'testing-project' AND zone = 'australia-southeast1-a' ORDER BY name ASC;" + SELECT_ACCELERATOR_TYPES_RUNNING_SUM = "SELECT name, SUM(maximumCardsPerInstance) OVER (ORDER BY name ASC) as running_sum FROM google.compute.acceleratorTypes WHERE project = 'testing-project' AND zone = 'australia-southeast1-a' ORDER BY name ASC;" + SELECT_ACCELERATOR_TYPES_RANK = "SELECT name, RANK() OVER (ORDER BY maximumCardsPerInstance) as rank FROM google.compute.acceleratorTypes WHERE project = 'testing-project' AND zone = 'australia-southeast1-a' ORDER BY name ASC;" + + # CTE (Common Table Expression) tests using accelerator types + SELECT_ACCELERATOR_TYPES_SIMPLE_CTE = "WITH accel_types AS (SELECT name, maximumCardsPerInstance FROM google.compute.acceleratorTypes WHERE project = 'testing-project' AND zone = 'australia-southeast1-a') SELECT name, maximumCardsPerInstance FROM accel_types ORDER BY name ASC;" + SELECT_AWS_CLOUD_CONTROL_EVENTS_MINIMAL = "SELECT DISTINCT EventTime, Identifier from aws.cloud_control.resource_requests where data__ResourceRequestStatusFilter='{}' and region = 'ap-southeast-1' order by Identifier, EventTime;" SELECT_AZURE_COMPUTE_PUBLIC_KEYS = "select id, location from azure.compute.ssh_public_keys where subscriptionId = '10001000-1000-1000-1000-100010001000' ORDER BY id ASC;" @@ -765,6 +774,15 @@ def get_native_table_count_by_name(table_name :str, sql_backend_str :str) -> str SELECT_ACCELERATOR_TYPES_DESC_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'compute-accelerator-type', 'select-zone-list-desc.txt')) + # Window function test expected results + SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'compute-accelerator-type', 'window-functions', 'row-number-by-name.txt')) + SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'compute-accelerator-type', 'window-functions', 'sum-over-all.txt')) + SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'compute-accelerator-type', 'window-functions', 'running-sum-by-name.txt')) + SELECT_ACCELERATOR_TYPES_RANK_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'compute-accelerator-type', 'window-functions', 'rank-by-cards.txt')) + + # CTE test expected results + SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'compute-accelerator-type', 'cte', 'simple-cte.txt')) + SELECT_MACHINE_TYPES_DESC_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'google', 'compute', 'instance-type-list-names-paginated-desc.txt')) SELECT_OKTA_APPS_ASC_EXPECTED = get_output_from_local_file(os.path.join('test', 'assets', 'expected', 'simple-select', 'okta', 'apps', 'select-apps-asc.txt')) @@ -987,6 +1005,16 @@ def get_registry_mock_url(execution_env :str) -> str: 'SELECT_ACCELERATOR_TYPES_DESC': SELECT_ACCELERATOR_TYPES_DESC, 'SELECT_ACCELERATOR_TYPES_DESC_EXPECTED': SELECT_ACCELERATOR_TYPES_DESC_EXPECTED, 'SELECT_ACCELERATOR_TYPES_DESC_SEQUENCE': [ SELECT_ACCELERATOR_TYPES_DESC, SELECT_ACCELERATOR_TYPES_DESC_FROM_INTEL_VIEWS, SELECT_ACCELERATOR_TYPES_DESC_FROM_INTEL_VIEWS_SUBQUERY ], + 'SELECT_ACCELERATOR_TYPES_RANK': SELECT_ACCELERATOR_TYPES_RANK, + 'SELECT_ACCELERATOR_TYPES_RANK_EXPECTED': SELECT_ACCELERATOR_TYPES_RANK_EXPECTED, + 'SELECT_ACCELERATOR_TYPES_ROW_NUMBER': SELECT_ACCELERATOR_TYPES_ROW_NUMBER, + 'SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED': SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED, + 'SELECT_ACCELERATOR_TYPES_RUNNING_SUM': SELECT_ACCELERATOR_TYPES_RUNNING_SUM, + 'SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED': SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED, + 'SELECT_ACCELERATOR_TYPES_SUM_OVER': SELECT_ACCELERATOR_TYPES_SUM_OVER, + 'SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED': SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED, + 'SELECT_ACCELERATOR_TYPES_SIMPLE_CTE': SELECT_ACCELERATOR_TYPES_SIMPLE_CTE, + 'SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED': SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED, 'SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_EXPECTED': SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_EXPECTED, 'SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_SIMPLE': SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_SIMPLE, 'SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_TRANSPARENT': SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_TRANSPARENT, diff --git a/test/robot/functional/stackql_mocked_from_cmd_line.robot b/test/robot/functional/stackql_mocked_from_cmd_line.robot index 05e7aa400..f8aa37eae 100644 --- a/test/robot/functional/stackql_mocked_from_cmd_line.robot +++ b/test/robot/functional/stackql_mocked_from_cmd_line.robot @@ -104,6 +104,66 @@ Google AcceleratorTypes SQL verb pre changeover ... ${SELECT_ACCELERATOR_TYPES_DESC} ... ${SELECT_ACCELERATOR_TYPES_DESC_EXPECTED} +Window Function ROW_NUMBER Over AcceleratorTypes + Should StackQL Exec Inline Equal + ... ${STACKQL_EXE} + ... ${OKTA_SECRET_STR} + ... ${GITHUB_SECRET_STR} + ... ${K8S_SECRET_STR} + ... ${REGISTRY_NO_VERIFY_CFG_STR} + ... ${AUTH_CFG_STR} + ... ${SQL_BACKEND_CFG_STR_CANONICAL} + ... ${SELECT_ACCELERATOR_TYPES_ROW_NUMBER} + ... ${SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED} + +Window Function SUM OVER All AcceleratorTypes + Should StackQL Exec Inline Equal + ... ${STACKQL_EXE} + ... ${OKTA_SECRET_STR} + ... ${GITHUB_SECRET_STR} + ... ${K8S_SECRET_STR} + ... ${REGISTRY_NO_VERIFY_CFG_STR} + ... ${AUTH_CFG_STR} + ... ${SQL_BACKEND_CFG_STR_CANONICAL} + ... ${SELECT_ACCELERATOR_TYPES_SUM_OVER} + ... ${SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED} + +Window Function Running SUM Over AcceleratorTypes + Should StackQL Exec Inline Equal + ... ${STACKQL_EXE} + ... ${OKTA_SECRET_STR} + ... ${GITHUB_SECRET_STR} + ... ${K8S_SECRET_STR} + ... ${REGISTRY_NO_VERIFY_CFG_STR} + ... ${AUTH_CFG_STR} + ... ${SQL_BACKEND_CFG_STR_CANONICAL} + ... ${SELECT_ACCELERATOR_TYPES_RUNNING_SUM} + ... ${SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED} + +Window Function RANK Over AcceleratorTypes + Should StackQL Exec Inline Equal + ... ${STACKQL_EXE} + ... ${OKTA_SECRET_STR} + ... ${GITHUB_SECRET_STR} + ... ${K8S_SECRET_STR} + ... ${REGISTRY_NO_VERIFY_CFG_STR} + ... ${AUTH_CFG_STR} + ... ${SQL_BACKEND_CFG_STR_CANONICAL} + ... ${SELECT_ACCELERATOR_TYPES_RANK} + ... ${SELECT_ACCELERATOR_TYPES_RANK_EXPECTED} + +CTE Simple Select Over AcceleratorTypes + Should StackQL Exec Inline Equal + ... ${STACKQL_EXE} + ... ${OKTA_SECRET_STR} + ... ${GITHUB_SECRET_STR} + ... ${K8S_SECRET_STR} + ... ${REGISTRY_NO_VERIFY_CFG_STR} + ... ${AUTH_CFG_STR} + ... ${SQL_BACKEND_CFG_STR_CANONICAL} + ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE} + ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED} + Google Machine Types Select Paginated Should Horrid Query StackQL Inline Equal ... ${STACKQL_EXE} From d674161042aaa2a84088c60fcb099924a5fee4e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Dec 2025 09:44:47 +0000 Subject: [PATCH 2/5] refactor: extract CTE processing into helper method to reduce nesting complexity Move CTE reference handling into separate processCTEReference method to satisfy nestif linter complexity requirements. --- .../astanalysis/earlyanalysis/ast_expand.go | 59 ++++++++++++------- 1 file changed, 37 insertions(+), 22 deletions(-) diff --git a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go index 2be8e080c..ff47503ef 100644 --- a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go +++ b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go @@ -180,6 +180,38 @@ func (v *indirectExpandAstVisitor) IsReadOnly() bool { return v.selectCount > 0 && v.mutateCount == 0 } +// processCTEReference handles CTE references by converting them to subquery indirects. +// Returns true if the table name was a CTE reference and was processed, false otherwise. +func (v *indirectExpandAstVisitor) processCTEReference( + node *sqlparser.AliasedTableExpr, + tableName string, +) (bool, error) { + cteSubquery, isCTE := v.cteRegistry[tableName] + if !isCTE { + return false, nil + } + // Use the CTE name as the effective alias if no explicit alias is set + effectiveAlias := node.As.GetRawVal() + if effectiveAlias == "" { + effectiveAlias = tableName + } + // Create a synthetic AliasedTableExpr with the CTE subquery + syntheticExpr := &sqlparser.AliasedTableExpr{ + Expr: cteSubquery, + As: sqlparser.NewTableIdent(effectiveAlias), + } + sq := internaldto.NewSubqueryDTO(syntheticExpr, cteSubquery) + indirect, err := astindirect.NewSubqueryIndirect(sq) + if err != nil { + return true, nil //nolint:nilerr //TODO: investigate + } + err = v.processIndirect(syntheticExpr, indirect) + if err != nil { + return true, nil //nolint:nilerr //TODO: investigate + } + return true, nil +} + func (v *indirectExpandAstVisitor) ContainsAnalyticsCacheMaterial() bool { return v.containsAnalyticsCacheMaterial } @@ -797,28 +829,11 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error { return nil case sqlparser.TableName: // Check if this table name is a CTE reference - tableName := n.GetRawVal() - if cteSubquery, isCTE := v.cteRegistry[tableName]; isCTE { - // Process CTE reference as a subquery - // Use the CTE name as the effective alias if no explicit alias is set - effectiveAlias := node.As.GetRawVal() - if effectiveAlias == "" { - effectiveAlias = tableName - } - // Create a synthetic AliasedTableExpr with the CTE subquery - syntheticExpr := &sqlparser.AliasedTableExpr{ - Expr: cteSubquery, - As: sqlparser.NewTableIdent(effectiveAlias), - } - sq := internaldto.NewSubqueryDTO(syntheticExpr, cteSubquery) - indirect, err := astindirect.NewSubqueryIndirect(sq) - if err != nil { - return nil //nolint:nilerr //TODO: investigate - } - err = v.processIndirect(syntheticExpr, indirect) - if err != nil { - return nil //nolint:nilerr //TODO: investigate - } + processed, err := v.processCTEReference(node, n.GetRawVal()) + if err != nil { + return err + } + if processed { return nil } } From 1d4d6cfedd1ae7116206d7e0ddd08f06d600576b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Dec 2025 09:54:46 +0000 Subject: [PATCH 3/5] refactor: further extract AliasedTableExpr handling to fix linter issues - Remove unused error return from processCTEReference (unparam fix) - Extract visitAliasedTableExpr helper method to reduce nesting (nestif fix) --- .../astanalysis/earlyanalysis/ast_expand.go | 99 +++++++++---------- 1 file changed, 46 insertions(+), 53 deletions(-) diff --git a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go index ff47503ef..cf4517a96 100644 --- a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go +++ b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go @@ -185,10 +185,10 @@ func (v *indirectExpandAstVisitor) IsReadOnly() bool { func (v *indirectExpandAstVisitor) processCTEReference( node *sqlparser.AliasedTableExpr, tableName string, -) (bool, error) { +) bool { cteSubquery, isCTE := v.cteRegistry[tableName] if !isCTE { - return false, nil + return false } // Use the CTE name as the effective alias if no explicit alias is set effectiveAlias := node.As.GetRawVal() @@ -203,13 +203,50 @@ func (v *indirectExpandAstVisitor) processCTEReference( sq := internaldto.NewSubqueryDTO(syntheticExpr, cteSubquery) indirect, err := astindirect.NewSubqueryIndirect(sq) if err != nil { - return true, nil //nolint:nilerr //TODO: investigate + return true //nolint:nilerr //TODO: investigate } - err = v.processIndirect(syntheticExpr, indirect) - if err != nil { - return true, nil //nolint:nilerr //TODO: investigate + _ = v.processIndirect(syntheticExpr, indirect) //nolint:errcheck // errors handled via indirect pattern + return true +} + +// visitAliasedTableExpr handles visiting an AliasedTableExpr node, including +// subqueries, CTE references, and regular table expressions. +func (v *indirectExpandAstVisitor) visitAliasedTableExpr(node *sqlparser.AliasedTableExpr) error { + if node.Expr != nil { + switch n := node.Expr.(type) { //nolint:gocritic // switch preferred for type assertions + case *sqlparser.Subquery: + sq := internaldto.NewSubqueryDTO(node, n) + indirect, err := astindirect.NewSubqueryIndirect(sq) + if err != nil { + return nil //nolint:nilerr //TODO: investigate + } + _ = v.processIndirect(node, indirect) //nolint:errcheck // errors handled via indirect pattern + return nil + case sqlparser.TableName: + if v.processCTEReference(node, n.GetRawVal()) { + return nil + } + } + if err := node.Expr.Accept(v); err != nil { + return err + } } - return true, nil + if node.Partitions != nil { + if err := node.Partitions.Accept(v); err != nil { + return err + } + } + if !node.As.IsEmpty() { + if err := node.As.Accept(v); err != nil { + return err + } + } + if node.Hints != nil { + if err := node.Hints.Accept(v); err != nil { + return err + } + } + return nil } func (v *indirectExpandAstVisitor) ContainsAnalyticsCacheMaterial() bool { @@ -813,52 +850,8 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error { } case *sqlparser.AliasedTableExpr: - if node.Expr != nil { - //nolint:gocritic // deferring cosmetics on visitors - switch n := node.Expr.(type) { - case *sqlparser.Subquery: - sq := internaldto.NewSubqueryDTO(node, n) - indirect, err := astindirect.NewSubqueryIndirect(sq) - if err != nil { - return nil //nolint:nilerr //TODO: investigate - } - err = v.processIndirect(node, indirect) - if err != nil { - return nil //nolint:nilerr //TODO: investigate - } - return nil - case sqlparser.TableName: - // Check if this table name is a CTE reference - processed, err := v.processCTEReference(node, n.GetRawVal()) - if err != nil { - return err - } - if processed { - return nil - } - } - err := node.Expr.Accept(v) - if err != nil { - return err - } - } - if node.Partitions != nil { - err := node.Partitions.Accept(v) - if err != nil { - return err - } - } - if !node.As.IsEmpty() { - err := node.As.Accept(v) - if err != nil { - return err - } - } - if node.Hints != nil { - err := node.Hints.Accept(v) - if err != nil { - return err - } + if err := v.visitAliasedTableExpr(node); err != nil { + return err } case sqlparser.TableNames: From 6c5f6c8009f6e3fed58ade20139ca49accaa7653 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Dec 2025 10:46:48 +0000 Subject: [PATCH 4/5] test: temporarily disable CTE test for debugging --- .../stackql_mocked_from_cmd_line.robot | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/test/robot/functional/stackql_mocked_from_cmd_line.robot b/test/robot/functional/stackql_mocked_from_cmd_line.robot index f8aa37eae..226bb8091 100644 --- a/test/robot/functional/stackql_mocked_from_cmd_line.robot +++ b/test/robot/functional/stackql_mocked_from_cmd_line.robot @@ -152,17 +152,18 @@ Window Function RANK Over AcceleratorTypes ... ${SELECT_ACCELERATOR_TYPES_RANK} ... ${SELECT_ACCELERATOR_TYPES_RANK_EXPECTED} -CTE Simple Select Over AcceleratorTypes - Should StackQL Exec Inline Equal - ... ${STACKQL_EXE} - ... ${OKTA_SECRET_STR} - ... ${GITHUB_SECRET_STR} - ... ${K8S_SECRET_STR} - ... ${REGISTRY_NO_VERIFY_CFG_STR} - ... ${AUTH_CFG_STR} - ... ${SQL_BACKEND_CFG_STR_CANONICAL} - ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE} - ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED} +# CTE test temporarily disabled for debugging - CTE support may need more work +# CTE Simple Select Over AcceleratorTypes +# Should StackQL Exec Inline Equal +# ... ${STACKQL_EXE} +# ... ${OKTA_SECRET_STR} +# ... ${GITHUB_SECRET_STR} +# ... ${K8S_SECRET_STR} +# ... ${REGISTRY_NO_VERIFY_CFG_STR} +# ... ${AUTH_CFG_STR} +# ... ${SQL_BACKEND_CFG_STR_CANONICAL} +# ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE} +# ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED} Google Machine Types Select Paginated Should Horrid Query StackQL Inline Equal From b226c63d524446bbcaa7c865e95d7f7ef7ef0fa6 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Dec 2025 19:49:44 +0000 Subject: [PATCH 5/5] fix: CTE table references now correctly modify AST for downstream resolution The CTE implementation was storing an indirect but not modifying the original AST node. This caused GetHIDs to still see a TableName type instead of a Subquery type, resulting in "could not locate table" errors. Fix: Modify the original node.Expr to be the CTE's Subquery instead of creating a synthetic expression. This ensures downstream code correctly recognizes the CTE reference as a subquery. Also re-enables the CTE robot test. --- .../astanalysis/earlyanalysis/ast_expand.go | 19 +++++++-------- .../stackql_mocked_from_cmd_line.robot | 23 +++++++++---------- 2 files changed, 19 insertions(+), 23 deletions(-) diff --git a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go index cf4517a96..d0551dd5e 100644 --- a/internal/stackql/astanalysis/earlyanalysis/ast_expand.go +++ b/internal/stackql/astanalysis/earlyanalysis/ast_expand.go @@ -190,22 +190,19 @@ func (v *indirectExpandAstVisitor) processCTEReference( if !isCTE { return false } - // Use the CTE name as the effective alias if no explicit alias is set - effectiveAlias := node.As.GetRawVal() - if effectiveAlias == "" { - effectiveAlias = tableName + // Modify the original node to replace the TableName with the CTE subquery + // This is critical: downstream code (GetHIDs) checks node.Expr type to identify subqueries + node.Expr = cteSubquery + // Set the alias to the CTE name if no explicit alias was provided + if node.As.IsEmpty() { + node.As = sqlparser.NewTableIdent(tableName) } - // Create a synthetic AliasedTableExpr with the CTE subquery - syntheticExpr := &sqlparser.AliasedTableExpr{ - Expr: cteSubquery, - As: sqlparser.NewTableIdent(effectiveAlias), - } - sq := internaldto.NewSubqueryDTO(syntheticExpr, cteSubquery) + sq := internaldto.NewSubqueryDTO(node, cteSubquery) indirect, err := astindirect.NewSubqueryIndirect(sq) if err != nil { return true //nolint:nilerr //TODO: investigate } - _ = v.processIndirect(syntheticExpr, indirect) //nolint:errcheck // errors handled via indirect pattern + _ = v.processIndirect(node, indirect) //nolint:errcheck // errors handled via indirect pattern return true } diff --git a/test/robot/functional/stackql_mocked_from_cmd_line.robot b/test/robot/functional/stackql_mocked_from_cmd_line.robot index 226bb8091..f8aa37eae 100644 --- a/test/robot/functional/stackql_mocked_from_cmd_line.robot +++ b/test/robot/functional/stackql_mocked_from_cmd_line.robot @@ -152,18 +152,17 @@ Window Function RANK Over AcceleratorTypes ... ${SELECT_ACCELERATOR_TYPES_RANK} ... ${SELECT_ACCELERATOR_TYPES_RANK_EXPECTED} -# CTE test temporarily disabled for debugging - CTE support may need more work -# CTE Simple Select Over AcceleratorTypes -# Should StackQL Exec Inline Equal -# ... ${STACKQL_EXE} -# ... ${OKTA_SECRET_STR} -# ... ${GITHUB_SECRET_STR} -# ... ${K8S_SECRET_STR} -# ... ${REGISTRY_NO_VERIFY_CFG_STR} -# ... ${AUTH_CFG_STR} -# ... ${SQL_BACKEND_CFG_STR_CANONICAL} -# ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE} -# ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED} +CTE Simple Select Over AcceleratorTypes + Should StackQL Exec Inline Equal + ... ${STACKQL_EXE} + ... ${OKTA_SECRET_STR} + ... ${GITHUB_SECRET_STR} + ... ${K8S_SECRET_STR} + ... ${REGISTRY_NO_VERIFY_CFG_STR} + ... ${AUTH_CFG_STR} + ... ${SQL_BACKEND_CFG_STR_CANONICAL} + ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE} + ... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED} Google Machine Types Select Paginated Should Horrid Query StackQL Inline Equal