partiql · johnedquinn · May 4, 2023 · May 1, 2023 · May 2, 2023 · May 4, 2023
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,15 +28,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Adds support for SQL's CURRENT_USER in the AST, EvaluatingCompiler, experimental planner implementation, and Schema Inferencer.
   - Adds the AST node `session_attribute`.
   - Adds the function `EvaluationSession.Builder::user()` to add the CURRENT_USER to the EvaluationSession
+- Adds support for parsing and planning of `INSERT INTO .. AS <alias> ... ON CONFLICT DO [UPDATE|REPLACE] EXCLUDED WHERE <expr>`
+- Adds the `statement.dml` and `dml_operation` node to the experimental PartiQL Physical Plan.
 
 ### Changed
 
+- **Breaking**: Adds a new property `as_alias` to the `insert` AST node.
+- **Breaking**: Adds new property `condition` to the AST nodes of `do_replace` and `do_update`
+- **Breaking**: Adds `target_alias` property to the `dml_insert`, `dml_replace`, and `dml_update` nodes within the
+  Logical and Logical Resolved plans
+- **Breaking**: Adds `condition` property to the `dml_replace` and `dml_update` nodes within the
+  Logical and Logical Resolved plans
+
 ### Deprecated
 
 ### Fixed
 
+- Parsing INSERT statements with aliases no longer loses the original table name. Closes #1043.
+
 ### Removed
 
+- **Breaking**: Removes node `statement.dml_query` from the experimental PartiQL Physical Plan. Please see the added
+  `statement.dml` and `dml_operation` nodes.
+
 ### Security
 
 ## [0.9.3] - 2023-04-12

diff --git a/partiql-lang/src/main/antlr/PartiQL.g4 b/partiql-lang/src/main/antlr/PartiQL.g4
@@ -161,6 +161,8 @@ removeCommand
 insertCommandReturning
     : INSERT INTO pathSimple VALUE value=expr ( AT pos=expr )? onConflictClause? returningClause?;
 
+// TODO: Update grammar to disallow the mixing and matching of `insert`, `insertLegacy`, `onConflict`, and `onConflictLegacy`
+//  See https://github.com/partiql/partiql-lang-kotlin/issues/1063 for more details.
 insertCommand
     : INSERT INTO pathSimple VALUE value=expr ( AT pos=expr )? onConflictClause?  # InsertLegacy
     // See the Grammar at https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md#2-proposed-grammar-and-semantics
@@ -197,7 +199,7 @@ conflictAction
    [ WHERE <condition> ]
 */
 doReplace
-    : EXCLUDED;
+    : EXCLUDED ( WHERE condition=expr )?;
     // :TODO add the rest of the grammar
 
 /*
@@ -207,7 +209,7 @@ doReplace
    [ WHERE <condition> ]
 */
 doUpdate
-    : EXCLUDED;
+    : EXCLUDED ( WHERE condition=expr )?;
     // :TODO add the rest of the grammar
 
 updateClause

diff --git a/partiql-lang/src/main/kotlin/org/partiql/lang/compiler/PartiQLCompilerDefault.kt b/partiql-lang/src/main/kotlin/org/partiql/lang/compiler/PartiQLCompilerDefault.kt
@@ -20,12 +20,8 @@ import org.partiql.lang.domains.PartiqlLogical
 import org.partiql.lang.domains.PartiqlLogicalResolved
 import org.partiql.lang.domains.PartiqlPhysical
 import org.partiql.lang.errors.PartiQLException
-import org.partiql.lang.eval.BindingCase
-import org.partiql.lang.eval.BindingName
-import org.partiql.lang.eval.Bindings
 import org.partiql.lang.eval.ExprFunction
 import org.partiql.lang.eval.ExprValue
-import org.partiql.lang.eval.Expression
 import org.partiql.lang.eval.PartiQLResult
 import org.partiql.lang.eval.PartiQLStatement
 import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
