Skip to content

Commit 45ab752

Browse files
claudeclaude
committed
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
1 parent e755b7a commit 45ab752

File tree

8 files changed

+179
-0
lines changed

8 files changed

+179
-0
lines changed

‎internal/stackql/astanalysis/earlyanalysis/ast_expand.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ type indirectExpandAstVisitor struct {
4949
selectCount int
5050
mutateCount int
5151
createBuilder []primitivebuilder.Builder
52+
cteRegistry map[string]*sqlparser.Subquery // CTE name -> subquery definition
5253
}
5354

5455
func newIndirectExpandAstVisitor(
@@ -75,6 +76,7 @@ func newIndirectExpandAstVisitor(
7576
tcc: tcc,
7677
whereParams: whereParams,
7778
indirectionDepth: indirectionDepth,
79+
cteRegistry: make(map[string]*sqlparser.Subquery),
7880
}
7981
return rv, nil
8082
}
@@ -214,6 +216,14 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error {
214216
addIf(node.StraightJoinHint, sqlparser.StraightJoinHint)
215217
addIf(node.SQLCalcFoundRows, sqlparser.SQLCalcFoundRowsStr)
216218

219+
// Extract CTEs from WITH clause and store in registry
220+
if node.With != nil {
221+
for _, cte := range node.With.CTEs {
222+
cteName := cte.Name.GetRawVal()
223+
v.cteRegistry[cteName] = cte.Subquery
224+
}
225+
}
226+
217227
if node.Comments != nil {
218228
node.Comments.Accept(v) //nolint:errcheck // future proof
219229
}
@@ -785,6 +795,32 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error {
785795
return nil //nolint:nilerr //TODO: investigate
786796
}
787797
return nil
798+
case sqlparser.TableName:
799+
// Check if this table name is a CTE reference
800+
tableName := n.GetRawVal()
801+
if cteSubquery, isCTE := v.cteRegistry[tableName]; isCTE {
802+
// Process CTE reference as a subquery
803+
// Use the CTE name as the effective alias if no explicit alias is set
804+
effectiveAlias := node.As.GetRawVal()
805+
if effectiveAlias == "" {
806+
effectiveAlias = tableName
807+
}
808+
// Create a synthetic AliasedTableExpr with the CTE subquery
809+
syntheticExpr := &sqlparser.AliasedTableExpr{
810+
Expr: cteSubquery,
811+
As: sqlparser.NewTableIdent(effectiveAlias),
812+
}
813+
sq := internaldto.NewSubqueryDTO(syntheticExpr, cteSubquery)
814+
indirect, err := astindirect.NewSubqueryIndirect(sq)
815+
if err != nil {
816+
return nil //nolint:nilerr //TODO: investigate
817+
}
818+
err = v.processIndirect(syntheticExpr, indirect)
819+
if err != nil {
820+
return nil //nolint:nilerr //TODO: investigate
821+
}
822+
return nil
823+
}
788824
}
789825
err := node.Expr.Accept(v)
790826
if err != nil {

‎test/assets/expected/simple-select/compute-accelerator-type/cte/simple-cte.txt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
|---------------------|-------------------------|
2+
| name | maximumCardsPerInstance |
3+
|---------------------|-------------------------|
4+
| nvidia-tesla-p4 | 4 |
5+
|---------------------|-------------------------|
6+
| nvidia-tesla-p4-vws | 4 |
7+
|---------------------|-------------------------|
8+
| nvidia-tesla-t4 | 4 |
9+
|---------------------|-------------------------|
10+
| nvidia-tesla-t4-vws | 4 |
11+
|---------------------|-------------------------|

‎test/assets/expected/simple-select/compute-accelerator-type/window-functions/rank-by-cards.txt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
|---------------------|------|
2+
| name | rank |
3+
|---------------------|------|
4+
| nvidia-tesla-p4 | 1 |
5+
|---------------------|------|
6+
| nvidia-tesla-p4-vws | 1 |
7+
|---------------------|------|
8+
| nvidia-tesla-t4 | 1 |
9+
|---------------------|------|
10+
| nvidia-tesla-t4-vws | 1 |
11+
|---------------------|------|

‎test/assets/expected/simple-select/compute-accelerator-type/window-functions/row-number-by-name.txt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
|---------------------|------------|
2+
| name | row_number |
3+
|---------------------|------------|
4+
| nvidia-tesla-p4 | 1 |
5+
|---------------------|------------|
6+
| nvidia-tesla-p4-vws | 2 |
7+
|---------------------|------------|
8+
| nvidia-tesla-t4 | 3 |
9+
|---------------------|------------|
10+
| nvidia-tesla-t4-vws | 4 |
11+
|---------------------|------------|

‎test/assets/expected/simple-select/compute-accelerator-type/window-functions/running-sum-by-name.txt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
|---------------------|-------------|
2+
| name | running_sum |
3+
|---------------------|-------------|
4+
| nvidia-tesla-p4 | 4 |
5+
|---------------------|-------------|
6+
| nvidia-tesla-p4-vws | 8 |
7+
|---------------------|-------------|
8+
| nvidia-tesla-t4 | 12 |
9+
|---------------------|-------------|
10+
| nvidia-tesla-t4-vws | 16 |
11+
|---------------------|-------------|

‎test/assets/expected/simple-select/compute-accelerator-type/window-functions/sum-over-all.txt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
|---------------------|-----------|
2+
| name | total_sum |
3+
|---------------------|-----------|
4+
| nvidia-tesla-p4 | 16 |
5+
|---------------------|-----------|
6+
| nvidia-tesla-p4-vws | 16 |
7+
|---------------------|-----------|
8+
| nvidia-tesla-t4 | 16 |
9+
|---------------------|-----------|
10+
| nvidia-tesla-t4-vws | 16 |
11+
|---------------------|-----------|

‎test/python/stackql_test_tooling/stackql_context.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -644,6 +644,15 @@ def generate_password() -> str:
644644
SELECT_MACHINE_TYPES_DESC = "select name from google.compute.machineTypes where project = 'testing-project' and zone = 'australia-southeast1-a' order by name desc;"
645645
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';"
646646

647+
# Window function tests using accelerator types
648+
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;"
649+
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;"
650+
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;"
651+
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;"
652+
653+
# CTE (Common Table Expression) tests using accelerator types
654+
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;"
655+
647656
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;"
648657

649658
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
765774

766775
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'))
767776

777+
# Window function test expected results
778+
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'))
779+
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'))
780+
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'))
781+
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'))
782+
783+
# CTE test expected results
784+
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'))
785+
768786
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'))
769787

770788
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:
9871005
'SELECT_ACCELERATOR_TYPES_DESC': SELECT_ACCELERATOR_TYPES_DESC,
9881006
'SELECT_ACCELERATOR_TYPES_DESC_EXPECTED': SELECT_ACCELERATOR_TYPES_DESC_EXPECTED,
9891007
'SELECT_ACCELERATOR_TYPES_DESC_SEQUENCE': [ SELECT_ACCELERATOR_TYPES_DESC, SELECT_ACCELERATOR_TYPES_DESC_FROM_INTEL_VIEWS, SELECT_ACCELERATOR_TYPES_DESC_FROM_INTEL_VIEWS_SUBQUERY ],
1008+
'SELECT_ACCELERATOR_TYPES_RANK': SELECT_ACCELERATOR_TYPES_RANK,
1009+
'SELECT_ACCELERATOR_TYPES_RANK_EXPECTED': SELECT_ACCELERATOR_TYPES_RANK_EXPECTED,
1010+
'SELECT_ACCELERATOR_TYPES_ROW_NUMBER': SELECT_ACCELERATOR_TYPES_ROW_NUMBER,
1011+
'SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED': SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED,
1012+
'SELECT_ACCELERATOR_TYPES_RUNNING_SUM': SELECT_ACCELERATOR_TYPES_RUNNING_SUM,
1013+
'SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED': SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED,
1014+
'SELECT_ACCELERATOR_TYPES_SUM_OVER': SELECT_ACCELERATOR_TYPES_SUM_OVER,
1015+
'SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED': SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED,
1016+
'SELECT_ACCELERATOR_TYPES_SIMPLE_CTE': SELECT_ACCELERATOR_TYPES_SIMPLE_CTE,
1017+
'SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED': SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED,
9901018
'SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_EXPECTED': SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_EXPECTED,
9911019
'SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_SIMPLE': SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_SIMPLE,
9921020
'SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_TRANSPARENT': SELECT_ANALYTICS_CACHE_GITHUB_REPOSITORIES_COLLABORATORS_TRANSPARENT,

‎test/robot/functional/stackql_mocked_from_cmd_line.robot‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,66 @@ Google AcceleratorTypes SQL verb pre changeover
104104
... ${SELECT_ACCELERATOR_TYPES_DESC}
105105
... ${SELECT_ACCELERATOR_TYPES_DESC_EXPECTED}
106106

107+
Window Function ROW_NUMBER Over AcceleratorTypes
108+
Should StackQL Exec Inline Equal
109+
... ${STACKQL_EXE}
110+
... ${OKTA_SECRET_STR}
111+
... ${GITHUB_SECRET_STR}
112+
... ${K8S_SECRET_STR}
113+
... ${REGISTRY_NO_VERIFY_CFG_STR}
114+
... ${AUTH_CFG_STR}
115+
... ${SQL_BACKEND_CFG_STR_CANONICAL}
116+
... ${SELECT_ACCELERATOR_TYPES_ROW_NUMBER}
117+
... ${SELECT_ACCELERATOR_TYPES_ROW_NUMBER_EXPECTED}
118+
119+
Window Function SUM OVER All AcceleratorTypes
120+
Should StackQL Exec Inline Equal
121+
... ${STACKQL_EXE}
122+
... ${OKTA_SECRET_STR}
123+
... ${GITHUB_SECRET_STR}
124+
... ${K8S_SECRET_STR}
125+
... ${REGISTRY_NO_VERIFY_CFG_STR}
126+
... ${AUTH_CFG_STR}
127+
... ${SQL_BACKEND_CFG_STR_CANONICAL}
128+
... ${SELECT_ACCELERATOR_TYPES_SUM_OVER}
129+
... ${SELECT_ACCELERATOR_TYPES_SUM_OVER_EXPECTED}
130+
131+
Window Function Running SUM Over AcceleratorTypes
132+
Should StackQL Exec Inline Equal
133+
... ${STACKQL_EXE}
134+
... ${OKTA_SECRET_STR}
135+
... ${GITHUB_SECRET_STR}
136+
... ${K8S_SECRET_STR}
137+
... ${REGISTRY_NO_VERIFY_CFG_STR}
138+
... ${AUTH_CFG_STR}
139+
... ${SQL_BACKEND_CFG_STR_CANONICAL}
140+
... ${SELECT_ACCELERATOR_TYPES_RUNNING_SUM}
141+
... ${SELECT_ACCELERATOR_TYPES_RUNNING_SUM_EXPECTED}
142+
143+
Window Function RANK Over AcceleratorTypes
144+
Should StackQL Exec Inline Equal
145+
... ${STACKQL_EXE}
146+
... ${OKTA_SECRET_STR}
147+
... ${GITHUB_SECRET_STR}
148+
... ${K8S_SECRET_STR}
149+
... ${REGISTRY_NO_VERIFY_CFG_STR}
150+
... ${AUTH_CFG_STR}
151+
... ${SQL_BACKEND_CFG_STR_CANONICAL}
152+
... ${SELECT_ACCELERATOR_TYPES_RANK}
153+
... ${SELECT_ACCELERATOR_TYPES_RANK_EXPECTED}
154+
155+
CTE Simple Select Over AcceleratorTypes
156+
Should StackQL Exec Inline Equal
157+
... ${STACKQL_EXE}
158+
... ${OKTA_SECRET_STR}
159+
... ${GITHUB_SECRET_STR}
160+
... ${K8S_SECRET_STR}
161+
... ${REGISTRY_NO_VERIFY_CFG_STR}
162+
... ${AUTH_CFG_STR}
163+
... ${SQL_BACKEND_CFG_STR_CANONICAL}
164+
... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE}
165+
... ${SELECT_ACCELERATOR_TYPES_SIMPLE_CTE_EXPECTED}
166+
107167
Google Machine Types Select Paginated
108168
Should Horrid Query StackQL Inline Equal
109169
... ${STACKQL_EXE}

0 commit comments

Comments
 (0)