Skip to content

Commit 260d5ce

Browse files
alancai98alancai98
authored
Add stored procedure calls with unnamed args (#345)
1 parent 2ca7702 commit 260d5ce

File tree

21 files changed

+818
-12
lines changed

21 files changed

+818
-12
lines changed

‎examples/Readme.md

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

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

+130
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
package org.partiql.examples
2+
3+
import com.amazon.ion.IonDecimal
4+
import com.amazon.ion.IonStruct
5+
import com.amazon.ion.system.IonSystemBuilder
6+
import org.partiql.examples.util.Example
7+
import org.partiql.lang.CompilerPipeline
8+
import org.partiql.lang.errors.ErrorCode
9+
import org.partiql.lang.errors.Property
10+
import org.partiql.lang.errors.PropertyValueMap
11+
import org.partiql.lang.eval.BindingCase
12+
import org.partiql.lang.eval.BindingName
13+
import org.partiql.lang.eval.Bindings
14+
import org.partiql.lang.eval.EvaluationException
15+
import org.partiql.lang.eval.EvaluationSession
16+
import org.partiql.lang.eval.ExprValue
17+
import org.partiql.lang.eval.ExprValueFactory
18+
import org.partiql.lang.eval.ExprValueType
19+
import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
20+
import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedureSignature
21+
import org.partiql.lang.eval.stringValue
22+
import java.io.PrintStream
23+
import java.math.BigDecimal
24+
import java.math.RoundingMode
25+
26+
private val ion = IonSystemBuilder.standard().build()
27+
28+
/**
29+
* A simple custom stored procedure that calculates the moon weight for each crewmate of the given crew, storing the
30+
* moon weight in the [EvaluationSession] global bindings. This procedure also returns the number of crewmates we
31+
* calculated the moon weight for, returning -1 if no crew is found.
32+
*
33+
* This example demonstrates how to create a custom stored procedure, check argument types, and modify the
34+
* [EvaluationSession].
35+
*/
36+
class CalculateCrewMoonWeight(private val valueFactory: ExprValueFactory): StoredProcedure {
37+
private val MOON_GRAVITATIONAL_CONSTANT = BigDecimal(1.622 / 9.81)
38+
39+
// [StoredProcedureSignature] takes two arguments:
40+
// 1. the name of the stored procedure
41+
// 2. the arity of this stored procedure. Checks to arity are taken care of by the evaluator. However, we must
42+
// still check that the passed arguments are of the right type in our implementation of the procedure.
43+
override val signature = StoredProcedureSignature(name = "calculate_crew_moon_weight", arity = 1)
44+
45+
// `call` is where you define the logic of the stored procedure given an [EvaluationSession] and a list of
46+
// arguments
47+
override fun call(session: EvaluationSession, args: List<ExprValue>): ExprValue {
48+
// We first check that the first argument is a string
49+
val crewName = args.first()
50+
// In the future the evaluator will also verify function argument types, but for now we must verify their type
51+
// manually
52+
if (crewName.type != ExprValueType.STRING) {
53+
val errorContext = PropertyValueMap().also {
54+
it[Property.EXPECTED_ARGUMENT_TYPES] = "STRING"
55+
it[Property.ACTUAL_ARGUMENT_TYPES] = crewName.type.name
56+
it[Property.FUNCTION_NAME] = signature.name
57+
}
58+
throw EvaluationException("First argument to ${signature.name} was not a string",
59+
ErrorCode.EVALUATOR_INCORRECT_TYPE_OF_ARGUMENTS_TO_PROCEDURE_CALL,
60+
errorContext,
61+
internal = false)
62+
}
63+
64+
// Next we check if the given `crewName` is in the [EvaluationSession]'s global bindings. If not, we return 0.
65+
val sessionGlobals = session.globals
66+
val crewBindings = sessionGlobals[BindingName(crewName.stringValue(), BindingCase.INSENSITIVE)]
67+
?: return valueFactory.newInt(-1)
68+
69+
// Now that we've confirmed the given `crewName` is in the session's global bindings, we calculate and store
70+
// the moon weight for each crewmate in the crew.
71+
// In addition, we keep a running a tally of how many crewmates we do this for.
72+
var numCalculated = 0
73+
for (crewmateBinding in crewBindings) {
74+
val crewmate = crewmateBinding.ionValue as IonStruct
75+
val mass = crewmate["mass"] as IonDecimal
76+
val moonWeight = (mass.decimalValue() * MOON_GRAVITATIONAL_CONSTANT).setScale(1, RoundingMode.HALF_UP)
77+
crewmate.add("moonWeight", ion.newDecimal(moonWeight))
78+
79+
numCalculated++
80+
}
81+
return valueFactory.newInt(numCalculated)
82+
}
83+
}
84+
85+
/**
86+
* Demonstrates the use of custom stored procedure [CalculateCrewMoonWeight] in PartiQL queries.
87+
*/
88+
class CustomProceduresExample(out: PrintStream) : Example(out) {
89+
override fun run() {
90+
/**
91+
* To make custom stored procedures available to the PartiQL query being executed, they must be passed to
92+
* [CompilerPipeline.Builder.addProcedure].
93+
*/
94+
val pipeline = CompilerPipeline.build(ion) {
95+
addProcedure(CalculateCrewMoonWeight(valueFactory))
96+
}
97+
98+
// Here, we initialize the crews to be stored in our global session bindings
99+
val initialCrews = Bindings.ofMap(
100+
mapOf(
101+
"crew1" to pipeline.valueFactory.newFromIonValue(
102+
ion.singleValue("""[ { name: "Neil", mass: 80.5 },
103+
{ name: "Buzz", mass: 72.3 },
104+
{ name: "Michael", mass: 89.9 } ]""")),
105+
"crew2" to pipeline.valueFactory.newFromIonValue(
106+
ion.singleValue("""[ { name: "James", mass: 77.1 },
107+
{ name: "Spock", mass: 81.6 } ]"""))
108+
)
109+
)
110+
val session = EvaluationSession.build { globals(initialCrews) }
111+
112+
val crew1BindingName = BindingName("crew1", BindingCase.INSENSITIVE)
113+
val crew2BindingName = BindingName("crew2", BindingCase.INSENSITIVE)
114+
115+
out.println("Initial global session bindings:")
116+
print("Crew 1:", "${session.globals[crew1BindingName]}")
117+
print("Crew 2:", "${session.globals[crew2BindingName]}")
118+
119+
// We call our custom stored procedure using PartiQL's `EXEC` clause. Here we call our stored procedure
120+
// 'calculate_crew_moon_weight' with the arg 'crew1', which outputs the number of crewmates we've calculated
121+
// the moon weight for
122+
val procedureCall = "EXEC calculate_crew_moon_weight 'crew1'"
123+
val procedureCallOutput = pipeline.compile(procedureCall).eval(session)
124+
print("Number of calculated moon weights:", "$procedureCallOutput")
125+
126+
out.println("Updated global session bindings:")
127+
print("Crew 1:", "${session.globals[crew1BindingName]}")
128+
print("Crew 2:", "${session.globals[crew2BindingName]}")
129+
}
130+
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ private val examples = mapOf(
1515
// Kotlin Examples
1616
CsvExprValueExample::class.java.simpleName to CsvExprValueExample(System.out),
1717
CustomFunctionsExample::class.java.simpleName to CustomFunctionsExample(System.out),
18+
CustomProceduresExample::class.java.simpleName to CustomProceduresExample(System.out),
1819
EvaluationWithBindings::class.java.simpleName to EvaluationWithBindings(System.out),
1920
EvaluationWithLazyBindings::class.java.simpleName to EvaluationWithLazyBindings(System.out),
2021
ParserErrorExample::class.java.simpleName to ParserErrorExample(System.out),

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

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package org.partiql.examples
2+
3+
import org.partiql.examples.util.Example
4+
import java.io.PrintStream
5+
6+
class CustomProceduresExampleTest : BaseExampleTest() {
7+
override fun example(out: PrintStream): Example = CustomProceduresExample(out)
8+
9+
override val expected = """
10+
|Initial global session bindings:
11+
|Crew 1:
12+
| [{'name': 'Neil', 'mass': 80.5}, {'name': 'Buzz', 'mass': 72.3}, {'name': 'Michael', 'mass': 89.9}]
13+
|Crew 2:
14+
| [{'name': 'James', 'mass': 77.1}, {'name': 'Spock', 'mass': 81.6}]
15+
|Number of calculated moon weights:
16+
| 3
17+
|Updated global session bindings:
18+
|Crew 1:
19+
| [{'name': 'Neil', 'mass': 80.5, 'moonWeight': 13.3}, {'name': 'Buzz', 'mass': 72.3, 'moonWeight': 12.0}, {'name': 'Michael', 'mass': 89.9, 'moonWeight': 14.9}]
20+
|Crew 2:
21+
| [{'name': 'James', 'mass': 77.1}, {'name': 'Spock', 'mass': 81.6}]
22+
|
23+
""".trimMargin()
24+
}

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

+6-1
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,12 @@
2626
(where where::(? expr)))
2727

2828
// Data definition operations also cannot be composed with other `expr` nodes.
29-
(ddl op::ddl_op))
29+
(ddl op::ddl_op)
30+
31+
// Stored procedure calls are only allowed at the top level of a query and cannot be used as an expression
32+
// Currently supports stored procedure calls with the unnamed argument syntax:
33+
// EXEC <symbol> [<expr>.*]
34+
(exec procedure_name::symbol args::(* expr 0)))
3035

