Merge pull request #594 from stackql/claude/fix-stackql-tests-01FaQ7YN6kjqRvqDwtymJdDR

Add window function (OVER clause) and CTE (WITH clause)

Author: jeffreyaven

.github/workflows/lint.yml

@@ -33,6 +33,8 @@ jobs:
         with:
           args: -color
 
+      - name: Update Go dependencies
+        run: go mod tidy
 
       - name: Run golangci-lint
         uses: golangci/golangci-lint-action@v8.0.0

go.mod

@@ -20,7 +20,7 @@ require (
 	github.com/stackql/any-sdk v0.3.1-beta01
 	github.com/stackql/go-suffix-map v0.0.1-alpha01
 	github.com/stackql/psql-wire v0.1.2-alpha01
-	github.com/stackql/stackql-parser v0.0.15-alpha06
+	github.com/stackql/stackql-parser v0.0.16-alpha01
 	github.com/stretchr/testify v1.10.0
 	golang.org/x/sync v0.15.0
 	gonum.org/v1/gonum v0.15.1

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
 }
@@ -103,6 +105,41 @@ func (v *indirectExpandAstVisitor) processMaterializedView(
 	return nil
 }
 
+// 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 {
+	cteSubquery, isCTE := v.cteRegistry[tableName]
+	if !isCTE {
+		return false
+	}
+	logging.GetLogger().Infof("processCTEReference: Converting CTE '%s' to subquery", tableName)
+	logging.GetLogger().Debugf("processCTEReference: CTE subquery = %s", sqlparser.String(cteSubquery))
+	// 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)
+	}
+	logging.GetLogger().Debugf("processCTEReference: Node alias set to '%s'", node.As.GetRawVal())
+	sq := internaldto.NewSubqueryDTO(node, cteSubquery)
+	indirect, err := astindirect.NewSubqueryIndirect(sq)
+	if err != nil {
+		logging.GetLogger().Errorf("processCTEReference: Failed to create subquery indirect: %v", err)
+		return true //nolint:nilerr //TODO: investigate
+	}
+	err = v.processIndirect(node, indirect)
+	if err != nil {
+		logging.GetLogger().Errorf("processCTEReference: processIndirect failed: %v", err)
+	} else {
+		logging.GetLogger().Infof("processCTEReference: Successfully processed CTE '%s' as subquery", tableName)
+	}
+	return true
+}
+
 func (v *indirectExpandAstVisitor) processIndirect(node sqlparser.SQLNode, indirect astindirect.Indirect) error {
 	err := indirect.Parse()
 	if err != nil {
@@ -214,6 +251,19 @@ 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 as Subqueries.
+		// CTEs are converted to subqueries at the AST level for uniform handling.
+		if node.With != nil {
+			logging.GetLogger().Infof("Registering %d CTEs from WITH clause", len(node.With.CTEs))
+			for _, cte := range node.With.CTEs {
+				cteName := cte.Name.GetRawVal()
+				// Wrap the CTE's SELECT statement in a Subquery struct
+				cteSubquery := &sqlparser.Subquery{Select: cte.Select}
+				v.cteRegistry[cteName] = cteSubquery
+				logging.GetLogger().Debugf("Registered CTE '%s' with subquery: %s", cteName, sqlparser.String(cteSubquery))
+			}
+		}
+
 		if node.Comments != nil {
 			node.Comments.Accept(v) //nolint:errcheck // future proof
 		}
@@ -785,6 +835,11 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error {
 					return nil //nolint:nilerr //TODO: investigate
 				}
 				return nil
+			case sqlparser.TableName:
+				// Check if this is a CTE reference - convert to subquery
+				if v.processCTEReference(node, n.GetRawVal()) {
+					return nil
+				}
 			}
 			err := node.Expr.Accept(v)
 			if err != nil {
@@ -822,6 +877,9 @@ func (v *indirectExpandAstVisitor) Visit(node sqlparser.SQLNode) error {
 		if node.IsEmpty() {
 			return nil
 		}
+		// Note: CTE references are handled in AliasedTableExpr case above,
+		// where they are converted to subqueries. This case only handles
+		// regular table names (provider.service.resource).
 		containsBackendMaterial := v.handlerCtx.GetDBMSInternalRouter().ExprIsRoutable(node)
 		if containsBackendMaterial {
 			v.containsNativeBackendMaterial = true

internal/stackql/astanalysis/earlyanalysis/cte_test.go

@@ -0,0 +1,228 @@
+package earlyanalysis_test
+
+import (
+	"testing"
+
+	"github.com/stackql/stackql-parser/go/vt/sqlparser"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCTEParsing(t *testing.T) {
+	t.Run("Simple CTE is parsed correctly", func(t *testing.T) {
+		query := "WITH cte AS (SELECT id, name FROM users) SELECT * FROM cte"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel, ok := stmt.(*sqlparser.Select)
+		require.True(t, ok, "Statement should be a SELECT")
+
+		// Check that With clause exists
+		require.NotNil(t, sel.With, "WITH clause should exist")
+		require.Len(t, sel.With.CTEs, 1, "Should have 1 CTE")
+
+		// Check CTE name
+		cte := sel.With.CTEs[0]
+		assert.Equal(t, "cte", cte.Name.GetRawVal(), "CTE name should be 'cte'")
+
+		// Check that CTE has a select statement
+		require.NotNil(t, cte.Select, "CTE should have a select statement")
+	})
+
+	t.Run("Multiple CTEs are parsed correctly", func(t *testing.T) {
+		query := "WITH a AS (SELECT 1 as x), b AS (SELECT 2 as y) SELECT * FROM a, b"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel, ok := stmt.(*sqlparser.Select)
+		require.True(t, ok, "Statement should be a SELECT")
+
+		// Check that With clause exists
+		require.NotNil(t, sel.With, "WITH clause should exist")
+		require.Len(t, sel.With.CTEs, 2, "Should have 2 CTEs")
+
+		// Check CTE names
+		assert.Equal(t, "a", sel.With.CTEs[0].Name.GetRawVal(), "First CTE name should be 'a'")
+		assert.Equal(t, "b", sel.With.CTEs[1].Name.GetRawVal(), "Second CTE name should be 'b'")
+	})
+
+	t.Run("Recursive CTE is parsed correctly", func(t *testing.T) {
+		query := "WITH RECURSIVE cte AS (SELECT 1 as n UNION ALL SELECT n + 1 FROM cte WHERE n < 10) SELECT * FROM cte"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel, ok := stmt.(*sqlparser.Select)
+		require.True(t, ok, "Statement should be a SELECT")
+
+		// Check that With clause exists with Recursive flag
+		require.NotNil(t, sel.With, "WITH clause should exist")
+		assert.True(t, sel.With.Recursive, "WITH clause should be RECURSIVE")
+		require.Len(t, sel.With.CTEs, 1, "Should have 1 CTE")
+
+		// Check CTE name
+		assert.Equal(t, "cte", sel.With.CTEs[0].Name.GetRawVal(), "CTE name should be 'cte'")
+	})
+
+	t.Run("CTE with column aliases", func(t *testing.T) {
+		query := "WITH cte(col1, col2) AS (SELECT id, name FROM users) SELECT * FROM cte"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel, ok := stmt.(*sqlparser.Select)
+		require.True(t, ok, "Statement should be a SELECT")
+
+		require.NotNil(t, sel.With, "WITH clause should exist")
+		require.Len(t, sel.With.CTEs, 1, "Should have 1 CTE")
+
+		cte := sel.With.CTEs[0]
+		assert.Equal(t, "cte", cte.Name.GetRawVal(), "CTE name should be 'cte'")
+
+		// Check column aliases if present
+		require.Len(t, cte.Columns, 2, "CTE should have 2 column aliases")
+		assert.Equal(t, "col1", cte.Columns[0].GetRawVal(), "First column alias should be 'col1'")
+		assert.Equal(t, "col2", cte.Columns[1].GetRawVal(), "Second column alias should be 'col2'")
+	})
+
+	t.Run("Nested CTEs - CTE referencing another CTE", func(t *testing.T) {
+		query := "WITH a AS (SELECT 1 as x), b AS (SELECT x * 2 as y FROM a) SELECT * FROM b"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel, ok := stmt.(*sqlparser.Select)
+		require.True(t, ok, "Statement should be a SELECT")
+
+		require.NotNil(t, sel.With, "WITH clause should exist")
+		require.Len(t, sel.With.CTEs, 2, "Should have 2 CTEs")
+	})
+}
+
+func TestCTERegistry(t *testing.T) {
+	t.Run("CTE registry stores CTEs correctly", func(t *testing.T) {
+		registry := make(map[string]*sqlparser.CommonTableExpr)
+
+		query := "WITH cte1 AS (SELECT 1), cte2 AS (SELECT 2) SELECT * FROM cte1, cte2"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel := stmt.(*sqlparser.Select)
+		require.NotNil(t, sel.With)
+
+		// Simulate what the visitor does - register CTEs
+		for _, cte := range sel.With.CTEs {
+			cteName := cte.Name.GetRawVal()
+			registry[cteName] = cte
+		}
+
+		// Verify registry contents
+		assert.Len(t, registry, 2, "Registry should have 2 CTEs")
+		assert.Contains(t, registry, "cte1", "Registry should contain 'cte1'")
+		assert.Contains(t, registry, "cte2", "Registry should contain 'cte2'")
+	})
+
+	t.Run("CTE lookup works correctly", func(t *testing.T) {
+		registry := make(map[string]*sqlparser.CommonTableExpr)
+
+		query := "WITH my_cte AS (SELECT id, name FROM users) SELECT * FROM my_cte"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel := stmt.(*sqlparser.Select)
+		require.NotNil(t, sel.With)
+
+		// Register the CTE
+		for _, cte := range sel.With.CTEs {
+			cteName := cte.Name.GetRawVal()
+			registry[cteName] = cte
+		}
+
+		// Verify we can look up the CTE
+		_, isCTE := registry["my_cte"]
+		assert.True(t, isCTE, "'my_cte' should be found in registry")
+
+		// Verify non-CTE names are not found
+		_, isNotCTE := registry["users"]
+		assert.False(t, isNotCTE, "'users' should not be found in registry")
+	})
+}
+
+func TestWindowFunctionParsing(t *testing.T) {
+	t.Run("Window function with OVER clause is parsed correctly", func(t *testing.T) {
+		query := "SELECT ROW_NUMBER() OVER (ORDER BY id) as row_num FROM t"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel, ok := stmt.(*sqlparser.Select)
+		require.True(t, ok, "Statement should be a SELECT")
+
+		require.Len(t, sel.SelectExprs, 1, "Should have 1 select expression")
+
+		aliased, ok := sel.SelectExprs[0].(*sqlparser.AliasedExpr)
+		require.True(t, ok, "Select expression should be aliased")
+
+		funcExpr, ok := aliased.Expr.(*sqlparser.FuncExpr)
+		require.True(t, ok, "Expression should be a FuncExpr")
+
+		assert.Equal(t, "row_number", funcExpr.Name.Lowered(), "Function name should be 'row_number'")
+		assert.NotNil(t, funcExpr.Over, "FuncExpr should have OVER clause")
+	})
+
+	t.Run("Window function with PARTITION BY is parsed correctly", func(t *testing.T) {
+		query := "SELECT SUM(amount) OVER (PARTITION BY category ORDER BY date) as running_sum FROM t"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel := stmt.(*sqlparser.Select)
+		aliased := sel.SelectExprs[0].(*sqlparser.AliasedExpr)
+		funcExpr := aliased.Expr.(*sqlparser.FuncExpr)
+
+		assert.Equal(t, "sum", funcExpr.Name.Lowered())
+		require.NotNil(t, funcExpr.Over, "FuncExpr should have OVER clause")
+
+		// Check partition by exists
+		require.NotNil(t, funcExpr.Over.PartitionBy, "OVER clause should have PARTITION BY")
+	})
+
+	t.Run("Window function with frame specification", func(t *testing.T) {
+		query := "SELECT SUM(value) OVER (ORDER BY date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) as cumsum FROM t"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel := stmt.(*sqlparser.Select)
+		aliased := sel.SelectExprs[0].(*sqlparser.AliasedExpr)
+		funcExpr := aliased.Expr.(*sqlparser.FuncExpr)
+
+		assert.NotNil(t, funcExpr.Over, "FuncExpr should have OVER clause")
+	})
+
+	t.Run("Multiple window functions in query", func(t *testing.T) {
+		query := "SELECT ROW_NUMBER() OVER (ORDER BY id) as rn, RANK() OVER (ORDER BY score DESC) as rank FROM t"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel := stmt.(*sqlparser.Select)
+		require.Len(t, sel.SelectExprs, 2, "Should have 2 select expressions")
+
+		// Check first window function
+		aliased1 := sel.SelectExprs[0].(*sqlparser.AliasedExpr)
+		funcExpr1 := aliased1.Expr.(*sqlparser.FuncExpr)
+		assert.NotNil(t, funcExpr1.Over, "First FuncExpr should have OVER clause")
+
+		// Check second window function
+		aliased2 := sel.SelectExprs[1].(*sqlparser.AliasedExpr)
+		funcExpr2 := aliased2.Expr.(*sqlparser.FuncExpr)
+		assert.NotNil(t, funcExpr2.Over, "Second FuncExpr should have OVER clause")
+	})
+
+	t.Run("Regular function without OVER clause", func(t *testing.T) {
+		query := "SELECT UPPER(name) as upper_name FROM t"
+		stmt, err := sqlparser.Parse(query)
+		require.NoError(t, err)
+
+		sel := stmt.(*sqlparser.Select)
+		aliased := sel.SelectExprs[0].(*sqlparser.AliasedExpr)
+		funcExpr := aliased.Expr.(*sqlparser.FuncExpr)
+
+		assert.Equal(t, "upper", funcExpr.Name.Lowered())
+		assert.Nil(t, funcExpr.Over, "UPPER() should not have OVER clause")
+	})
+}