@@ -35,10 +31,6 @@ import org.partiql.lang.eval.physical.PhysicalPlanCompilerImpl
 import org.partiql.lang.eval.physical.PhysicalPlanThunk
 import org.partiql.lang.eval.physical.operators.RelationalOperatorFactory
 import org.partiql.lang.eval.physical.operators.RelationalOperatorFactoryKey
-import org.partiql.lang.planner.DML_COMMAND_FIELD_ACTION
-import org.partiql.lang.planner.DML_COMMAND_FIELD_ROWS
-import org.partiql.lang.planner.DML_COMMAND_FIELD_TARGET_UNIQUE_ID
-import org.partiql.lang.planner.DmlAction
 import org.partiql.lang.planner.EvaluatorOptions
 import org.partiql.lang.planner.PartiQLPlanner
 import org.partiql.lang.types.TypedOpParameter
@@ -71,18 +63,20 @@ internal class PartiQLCompilerDefault(
     }
 
     override fun compile(statement: PartiqlPhysical.Plan): PartiQLStatement {
-        val expression = exprConverter.compile(statement)
         return when (statement.stmt) {
-            is PartiqlPhysical.Statement.DmlQuery -> PartiQLStatement { expression.eval(it).toDML() }
+            is PartiqlPhysical.Statement.Dml -> compileDml(statement.stmt, statement.locals.size)
             is PartiqlPhysical.Statement.Exec,
-            is PartiqlPhysical.Statement.Query -> PartiQLStatement { expression.eval(it).toValue() }
+            is PartiqlPhysical.Statement.Query -> {
+                val expression = exprConverter.compile(statement)
+                PartiQLStatement { expression.eval(it).toValue() }
+            }
             is PartiqlPhysical.Statement.Explain -> throw PartiQLException("Unable to compile EXPLAIN without details.")
         }
     }
 
     override fun compile(statement: PartiqlPhysical.Plan, details: PartiQLPlanner.PlanningDetails): PartiQLStatement {
         return when (statement.stmt) {
-            is PartiqlPhysical.Statement.DmlQuery,
+            is PartiqlPhysical.Statement.Dml -> compileDml(statement.stmt, statement.locals.size)
             is PartiqlPhysical.Statement.Exec,
             is PartiqlPhysical.Statement.Query -> compile(statement)
             is PartiqlPhysical.Statement.Explain -> PartiQLStatement { compileExplain(statement.stmt, details) }
@@ -100,6 +94,18 @@ internal class PartiQLCompilerDefault(
         PHYSICAL_TRANSFORMED
     }
 
+    private fun compileDml(dml: PartiqlPhysical.Statement.Dml, localsSize: Int): PartiQLStatement {
+        val rows = exprConverter.compile(dml.rows, localsSize)
+        return PartiQLStatement { session ->
+            when (dml.operation) {
+                is PartiqlPhysical.DmlOperation.DmlReplace -> PartiQLResult.Replace(dml.uniqueId.text, rows.eval(session))
+                is PartiqlPhysical.DmlOperation.DmlInsert -> PartiQLResult.Insert(dml.uniqueId.text, rows.eval(session))
+                is PartiqlPhysical.DmlOperation.DmlDelete -> PartiQLResult.Delete(dml.uniqueId.text, rows.eval(session))
+                is PartiqlPhysical.DmlOperation.DmlUpdate -> TODO("DML Update compilation not supported yet.")
+            }
+        }
+    }
+
     private fun compileExplain(statement: PartiqlPhysical.Statement.Explain, details: PartiQLPlanner.PlanningDetails): PartiQLResult.Explain.Domain {
         return when (statement.target) {
             is PartiqlPhysical.ExplainTarget.Domain -> compileExplainDomain(statement.target, details)
@@ -152,44 +158,5 @@ internal class PartiQLCompilerDefault(
         }
     }
 
-    /**
-     * The physical expr converter is EvaluatingCompiler with s/Ast/Physical and `bindingsToValues -> Thunk`
-     *   so it returns the [Expression] rather than a [PartiQLStatement]. This method parses a DML Command from the result.
-     *
-     * {
-     *     'action': <action>,
-     *     'target_unique_id': <unique_id>
-     *     'rows': <rows>
-     * }
-     *
-     * Later refactors will rework the Compiler to use [PartiQLStatement], but this is an acceptable workaround for now.
-     */
-    private fun ExprValue.toDML(): PartiQLResult {
-        val action = bindings string DML_COMMAND_FIELD_ACTION
-        val target = bindings string DML_COMMAND_FIELD_TARGET_UNIQUE_ID
-        val rows = bindings seq DML_COMMAND_FIELD_ROWS
-        return when (DmlAction.safeValueOf(action)) {
-            DmlAction.INSERT -> PartiQLResult.Insert(target, rows)
-            DmlAction.DELETE -> PartiQLResult.Delete(target, rows)
-            DmlAction.REPLACE -> PartiQLResult.Replace(target, rows)
-            null -> error("Unknown DML Action `$action`")
-        }
-    }
-
     private fun ExprValue.toValue(): PartiQLResult = PartiQLResult.Value(this)
-
-    private infix fun Bindings<ExprValue>.string(field: String): String {
-        return this[BindingName(field, BindingCase.SENSITIVE)]?.scalar?.stringValue() ?: missing(field)
-    }
-
-    private infix fun Bindings<ExprValue>.seq(field: String): Iterable<ExprValue> {
-        val v = this[BindingName(field, BindingCase.SENSITIVE)] ?: missing(field)
-        if (!v.type.isSequence) {
-            error("DML command struct '$DML_COMMAND_FIELD_ROWS' field must be a bag or list")
-        }
-        return v.asIterable()
-    }
-
-    private fun missing(field: String): Nothing =
-        error("Field `$field` missing from DML command struct or has incorrect Ion type")
 }
diff --git a/partiql-lang/src/main/kotlin/org/partiql/lang/eval/physical/PhysicalPlanCompilerImpl.kt b/partiql-lang/src/main/kotlin/org/partiql/lang/eval/physical/PhysicalPlanCompilerImpl.kt
@@ -173,6 +173,27 @@ internal class PhysicalPlanCompilerImpl(
         }
     }
 
+    /**
+     * Compiles a [PartiqlPhysical.Expr] tree to an [Expression].
+     *
+     * Checks [Thread.interrupted] before every expression and sub-expression is compiled
+     * and throws [InterruptedException] if [Thread.interrupted] it has been set in the
+     * hope that long-running compilations may be aborted by the caller.
+     */
+    internal fun compile(expr: PartiqlPhysical.Expr, localsSize: Int): Expression {
+        val thunk = compileAstExpr(expr)
+
+        return object : Expression {
+            override fun eval(session: EvaluationSession): ExprValue {
+                val env = EvaluatorState(
+                    session = session,
+                    registers = Array(localsSize) { ExprValue.missingValue }
+                )
+                return thunk(env)
+            }
+        }
+    }
+
     override fun convert(expr: PartiqlPhysical.Expr): PhysicalPlanThunk = this.compileAstExpr(expr)
 
     /**
@@ -183,8 +204,8 @@ internal class PhysicalPlanCompilerImpl(
     private fun compileAstStatement(ast: PartiqlPhysical.Statement): PhysicalPlanThunk {
         return when (ast) {
             is PartiqlPhysical.Statement.Query -> compileAstExpr(ast.expr)
-            is PartiqlPhysical.Statement.DmlQuery -> compileAstExpr(ast.expr)
             is PartiqlPhysical.Statement.Exec -> compileExec(ast)
+            is PartiqlPhysical.Statement.Dml,
             is PartiqlPhysical.Statement.Explain -> {
                 val value = ExprValue.newBoolean(true)
                 thunkFactory.thunkEnv(emptyMetaContainer()) { value }

diff --git a/partiql-lang/src/main/kotlin/org/partiql/lang/planner/QueryPlan.kt b/partiql-lang/src/main/kotlin/org/partiql/lang/planner/QueryPlan.kt
@@ -5,7 +5,6 @@ import org.partiql.lang.eval.BindingName
 import org.partiql.lang.eval.Bindings
 import org.partiql.lang.eval.EvaluationSession
 import org.partiql.lang.eval.ExprValue
-import org.partiql.lang.eval.ExprValueType
 
 /** A query plan that has been compiled and is ready to be evaluated. */
 fun interface QueryPlan {
@@ -29,7 +28,7 @@ sealed class QueryResult {
      *
      * The primary benefit of this class is that it ensures that the [rows] property is evaluated lazily.  It also
      * provides a cleaner API that is easier to work with for PartiQL embedders.  Without this, the user would have to
-     * consume the [ExprValue] directly and use code similar to that in [toDmlCommand] or convert it to Ion.  Neither
+     * consume the [ExprValue] directly and convert it to Ion.  Neither
      * option is particularly developer friendly, efficient or maintainable.
      *
      * This is currently only factored to support `INSERT INTO` and `DELETE FROM` as `UPDATE` and `FROM ... UPDATE` is
@@ -67,54 +66,5 @@ enum class DmlAction {
     }
 }
 
-internal const val DML_COMMAND_FIELD_ACTION = "action"
-internal const val DML_COMMAND_FIELD_TARGET_UNIQUE_ID = "target_unique_id"
-internal const val DML_COMMAND_FIELD_ROWS = "rows"
-
 private operator fun Bindings<ExprValue>.get(fieldName: String): ExprValue? =
     this[BindingName(fieldName, BindingCase.SENSITIVE)]
-
-private fun errMissing(fieldName: String): Nothing =
-    error("'$fieldName' missing from DML command struct or has incorrect Ion type")
-
-/**
- * Converts an [ExprValue] which is the result of a DML query to an instance of [DmlCommand].
- *
- * Format of a such an [ExprValue]:
- *
- * ```
- * {
- *     'action': <action>,
- *     'target_unique_id': <unique_id>
- *     'rows': <rows>
- * }
- * ```
- *
- * Where:
- *  - `<action>` is either `insert` or `delete`
- *  - `<target_unique_id>` is a string or symbol containing the unique identifier of the table to be effected
- *  by the DML statement.
- *  - `<rows>` is a bag or list containing the rows (structs) effected by the DML statement.
- *      - When `<action>` is `insert` this is the rows to be inserted.
- *      - When `<action>` is `delete` this is the rows to be deleted.  Non-primary key fields may be elided, but the
- *      default behavior is to include all fields because PartiQL does not yet know about primary keys.
- */
-internal fun ExprValue.toDmlCommand(): QueryResult.DmlCommand {
-    require(this.type == ExprValueType.STRUCT) { "'row' must be a struct" }
-
-    val actionString = this.bindings[DML_COMMAND_FIELD_ACTION]?.scalar?.stringValue()?.toUpperCase()
-        ?: errMissing(DML_COMMAND_FIELD_ACTION)
-
-    val dmlAction = DmlAction.safeValueOf(actionString)
-        ?: error("Unknown DmlAction in DML command struct: '$actionString'")
-
-    val targetUniqueId = this.bindings[DML_COMMAND_FIELD_TARGET_UNIQUE_ID]?.scalar?.stringValue()
-        ?: errMissing(DML_COMMAND_FIELD_TARGET_UNIQUE_ID)
-
-    val rows = this.bindings[DML_COMMAND_FIELD_ROWS] ?: errMissing(DML_COMMAND_FIELD_ROWS)
-    if (!rows.type.isSequence) {
-        error("DML command struct '$DML_COMMAND_FIELD_ROWS' field must be a bag or list")
-    }
-
-    return QueryResult.DmlCommand(dmlAction, targetUniqueId, rows)
-}
diff --git a/...-lang/src/main/kotlin/org/partiql/lang/planner/transforms/AstToLogicalVisitorTransform.kt b/...-lang/src/main/kotlin/org/partiql/lang/planner/transforms/AstToLogicalVisitorTransform.kt
@@ -363,46 +363,34 @@ internal class AstToLogicalVisitorTransform(
                     }
                 }
 
-                when (val conflictAction = dmlOp.conflictAction) {
-                    null -> {
-                        PartiqlLogical.build {
-                            dml(
-                                target = dmlOp.target.toDmlTargetId(),
-                                operation = dmlInsert(),
-                                rows = transformExpr(dmlOp.values),
-                                metas = node.metas
-                            )
-                        }
-                    }
-                    is PartiqlAst.ConflictAction.DoReplace -> {
-                        when (conflictAction.value) {
-                            PartiqlAst.OnConflictValue.Excluded() -> PartiqlLogical.build {
-                                dml(
-                                    target = dmlOp.target.toDmlTargetId(),
-                                    operation = dmlReplace(),
-                                    rows = transformExpr(dmlOp.values),
-                                    metas = node.metas
-                                )
-                            } else -> TODO("Only `DO REPLACE EXCLUDED` is supported in logical plan at the moment.")
-                        }
+                val target = dmlOp.target.toDmlTargetId()
+                val alias = dmlOp.asAlias?.let {
+                    PartiqlLogical.VarDecl(it)
+                } ?: PartiqlLogical.VarDecl(target.name)
+
+                val operation = when (val conflictAction = dmlOp.conflictAction) {
+                    null -> PartiqlLogical.DmlOperation.DmlInsert(targetAlias = alias)
+                    is PartiqlAst.ConflictAction.DoReplace -> when (conflictAction.value) {
+                        is PartiqlAst.OnConflictValue.Excluded -> PartiqlLogical.DmlOperation.DmlReplace(
+                            targetAlias = alias,
+                            condition = conflictAction.condition?.let { transformExpr(conflictAction.condition) }
+                        )
                     }
-                    is PartiqlAst.ConflictAction.DoUpdate -> {
-                        when (conflictAction.value) {
-                            PartiqlAst.OnConflictValue.Excluded() -> PartiqlLogical.build {
-                                dml(
-                                    target = dmlOp.target.toDmlTargetId(),
-                                    operation = dmlUpdate(),
-                                    rows = transformExpr(dmlOp.values),
-                                    metas = node.metas
-                                )
-                            }
-                            else -> TODO("Only `DO UPDATE EXCLUDED` is supported in logical plan at the moment.")
-                        }
+                    is PartiqlAst.ConflictAction.DoUpdate -> when (conflictAction.value) {
+                        is PartiqlAst.OnConflictValue.Excluded -> PartiqlLogical.DmlOperation.DmlUpdate(
+                            targetAlias = alias,
+                            condition = conflictAction.condition?.let { transformExpr(conflictAction.condition) }
+                        )
                     }
-                    is PartiqlAst.ConflictAction.DoNothing -> TODO(
-                        "`ON CONFLICT DO NOTHING` is not supported in logical plan yet."
-                    )
+                    is PartiqlAst.ConflictAction.DoNothing -> TODO("`ON CONFLICT DO NOTHING` is not supported in logical plan yet.")
                 }
+
+                PartiqlLogical.Statement.Dml(
+                    target = target,
+                    operation = operation,
+                    rows = transformExpr(dmlOp.values),
+                    metas = node.metas
+                )
             }
             // INSERT single row with VALUE is disallowed. (This variation of INSERT might be removed in a future
             // release of PartiQL.)