Add stored procedure calls with unnamed args (#345)

examples/Readme.md

@@ -22,6 +22,7 @@ List of Examples:
 * Kotlin:
     * CsvExprValueExample: how to create an `ExprValue` for a custom data format, in this case CSV  
     * CustomFunctionsExample: how to create and register user defined functions (UDF)
+    * CustomProceduresExample: how to create and register stored procedures
     * EvaluationWithBindings: query evaluation with global bindings
     * EvaluationWithLazyBindings: query evaluation with global bindings that are lazily evaluated
     * ParserErrorExample: inspecting errors thrown by the `Parser`

examples/src/kotlin/org/partiql/examples/CustomProceduresExample.kt

@@ -0,0 +1,130 @@
+package org.partiql.examples
+
+import com.amazon.ion.IonDecimal
+import com.amazon.ion.IonStruct
+import com.amazon.ion.system.IonSystemBuilder
+import org.partiql.examples.util.Example
+import org.partiql.lang.CompilerPipeline
+import org.partiql.lang.errors.ErrorCode
+import org.partiql.lang.errors.Property
+import org.partiql.lang.errors.PropertyValueMap
+import org.partiql.lang.eval.BindingCase
+import org.partiql.lang.eval.BindingName
+import org.partiql.lang.eval.Bindings
+import org.partiql.lang.eval.EvaluationException
+import org.partiql.lang.eval.EvaluationSession
+import org.partiql.lang.eval.ExprValue
+import org.partiql.lang.eval.ExprValueFactory
+import org.partiql.lang.eval.ExprValueType
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedureSignature
+import org.partiql.lang.eval.stringValue
+import java.io.PrintStream
+import java.math.BigDecimal
+import java.math.RoundingMode
+
+private val ion = IonSystemBuilder.standard().build()
+
+/**
+ * A simple custom stored procedure that calculates the moon weight for each crewmate of the given crew, storing the
+ * moon weight in the [EvaluationSession] global bindings. This procedure also returns the number of crewmates we
+ * calculated the moon weight for, returning -1 if no crew is found.
+ *
+ * This example demonstrates how to create a custom stored procedure, check argument types, and modify the
+ * [EvaluationSession].
+ */
+class CalculateCrewMoonWeight(private val valueFactory: ExprValueFactory): StoredProcedure {
+    private val MOON_GRAVITATIONAL_CONSTANT = BigDecimal(1.622 / 9.81)
+
+    // [StoredProcedureSignature] takes two arguments:
+    //   1. the name of the stored procedure
+    //   2. the arity of this stored procedure. Checks to arity are taken care of by the evaluator. However, we must
+    //      still check that the passed arguments are of the right type in our implementation of the procedure.
+    override val signature = StoredProcedureSignature(name = "calculate_crew_moon_weight", arity = 1)
+
+    // `call` is where you define the logic of the stored procedure given an [EvaluationSession] and a list of
+    // arguments
+    override fun call(session: EvaluationSession, args: List<ExprValue>): ExprValue {
+        // We first check that the first argument is a string
+        val crewName = args.first()
+        // In the future the evaluator will also verify function argument types, but for now we must verify their type
+        // manually
+        if (crewName.type != ExprValueType.STRING) {
+            val errorContext = PropertyValueMap().also {
+                it[Property.EXPECTED_ARGUMENT_TYPES] = "STRING"
+                it[Property.ACTUAL_ARGUMENT_TYPES] = crewName.type.name
+                it[Property.FUNCTION_NAME] = signature.name
+            }
+            throw EvaluationException("First argument to ${signature.name} was not a string",
+                ErrorCode.EVALUATOR_INCORRECT_TYPE_OF_ARGUMENTS_TO_PROCEDURE_CALL,
+                errorContext,
+                internal = false)
+        }
+
+        // Next we check if the given `crewName` is in the [EvaluationSession]'s global bindings. If not, we return 0.
+        val sessionGlobals = session.globals
+        val crewBindings = sessionGlobals[BindingName(crewName.stringValue(), BindingCase.INSENSITIVE)]
+            ?: return valueFactory.newInt(-1)
+
+        // Now that we've confirmed the given `crewName` is in the session's global bindings, we calculate and store
+        // the moon weight for each crewmate in the crew.
+        // In addition, we keep a running a tally of how many crewmates we do this for.
+        var numCalculated = 0
+        for (crewmateBinding in crewBindings) {
+            val crewmate = crewmateBinding.ionValue as IonStruct
+            val mass = crewmate["mass"] as IonDecimal
+            val moonWeight = (mass.decimalValue() * MOON_GRAVITATIONAL_CONSTANT).setScale(1, RoundingMode.HALF_UP)
+            crewmate.add("moonWeight", ion.newDecimal(moonWeight))
+
+            numCalculated++
+        }
+        return valueFactory.newInt(numCalculated)
+    }
+}
+
+/**
+ * Demonstrates the use of custom stored procedure [CalculateCrewMoonWeight] in PartiQL queries.
+ */
+class CustomProceduresExample(out: PrintStream) : Example(out) {
+    override fun run() {
+        /**
+         * To make custom stored procedures available to the PartiQL query being executed, they must be passed to
+         * [CompilerPipeline.Builder.addProcedure].
+         */
+        val pipeline = CompilerPipeline.build(ion) {
+            addProcedure(CalculateCrewMoonWeight(valueFactory))
+        }
+
+        // Here, we initialize the crews to be stored in our global session bindings
+        val initialCrews = Bindings.ofMap(
+            mapOf(
+                "crew1" to pipeline.valueFactory.newFromIonValue(
+                    ion.singleValue("""[ { name: "Neil",    mass: 80.5 }, 
+                                         { name: "Buzz",    mass: 72.3 },
+                                         { name: "Michael", mass: 89.9 } ]""")),
+                "crew2" to pipeline.valueFactory.newFromIonValue(
+                    ion.singleValue("""[ { name: "James", mass: 77.1 }, 
+                                         { name: "Spock", mass: 81.6 } ]"""))
+            )
+        )
+        val session = EvaluationSession.build { globals(initialCrews) }
+
+        val crew1BindingName = BindingName("crew1", BindingCase.INSENSITIVE)
+        val crew2BindingName = BindingName("crew2", BindingCase.INSENSITIVE)
+
+        out.println("Initial global session bindings:")
+        print("Crew 1:", "${session.globals[crew1BindingName]}")
+        print("Crew 2:", "${session.globals[crew2BindingName]}")
+
+        // We call our custom stored procedure using PartiQL's `EXEC` clause. Here we call our stored procedure
+        // 'calculate_crew_moon_weight' with the arg 'crew1', which outputs the number of crewmates we've calculated
+        // the moon weight for
+        val procedureCall = "EXEC calculate_crew_moon_weight 'crew1'"
+        val procedureCallOutput = pipeline.compile(procedureCall).eval(session)
+        print("Number of calculated moon weights:", "$procedureCallOutput")
+
+        out.println("Updated global session bindings:")
+        print("Crew 1:", "${session.globals[crew1BindingName]}")
+        print("Crew 2:", "${session.globals[crew2BindingName]}")
+    }
+}

examples/src/kotlin/org/partiql/examples/util/Main.kt

@@ -15,6 +15,7 @@ private val examples = mapOf(
         // Kotlin Examples
         CsvExprValueExample::class.java.simpleName to CsvExprValueExample(System.out),
         CustomFunctionsExample::class.java.simpleName to CustomFunctionsExample(System.out),
+        CustomProceduresExample::class.java.simpleName to CustomProceduresExample(System.out),
         EvaluationWithBindings::class.java.simpleName to EvaluationWithBindings(System.out),
         EvaluationWithLazyBindings::class.java.simpleName to EvaluationWithLazyBindings(System.out),
         ParserErrorExample::class.java.simpleName to ParserErrorExample(System.out),

examples/test/org/partiql/examples/CustomProceduresExampleTest.kt

@@ -0,0 +1,24 @@
+package org.partiql.examples
+
+import org.partiql.examples.util.Example
+import java.io.PrintStream
+
+class CustomProceduresExampleTest : BaseExampleTest() {
+    override fun example(out: PrintStream): Example = CustomProceduresExample(out)
+
+    override val expected = """
+        |Initial global session bindings:
+        |Crew 1:
+        |    [{'name': 'Neil', 'mass': 80.5}, {'name': 'Buzz', 'mass': 72.3}, {'name': 'Michael', 'mass': 89.9}]
+        |Crew 2:
+        |    [{'name': 'James', 'mass': 77.1}, {'name': 'Spock', 'mass': 81.6}]
+        |Number of calculated moon weights:
+        |    3
+        |Updated global session bindings:
+        |Crew 1:
+        |    [{'name': 'Neil', 'mass': 80.5, 'moonWeight': 13.3}, {'name': 'Buzz', 'mass': 72.3, 'moonWeight': 12.0}, {'name': 'Michael', 'mass': 89.9, 'moonWeight': 14.9}]
+        |Crew 2:
+        |    [{'name': 'James', 'mass': 77.1}, {'name': 'Spock', 'mass': 81.6}]
+        |
+    """.trimMargin()
+}

lang/resources/org/partiql/type-domains/partiql.ion

@@ -26,7 +26,12 @@
                 (where where::(? expr)))
 
             // Data definition operations also cannot be composed with other `expr` nodes.
-            (ddl op::ddl_op))
+            (ddl op::ddl_op)
+
+            // Stored procedure calls are only allowed at the top level of a query and cannot be used as an expression
+            // Currently supports stored procedure calls with the unnamed argument syntax:
+            //     EXEC <symbol> [<expr>.*]
+            (exec procedure_name::symbol args::(* expr 0)))
 
         // The expressions that can result in values.
         (sum expr

lang/src/org/partiql/lang/CompilerPipeline.kt

@@ -18,6 +18,7 @@ import com.amazon.ion.*
 import org.partiql.lang.ast.*
 import org.partiql.lang.eval.*
 import org.partiql.lang.eval.builtins.*
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
 import org.partiql.lang.syntax.*
 
 /**
@@ -35,7 +36,13 @@ data class StepContext(
      * Includes built-in functions as well as custom functions added while the [CompilerPipeline]
      * was being built.
      */
-    val functions: @JvmSuppressWildcards Map<String, ExprFunction>
+    val functions: @JvmSuppressWildcards Map<String, ExprFunction>,
+
+    /**
+     * Returns a list of all stored procedures which are available for execution.
+     * Only includes the custom stored procedures added while the [CompilerPipeline] was being built.
+     */
+    val procedures: @JvmSuppressWildcards Map<String, StoredProcedure>
 )
 
 /**
@@ -65,6 +72,12 @@ interface CompilerPipeline  {
      */
     val functions: @JvmSuppressWildcards Map<String, ExprFunction>
 
