Skip to content

Commit a6bcfc3

Browse files
jeffreyavenjeffreyaven
authored
Merge pull request #115 from stackql/claude/fix-window-function-execution-01WbDt4A6kMPzL5KJmLWAHBU
feat: add CTE support and robot tests for window functions
2 parents 407ceaf + b226c63 commit a6bcfc3

File tree

8 files changed

+221
-37
lines changed

8 files changed

+221
-37
lines changed

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

Lines changed: 78 additions & 37 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
}
@@ -178,6 +180,72 @@ func (v *indirectExpandAstVisitor) IsReadOnly() bool {
178180
return v.selectCount > 0 && v.mutateCount == 0
179181
}
180182

183+
// processCTEReference handles CTE references by converting them to subquery indirects.
184+
// Returns true if the table name was a CTE reference and was processed, false otherwise.
185+
func (v *indirectExpandAstVisitor) processCTEReference(
186+
node *sqlparser.AliasedTableExpr,
187+
tableName string,
188+
) bool {
189+
cteSubquery, isCTE := v.cteRegistry[tableName]
190+
if !isCTE {
191+
return false
192+
}
193+
// Modify the original node to replace the TableName with the CTE subquery
194+
// This is critical: downstream code (GetHIDs) checks node.Expr type to identify subqueries
195+
node.Expr = cteSubquery
196+
// Set the alias to the CTE name if no explicit alias was provided
197+
if node.As.IsEmpty() {
198+
node.As = sqlparser.NewTableIdent(tableName)
199+
}
200+
sq := internaldto.NewSubqueryDTO(node, cteSubquery)
201+
indirect, err := astindirect.NewSubqueryIndirect(sq)
202+
if err != nil {
203+
return true //nolint:nilerr //TODO: investigate
204+
}
205+
_ = v.processIndirect(node, indirect) //nolint:errcheck // errors handled via indirect pattern
206+
return true
207+
}
208+
209+
// visitAliasedTableExpr handles visiting an AliasedTableExpr node, including
210+
// subqueries, CTE references, and regular table expressions.
211+
func (v *indirectExpandAstVisitor) visitAliasedTableExpr(node *sqlparser.AliasedTableExpr) error {
212+
if node.Expr != nil {
213+
switch n := node.Expr.(type) { //nolint:gocritic // switch preferred for type assertions
214+
case *sqlparser.Subquery:
215+
sq := internaldto.NewSubqueryDTO(node, n)
216+
indirect, err := astindirect.NewSubqueryIndirect(sq)
217+
if err != nil {
218+
return nil //nolint:nilerr //TODO: investigate
219+
}
220+
_ = v.processIndirect(node, indirect) //nolint:errcheck // errors handled via indirect pattern
221+
return nil
222+
case sqlparser.TableName:
223+
if v.processCTEReference(node, n.GetRawVal()) {
224+
return nil
225+
}
226+
}
227+
if err := node.Expr.Accept(v); err != nil {
228+
return err
229+
}
230+
}
231+
if node.Partitions != nil {
232+
if err := node.Partitions.Accept(v); err != nil {
233+
return err
234+
}
235+
}
236+
if !node.As.IsEmpty() {
237+
if err := node.As.Accept(v); err != nil {
238+
return err
239+
}
240+
}
241+
if node.Hints != nil {
242+
if err := node.Hints.Accept(v); err != nil {
243+
return err
244+
}
245+
}
246+
return nil
247+
}
248+
181249
func (v *indirectExpandAstVisitor) ContainsAnalyticsCacheMaterial() bool {
182250
return v.containsAnalyticsCacheMaterial
183251
}
@@ -214,6 +282,14 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error {
214282
addIf(node.StraightJoinHint, sqlparser.StraightJoinHint)
215283
addIf(node.SQLCalcFoundRows, sqlparser.SQLCalcFoundRowsStr)
216284

285+
// Extract CTEs from WITH clause and store in registry
286+
if node.With != nil {
287+
for _, cte := range node.With.CTEs {
288+
cteName := cte.Name.GetRawVal()
289+
v.cteRegistry[cteName] = cte.Subquery
290+
}
291+
}
292+
217293
if node.Comments != nil {
218294
node.Comments.Accept(v) //nolint:errcheck // future proof
219295
}
@@ -771,43 +847,8 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error {
771847
}
772848

773849
case *sqlparser.AliasedTableExpr:
774-
if node.Expr != nil {
775-
//nolint:gocritic // deferring cosmetics on visitors
776-
switch n := node.Expr.(type) {
777-
case *sqlparser.Subquery:
778-
sq := internaldto.NewSubqueryDTO(node, n)
779-
indirect, err := astindirect.NewSubqueryIndirect(sq)
780-
if err != nil {
781-
return nil //nolint:nilerr //TODO: investigate
782-
}
783-
err = v.processIndirect(node, indirect)
784-
if err != nil {
785-
return nil //nolint:nilerr //TODO: investigate
786-
}
787-
return nil
788-
}
789-
err := node.Expr.Accept(v)
790-
if err != nil {
791-
return err
792-
}
793-
}
794-
if node.Partitions != nil {
795-
err := node.Partitions.Accept(v)
796-
if err != nil {
797-
return err
798-
}
799-
}
800-
if !node.As.IsEmpty() {
801-
err := node.As.Accept(v)
802-
if err != nil {
803-
return err
804-
}
805-
}
806-
if node.Hints != nil {
807-
err := node.Hints.Accept(v)
808-
if err != nil {
809-
return err
810-
}
850+
if err := v.visitAliasedTableExpr(node); err != nil {
851+
return err
811852
}
812853

813854
case sqlparser.TableNames:

‎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)