3136
// The expressions that can result in values.
3237
(sum expr

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

+26-3
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import com.amazon.ion.*
1818
import org.partiql.lang.ast.*
1919
import org.partiql.lang.eval.*
2020
import org.partiql.lang.eval.builtins.*
21+
import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
2122
import org.partiql.lang.syntax.*
2223

2324
/**
@@ -35,7 +36,13 @@ data class StepContext(
3536
* Includes built-in functions as well as custom functions added while the [CompilerPipeline]
3637
* was being built.
3738
*/
38-
val functions: @JvmSuppressWildcards Map<String, ExprFunction>
39+
val functions: @JvmSuppressWildcards Map<String, ExprFunction>,
40+
41+
/**
42+
* Returns a list of all stored procedures which are available for execution.
43+
* Only includes the custom stored procedures added while the [CompilerPipeline] was being built.
44+
*/
45+
val procedures: @JvmSuppressWildcards Map<String, StoredProcedure>
3946
)
4047

4148
/**
@@ -65,6 +72,12 @@ interface CompilerPipeline {
6572
*/
6673
val functions: @JvmSuppressWildcards Map<String, ExprFunction>
6774

75+
/**
76+
* Returns a list of all stored procedures which are available for execution.
77+
* Only includes the custom stored procedures added while the [CompilerPipeline] was being built.
78+
*/
79+
val procedures: @JvmSuppressWildcards Map<String, StoredProcedure>
80+
6881
/** Compiles the specified PartiQL query using the configured parser. */
6982
fun compile(query: String): Expression
7083