+    /**
+     * Returns a list of all stored procedures which are available for execution.
+     * Only includes the custom stored procedures added while the [CompilerPipeline] was being built.
+     */
+    val procedures: @JvmSuppressWildcards Map<String, StoredProcedure>
+
     /** Compiles the specified PartiQL query using the configured parser. */
     fun compile(query: String): Expression
 
@@ -106,6 +119,7 @@ interface CompilerPipeline  {
         private var parser: Parser? = null
         private var compileOptions: CompileOptions? = null
         private val customFunctions: MutableMap<String, ExprFunction> = HashMap()
+        private val customProcedures: MutableMap<String, StoredProcedure> = HashMap()
         private val preProcessingSteps: MutableList<ProcessingStep> = ArrayList()
 
         /**
@@ -137,6 +151,13 @@ interface CompilerPipeline  {
          */
         fun addFunction(function: ExprFunction): Builder = this.apply { customFunctions[function.name] = function }
 
+        /**
+         * Add a custom stored procedure which will be callable by the compiled queries.
+         *
+         * Stored procedures added here will replace any built-in procedure with the same name.
+         */
+        fun addProcedure(procedure: StoredProcedure): Builder = this.apply { customProcedures[procedure.signature.name] = procedure }
+
         /** Adds a preprocessing step to be executed after parsing but before compilation. */
         fun addPreprocessingStep(step: ProcessingStep): Builder = this.apply { preProcessingSteps.add(step) }
 
@@ -153,6 +174,7 @@ interface CompilerPipeline  {
                 parser ?: SqlParser(valueFactory.ion),
                 compileOptions ?: CompileOptions.standard(),
                 allFunctions,
+                customProcedures,
                 preProcessingSteps)
         }
     }
