Adds planning error for the with clause

Updates the AST Rewriter for the WITH clause Adds test for the unsupported feature of planning a with clause
partiql · Feb 5, 2025 · 7f0b848 · 7f0b848
1 parent af0995b
commit 7f0b848
Show file tree

Hide file tree

Showing 10 changed files with 76 additions and 14 deletions.
diff --git a/partiql-ast/api/partiql-ast.api b/partiql-ast/api/partiql-ast.api
@@ -433,6 +433,10 @@ public abstract class org/partiql/ast/AstRewriter : org/partiql/ast/AstVisitor {
 	public fun visitUpdateTargetStepField (Lorg/partiql/ast/dml/UpdateTargetStep$Field;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
 	public synthetic fun visitUpsert (Lorg/partiql/ast/dml/Upsert;Ljava/lang/Object;)Ljava/lang/Object;
 	public fun visitUpsert (Lorg/partiql/ast/dml/Upsert;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
+	public synthetic fun visitWith (Lorg/partiql/ast/With;Ljava/lang/Object;)Ljava/lang/Object;
+	public fun visitWith (Lorg/partiql/ast/With;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
+	public synthetic fun visitWithListElement (Lorg/partiql/ast/WithListElement;Ljava/lang/Object;)Ljava/lang/Object;
+	public fun visitWithListElement (Lorg/partiql/ast/WithListElement;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
 }
 
 public abstract class org/partiql/ast/AstVisitor {

diff --git a/partiql-ast/src/main/kotlin/org/partiql/ast/AstRewriter.kt b/partiql-ast/src/main/kotlin/org/partiql/ast/AstRewriter.kt
@@ -346,10 +346,31 @@ public abstract class AstRewriter<C> : AstVisitor<AstNode, C>() {
         val orderBy = node.orderBy?.let { visitOrderBy(it, ctx) as OrderBy? }
         val limit = node.limit?.let { visitExpr(it, ctx) as Expr? }
         val offset = node.offset?.let { visitExpr(it, ctx) as Expr? }
+        val with = node.with?.let { visitWith(it, ctx) as With? }
         return if (body !== node.body || orderBy !== node.orderBy || limit !== node.limit || offset !==
-            node.offset
+            node.offset || with !== node.with
         ) {
-            exprQuerySet(body, orderBy, limit, offset)
+            exprQuerySet(body, orderBy, limit, offset, with)
+        } else {
+            node
+        }
+    }
+
+    override fun visitWith(node: With, ctx: C): AstNode {
+        val elements = _visitList(node.elements, ctx, ::visitWithListElement)
+        return if (elements !== node.elements) {
+            With(elements)
+        } else {
+            node
+        }
+    }
+
+    override fun visitWithListElement(node: WithListElement, ctx: C): AstNode {
+        val columns = node.columnList?.let { _visitList(it, ctx, ::visitIdentifierSimple) }
+        val queryName = visitIdentifierSimple(node.queryName, ctx) as Identifier.Simple
+        val query = visitExprQuerySet(node.asQuery, ctx) as ExprQuerySet
+        return if (columns !== node.columnList || queryName !== node.queryName || query !== node.asQuery) {
+            WithListElement(queryName, query, columns)
         } else {
             node
         }

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/Env.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/Env.kt
@@ -20,6 +20,7 @@ import org.partiql.spi.catalog.Catalogs
 import org.partiql.spi.catalog.Identifier
 import org.partiql.spi.catalog.Name
 import org.partiql.spi.catalog.Session
+import org.partiql.spi.errors.PErrorListener
 import org.partiql.spi.function.Agg
 import org.partiql.spi.function.AggOverload
 import org.partiql.spi.function.FnOverload
@@ -35,7 +36,7 @@ import org.partiql.spi.types.PType
  *
  * @property session
  */
-internal class Env(private val session: Session) {
+internal class Env(private val session: Session, internal val listener: PErrorListener) {
 
     /**
      * Catalogs provider.

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PErrors.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PErrors.kt
@@ -206,6 +206,16 @@ internal object PErrors {
         return PRuntimeException(internalError(cause))
     }
 
+    internal fun featureNotSupported(feature: String, location: SourceLocation? = null): PError {
+        return PError(
+            PError.FEATURE_NOT_SUPPORTED,
+            Severity.ERROR(),
+            PErrorKind.SEMANTIC(),
+            location,
+            mapOf("FEATURE_NAME" to feature),
+        )
+    }
+
     private fun internalError(cause: Throwable): PError = PError(
         PError.INTERNAL_ERROR,
         Severity.ERROR(),

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/SqlPlanner.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/SqlPlanner.kt
@@ -33,7 +33,7 @@ internal class SqlPlanner(
     override fun plan(statement: Statement, session: Session, ctx: Context): Result {
         try {
             // 0. Initialize the planning environment
-            val env = Env(session)
+            val env = Env(session, ctx.errorListener)
 
             // 1. Normalize
             val ast = statement.normalize()

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/NormalizeSelect.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/NormalizeSelect.kt
@@ -42,6 +42,8 @@ import org.partiql.ast.SelectItem
 import org.partiql.ast.SelectList
 import org.partiql.ast.SelectStar
 import org.partiql.ast.SelectValue
+import org.partiql.ast.With
+import org.partiql.ast.WithListElement
 import org.partiql.ast.expr.Expr
 import org.partiql.ast.expr.ExprCase
 import org.partiql.ast.expr.ExprQuerySet
@@ -110,14 +112,20 @@ import org.partiql.planner.internal.util.BinderUtils.toBinder
 internal object NormalizeSelect {
 
     internal fun normalize(node: ExprQuerySet): ExprQuerySet {
+        val with = node.with?.elements?.map { element ->
+            val elementQuery = normalize(element.asQuery)
+            WithListElement(element.queryName, elementQuery, element.columnList)
+        }?.let { With(it) }
+
         return when (val body = node.body) {
             is QueryBody.SFW -> {
                 val sfw = Visitor.visitSFW(body, newCtx())
                 exprQuerySet(
                     body = sfw,
                     orderBy = node.orderBy,
                     limit = node.limit,
-                    offset = node.offset
+                    offset = node.offset,
+                    with = with
                 )
             }
             is QueryBody.SetOp -> {
@@ -132,7 +140,8 @@ internal object NormalizeSelect {
                     ),
                     orderBy = node.orderBy,
                     limit = node.limit,
-                    offset = node.offset
+                    offset = node.offset,
+                    with = with
                 )
             }
             else -> error("Unexpected QueryBody type: $body")

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt
@@ -47,6 +47,7 @@ import org.partiql.ast.expr.Expr
 import org.partiql.ast.expr.ExprCall
 import org.partiql.ast.expr.ExprQuerySet
 import org.partiql.planner.internal.Env
+import org.partiql.planner.internal.PErrors
 import org.partiql.planner.internal.ir.Rel
 import org.partiql.planner.internal.ir.Rex
 import org.partiql.planner.internal.ir.rel
@@ -171,6 +172,9 @@ internal object RelConverter {
             val orderBy = node.orderBy
             val limit = node.limit
             val offset = node.offset
+            if (node.with != null) {
+                env.listener.report(PErrors.featureNotSupported("WITH clause"))
+            }
             when (body) {
                 is QueryBody.SFW -> {
                     var sel = body

diff --git a/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTest.kt b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTest.kt
@@ -125,7 +125,8 @@ class PlanTyperTest {
                             .root(root)
                             .build()
                     )
-                    .build()
+                    .build(),
+                config.errorListener
             )
             return PlanTyperWrapper(PlanTyper(env, config))
         }

diff --git a/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTestsPorted.kt b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTestsPorted.kt
@@ -3740,13 +3740,16 @@ internal class PlanTyperTestsPorted {
     }
 
     @Test
-    @Disabled
+    // @Disabled
     fun developmentTest() {
-        val tc = SuccessTestCase(
-            name = "DEV TEST",
-            query = "CAST('' AS STRING) < CAST('' AS SYMBOL);",
-            expected = PType.bool().toCType()
-        )
+        val tc =
+            ErrorTestCase(
+                name = "WITH not supported yet",
+                query = "WITH x AS (SELECT * FROM <<1, 2, 3>>) SELECT * FROM x",
+                problemHandler = assertProblemExists(
+                    PErrors.featureNotSupported("WITH clause")
+                )
+            )
         runTest(tc)
     }
 
@@ -4057,6 +4060,13 @@ internal class PlanTyperTestsPorted {
         }
 
         private val parameters = listOf(
+            ErrorTestCase(
+                name = "WITH not supported yet",
+                query = "WITH x AS (SELECT * FROM <<1, 2, 3>>) SELECT * FROM x",
+                problemHandler = assertProblemExists(
+                    PErrors.featureNotSupported("WITH clause")
+                )
+            ),
             ErrorTestCase(
                 name = "Pets should not be accessible #1",
                 query = "SELECT * FROM pets",

diff --git a/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/ScopeTest.kt b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/ScopeTest.kt
@@ -11,6 +11,7 @@ import org.partiql.planner.internal.typer.PlanTyper.Companion.toCType
 import org.partiql.spi.catalog.Catalog
 import org.partiql.spi.catalog.Identifier
 import org.partiql.spi.catalog.Session
+import org.partiql.spi.errors.PErrorListener
 import org.partiql.spi.types.PType
 import kotlin.test.assertEquals
 import kotlin.test.fail
@@ -39,7 +40,8 @@ internal class ScopeTest {
                 Session.builder()
                     .catalog("currentCatalog")
                     .catalogs(catalog)
-                    .build()
+                    .build(),
+                PErrorListener.abortOnError()
             ),
             Scope(
                 listOf(