@@ -106,6 +119,7 @@ interface CompilerPipeline {
106119
private var parser: Parser? = null
107120
private var compileOptions: CompileOptions? = null
108121
private val customFunctions: MutableMap<String, ExprFunction> = HashMap()
122+
private val customProcedures: MutableMap<String, StoredProcedure> = HashMap()
109123
private val preProcessingSteps: MutableList<ProcessingStep> = ArrayList()
110124

111125
/**
@@ -137,6 +151,13 @@ interface CompilerPipeline {
137151
*/
138152
fun addFunction(function: ExprFunction): Builder = this.apply { customFunctions[function.name] = function }
139153

154+
/**
155+
* Add a custom stored procedure which will be callable by the compiled queries.
156+
*
157+
* Stored procedures added here will replace any built-in procedure with the same name.
158+
*/
159+
fun addProcedure(procedure: StoredProcedure): Builder = this.apply { customProcedures[procedure.signature.name] = procedure }
160+
140161
/** Adds a preprocessing step to be executed after parsing but before compilation. */
141162
fun addPreprocessingStep(step: ProcessingStep): Builder = this.apply { preProcessingSteps.add(step) }
142163

@@ -153,6 +174,7 @@ interface CompilerPipeline {
153174
parser ?: SqlParser(valueFactory.ion),
154175
compileOptions ?: CompileOptions.standard(),
155176
allFunctions,
177+
customProcedures,
156178
preProcessingSteps)
157179
}
158180
}
@@ -163,17 +185,18 @@ private class CompilerPipelineImpl(
163185
private val parser: Parser,
164186
override val compileOptions: CompileOptions,
165187
override val functions: Map<String, ExprFunction>,
188+
override val procedures: Map<String, StoredProcedure>,
166189
private val preProcessingSteps: List<ProcessingStep>
167190
) : CompilerPipeline {
168191

169-
private val compiler = EvaluatingCompiler(valueFactory, functions, compileOptions)
192+
private val compiler = EvaluatingCompiler(valueFactory, functions, procedures, compileOptions)
170193

171194
override fun compile(query: String): Expression {
172195
return compile(parser.parseExprNode(query))
173196
}
174197

175198
override fun compile(query: ExprNode): Expression {
176-
val context = StepContext(valueFactory, compileOptions, functions)
199+
val context = StepContext(valueFactory, compileOptions, functions, procedures)
177200

178201
val preProcessedQuery = preProcessingSteps.fold(query) { currentExprNode, step ->
179202
step(currentExprNode, context)

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

+1
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ private class AstSerializerImpl(val astVersion: AstVersion, val ion: IonSystem):
9494
is DropTable -> case { writeDropTable(expr) }
9595
is DropIndex -> case { writeDropIndex(expr) }
9696
is Parameter -> case { writeParameter(expr)}
97+
is Exec -> throw UnsupportedOperationException("EXEC clause not supported by the V0 AST")
9798
}.toUnit()
9899
}
99100
}

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