@@ -163,17 +185,18 @@ private class CompilerPipelineImpl(
     private val parser: Parser,
     override val compileOptions: CompileOptions,
     override val functions: Map<String, ExprFunction>,
+    override val procedures: Map<String, StoredProcedure>,
     private val preProcessingSteps: List<ProcessingStep>
 ) : CompilerPipeline {
 
-    private val compiler = EvaluatingCompiler(valueFactory, functions, compileOptions)
+    private val compiler = EvaluatingCompiler(valueFactory, functions, procedures, compileOptions)
 
     override fun compile(query: String): Expression {
         return compile(parser.parseExprNode(query))
     }
 
     override fun compile(query: ExprNode): Expression {
-        val context = StepContext(valueFactory, compileOptions, functions)
+        val context = StepContext(valueFactory, compileOptions, functions, procedures)
 
         val preProcessedQuery = preProcessingSteps.fold(query) { currentExprNode, step ->
             step(currentExprNode, context)

lang/src/org/partiql/lang/ast/AstSerialization.kt

@@ -94,6 +94,7 @@ private class AstSerializerImpl(val astVersion: AstVersion, val ion: IonSystem):
                     is DropTable         -> case { writeDropTable(expr) }
                     is DropIndex         -> case { writeDropIndex(expr) }
                     is Parameter         -> case { writeParameter(expr)}
+                    is Exec              -> throw UnsupportedOperationException("EXEC clause not supported by the V0 AST")
                 }.toUnit()
             }
         }

lang/src/org/partiql/lang/ast/ExprNodeToStatement.kt

@@ -2,6 +2,7 @@ package org.partiql.lang.ast
 
 import com.amazon.ionelement.api.toIonElement
 import org.partiql.lang.domains.PartiqlAst
+import org.partiql.pig.runtime.SymbolPrimitive
 import org.partiql.pig.runtime.asPrimitive
 
 /** Converts an [ExprNode] to a [PartiqlAst.statement]. */
@@ -16,12 +17,16 @@ fun ExprNode.toAstStatement(): PartiqlAst.Statement {
 
         is CreateTable, is CreateIndex, is DropTable, is DropIndex -> toAstDdl()
 
+        is Exec -> toAstExec()
     }
 }
 
 private fun PartiQlMetaContainer.toElectrolyteMetaContainer(): ElectrolyteMetaContainer =
     com.amazon.ionelement.api.metaContainerOf(map { it.tag to it })
 
+private fun SymbolicName.toSymbolPrimitive() : SymbolPrimitive =
+    SymbolPrimitive(this.name, this.metas.toElectrolyteMetaContainer())
+
 private fun ExprNode.toAstDdl(): PartiqlAst.Statement {
     val thiz = this
     val metas = metas.toElectrolyteMetaContainer()
@@ -30,7 +35,7 @@ private fun ExprNode.toAstDdl(): PartiqlAst.Statement {
         when(thiz) {
             is Literal, is LiteralMissing, is VariableReference, is Parameter, is NAry, is CallAgg, is Typed,
             is Path, is SimpleCase, is SearchedCase, is Select, is Struct, is Seq,
-            is DataManipulation -> error("Can't convert ${thiz.javaClass} to PartiqlAst.ddl")
+            is DataManipulation, is Exec -> error("Can't convert ${thiz.javaClass} to PartiqlAst.ddl")
 
             is CreateTable -> ddl(createTable(thiz.tableName), metas)
             is CreateIndex -> ddl(createIndex(identifier(thiz.tableName, caseSensitive()), thiz.keys.map { it.toAstExpr() }), metas)
@@ -48,6 +53,18 @@ private fun ExprNode.toAstDdl(): PartiqlAst.Statement {
     }
 }
 
+private fun ExprNode.toAstExec() : PartiqlAst.Statement {
+    val node = this
+    val metas = metas.toElectrolyteMetaContainer()
+
+    return PartiqlAst.build {
+        when (node) {
+            is Exec -> exec_(node.procedureName.toSymbolPrimitive(), node.args.map { it.toAstExpr() }, metas)
+            else -> error("Can't convert ${node.javaClass} to PartiqlAst.Statement.Exec")
+        }
+    }
+}
+
 fun ExprNode.toAstExpr(): PartiqlAst.Expr {
     val node = this
     val metas = this.metas.toElectrolyteMetaContainer()
@@ -147,8 +164,8 @@ fun ExprNode.toAstExpr(): PartiqlAst.Expr {
                     SeqType.BAG -> bag(node.values.map { it.toAstExpr() })
                 }
 
-            // These are handled by `toAstDml()`
-            is DataManipulation, is CreateTable, is CreateIndex, is DropTable, is DropIndex ->
+            // These are handled by `toAstDml()`, `toAstDdl()`, and `toAstExec()`
+            is DataManipulation, is CreateTable, is CreateIndex, is DropTable, is DropIndex, is Exec ->
                 error("Can't transform ${node.javaClass} to a PartiqlAst.expr }")
         }
     }

lang/src/org/partiql/lang/ast/StatementToExprNode.kt

@@ -2,6 +2,7 @@ package org.partiql.lang.ast
 
 import com.amazon.ion.IonSystem
 import com.amazon.ionelement.api.toIonValue
+import org.partiql.lang.domains.PartiqlAst
 import org.partiql.lang.domains.PartiqlAst.CaseSensitivity
 import org.partiql.lang.domains.PartiqlAst.DdlOp
 import org.partiql.lang.domains.PartiqlAst.DmlOp
@@ -33,6 +34,7 @@ private class StatementTransformer(val ion: IonSystem) {
             is Statement.Query -> stmt.toExprNode()
             is Statement.Dml -> stmt.toExprNode()
             is Statement.Ddl -> stmt.toExprNode()
+            is Statement.Exec -> stmt.toExprNode()
         }
 
     private fun ElectrolyteMetaContainer.toPartiQlMetaContainer(): PartiQlMetaContainer {
@@ -344,4 +346,8 @@ private class StatementTransformer(val ion: IonSystem) {
                     metas = metas)
         }
     }
+
+    private fun Statement.Exec.toExprNode(): ExprNode {
+        return Exec(procedureName.toSymbolicName(), this.args.toExprNodeList(), metas.toPartiQlMetaContainer())
+    }
 }

lang/src/org/partiql/lang/ast/ast.kt

@@ -103,6 +103,9 @@ sealed class ExprNode : AstNode(), HasMetas {
             is Parameter       -> {
                 copy(metas = metas)
             }
+            is Exec            -> {
+                copy(metas = metas)
+            }
         }
     }
 }
@@ -210,6 +213,19 @@ data class Typed(
     override val children: List<AstNode> = listOf(expr, type)
 }
 
+//********************************
+// Stored procedure clauses
+//********************************
+
+/** Represents a call to a stored procedure, i.e. `EXEC stored_procedure [<expr>.*]` */
+data class Exec(
+    val procedureName: SymbolicName,
+    val args: List<ExprNode>,
+    override val metas: MetaContainer
+) : ExprNode() {
+    override val children: List<AstNode> = args
+}
+
 //********************************
 // Path expressions
 //********************************