+20-3
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package org.partiql.lang.ast
22

33
import com.amazon.ionelement.api.toIonElement
44
import org.partiql.lang.domains.PartiqlAst
5+
import org.partiql.pig.runtime.SymbolPrimitive
56
import org.partiql.pig.runtime.asPrimitive
67

78
/** Converts an [ExprNode] to a [PartiqlAst.statement]. */
@@ -16,12 +17,16 @@ fun ExprNode.toAstStatement(): PartiqlAst.Statement {
1617

1718
is CreateTable, is CreateIndex, is DropTable, is DropIndex -> toAstDdl()
1819

20+
is Exec -> toAstExec()
1921
}
2022
}
2123

2224
private fun PartiQlMetaContainer.toElectrolyteMetaContainer(): ElectrolyteMetaContainer =
2325
com.amazon.ionelement.api.metaContainerOf(map { it.tag to it })
2426

27+
private fun SymbolicName.toSymbolPrimitive() : SymbolPrimitive =
28+
SymbolPrimitive(this.name, this.metas.toElectrolyteMetaContainer())
29+
2530
private fun ExprNode.toAstDdl(): PartiqlAst.Statement {
2631
val thiz = this
2732
val metas = metas.toElectrolyteMetaContainer()
@@ -30,7 +35,7 @@ private fun ExprNode.toAstDdl(): PartiqlAst.Statement {
3035
when(thiz) {
3136
is Literal, is LiteralMissing, is VariableReference, is Parameter, is NAry, is CallAgg, is Typed,
3237
is Path, is SimpleCase, is SearchedCase, is Select, is Struct, is Seq,
33-
is DataManipulation -> error("Can't convert ${thiz.javaClass} to PartiqlAst.ddl")
38+
is DataManipulation, is Exec -> error("Can't convert ${thiz.javaClass} to PartiqlAst.ddl")
3439

3540
is CreateTable -> ddl(createTable(thiz.tableName), metas)
3641
is CreateIndex -> ddl(createIndex(identifier(thiz.tableName, caseSensitive()), thiz.keys.map { it.toAstExpr() }), metas)
@@ -48,6 +53,18 @@ private fun ExprNode.toAstDdl(): PartiqlAst.Statement {
4853
}
4954
}
5055

56+
private fun ExprNode.toAstExec() : PartiqlAst.Statement {
57+
val node = this
58+
val metas = metas.toElectrolyteMetaContainer()
59+
60+
return PartiqlAst.build {
61+
when (node) {
62+
is Exec -> exec_(node.procedureName.toSymbolPrimitive(), node.args.map { it.toAstExpr() }, metas)
63+
else -> error("Can't convert ${node.javaClass} to PartiqlAst.Statement.Exec")
64+
}
65+
}
66+
}
67+
5168
fun ExprNode.toAstExpr(): PartiqlAst.Expr {
5269
val node = this
5370
val metas = this.metas.toElectrolyteMetaContainer()
@@ -147,8 +164,8 @@ fun ExprNode.toAstExpr(): PartiqlAst.Expr {
147164
SeqType.BAG -> bag(node.values.map { it.toAstExpr() })
148165
}
149166

150-
// These are handled by `toAstDml()`
151-
is DataManipulation, is CreateTable, is CreateIndex, is DropTable, is DropIndex ->
167+
// These are handled by `toAstDml()`, `toAstDdl()`, and `toAstExec()`
168+
is DataManipulation, is CreateTable, is CreateIndex, is DropTable, is DropIndex, is Exec ->
152169
error("Can't transform ${node.javaClass} to a PartiqlAst.expr }")
153170
}
154171
}

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

+6
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package org.partiql.lang.ast
22

33
import com.amazon.ion.IonSystem
44
import com.amazon.ionelement.api.toIonValue
5+
import org.partiql.lang.domains.PartiqlAst
56
import org.partiql.lang.domains.PartiqlAst.CaseSensitivity
67
import org.partiql.lang.domains.PartiqlAst.DdlOp
78
import org.partiql.lang.domains.PartiqlAst.DmlOp
@@ -33,6 +34,7 @@ private class StatementTransformer(val ion: IonSystem) {
3334
is Statement.Query -> stmt.toExprNode()
3435
is Statement.Dml -> stmt.toExprNode()
3536
is Statement.Ddl -> stmt.toExprNode()
37+
is Statement.Exec -> stmt.toExprNode()
3638
}
3739

3840
private fun ElectrolyteMetaContainer.toPartiQlMetaContainer(): PartiQlMetaContainer {
@@ -344,4 +346,8 @@ private class StatementTransformer(val ion: IonSystem) {
344346
metas = metas)
345347
}
346348
}
349+
350+
private fun Statement.Exec.toExprNode(): ExprNode {
351+
return Exec(procedureName.toSymbolicName(), this.args.toExprNodeList(), metas.toPartiQlMetaContainer())
352+
}
347353
}

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

+16
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@ sealed class ExprNode : AstNode(), HasMetas {
103103
is Parameter -> {
104104
copy(metas = metas)
105105
}
106+
is Exec -> {
107+
copy(metas = metas)
108+
}
106109
}
107110
}
108111
}
@@ -210,6 +213,19 @@ data class Typed(
210213
override val children: List<AstNode> = listOf(expr, type)
211214
}
212215

216+
//********************************
217+
// Stored procedure clauses
218+
//********************************
219+
220+
/** Represents a call to a stored procedure, i.e. `EXEC stored_procedure [<expr>.*]` */
221+
data class Exec(
222+
val procedureName: SymbolicName,
223+
val args: List<ExprNode>,
224+
override val metas: MetaContainer
225+
) : ExprNode() {
226+
override val children: List<AstNode> = args
227+
}
228+
213229
//********************************
214230
// Path expressions
215231
//********************************

0 commit comments

Comments
 (0)