Support variables of array type (#263)

### What

This PR extends the support for query variables to include arrays.

### How

We need to go a bit out of our way to be able to translate JSON arrays
to Postgres arrays, since there is no single builtin function that does
this.

As an example taken from the tests, a variable of type
array-of-`organization` gets translated to the below SQL:

```
(
  SELECT
    array_agg(
      jsonb_populate_record(cast(null as organization), "%0_arr"."elem")
    ) AS "elem"
  FROM
    jsonb_array_elements(("%0_%variables_table"."%variables" -> $2)) AS "%0_arr"("elem")
)

```

Which, being a single-element set-returning expression, can feature both
in a `FROM` clause and in a select-list or other expression.

---------

Co-authored-by: Gil Mizrahi <[email protected]>

crates/query-engine/sql/src/sql/ast.rs

@@ -103,6 +103,11 @@ pub enum From {
         alias: TableAlias,
         columns: Vec<(ColumnAlias, ScalarType)>,
     },
+    JsonbArrayElements {
+        expression: Expression,
+        alias: TableAlias,
+        column: ColumnAlias,
+    },
 }
 
 /// A JOIN clause
@@ -234,6 +239,7 @@ pub enum Expression {
     /// A COUNT clause
     Count(CountType),
     ArrayConstructor(Vec<Expression>),
+    CorrelatedSubSelect(Box<Select>),
 }
 
 /// An unary operator

crates/query-engine/sql/src/sql/convert.rs

@@ -222,6 +222,21 @@ impl From {
                 }
                 sql.append_syntax(")");
             }
+            From::JsonbArrayElements {
+                expression,
+                alias,
+                column,
+            } => {
+                sql.append_syntax("jsonb_array_elements");
+                sql.append_syntax("(");
+                expression.to_sql(sql);
+                sql.append_syntax(")");
+                sql.append_syntax(" AS ");
+                alias.to_sql(sql);
+                sql.append_syntax("(");
+                column.to_sql(sql);
+                sql.append_syntax(")");
+            }
         }
     }
 }
@@ -407,6 +422,11 @@ impl Expression {
                 }
                 sql.append_syntax("]");
             }
+            Expression::CorrelatedSubSelect(select) => {
+                sql.append_syntax("(");
+                select.to_sql(sql);
+                sql.append_syntax(")");
+            }
         }
     }
 }

crates/query-engine/translation/src/translation/query/filtering.rs

@@ -380,6 +380,7 @@ fn translate_comparison_value(
         )),
         models::ComparisonValue::Variable { name: var } => Ok((
             values::translate_variable(
+                state,
                 env.get_variables_table()?,
                 var.clone(),
                 &database::Type::ScalarType(typ.clone()),

crates/query-engine/translation/src/translation/query/native_queries.rs

@@ -38,6 +38,16 @@ pub fn translate(env: &Env, state: State) -> Result<Vec<sql::ast::CommonTableExp
                             models::Argument::Variable { name } => match &variables_table {
                                 Err(err) => Err(err.clone()),
                                 Ok(variables_table) => Ok(values::translate_variable(
+                                    // We need a 'State' value when translating variables in order
+                                    // to be able to generate fresh names for bound relational
+                                    // expressions.
+                                    // However, we cannot readily re-use the 'state' argument we're
+                                    // given, since that is an "owned" value which is "moved" in
+                                    // the course of iterating over the native queries accumulated
+                                    // within.
+                                    // Ideally we should resolve this by tracking native queries
+                                    // separately from the fresh table name counter.
+                                    &mut State::new(),
                                     variables_table.clone(),
                                     name.clone(),
                                     &typ,

crates/query-engine/translation/src/translation/query/values.rs

@@ -1,8 +1,8 @@
 //! Handle the translation of literal values.
 
-use crate::translation::error::Error;
+use crate::translation::{error::Error, helpers::State};
 use query_engine_metadata::metadata::database;
-use query_engine_sql::sql;
+use query_engine_sql::sql::{self, ast::ColumnReference, helpers::simple_select};
 use sql::ast::{Expression, Value};
 
 /// Convert a JSON value into a SQL value.
@@ -84,48 +84,108 @@ fn type_to_ast_scalar_type(typ: &database::Type) -> sql::ast::ScalarType {
 
 /// Convert a variable into a SQL value.
 pub fn translate_variable(
+    state: &mut State,
     variables_table: sql::ast::TableReference,
     variable: String,
     r#type: &database::Type,
 ) -> sql::ast::Expression {
-    let variables_reference =
-        Expression::ColumnReference(sql::ast::ColumnReference::AliasedColumn {
-            table: variables_table,
-            column: sql::helpers::make_column_alias(sql::helpers::VARIABLES_FIELD.to_string()),
-        });
+    let variables_reference = Expression::ColumnReference(ColumnReference::AliasedColumn {
+        table: variables_table,
+        column: sql::helpers::make_column_alias(sql::helpers::VARIABLES_FIELD.to_string()),
+    });
 
+    // We use the binop '->' to project (as jsonb) the value of a variable from the data column of
+    // the variable table.
+    let projected_variable_exp = sql::ast::Expression::BinaryOperation {
+        left: Box::new(variables_reference),
+        operator: sql::ast::BinaryOperator("->".to_string()),
+        right: Box::new(sql::ast::Expression::Value(sql::ast::Value::String(
+            variable.clone(),
+        ))),
+    };
+
+    translate_projected_variable(state, r#type, projected_variable_exp)
+}
+
+/// Produce a SQL expression that translates an expression of Postgres type 'jsonb' into a given
+/// type.
+///
+/// For scalar types and object types this is a simple operation, since we can rely on builtin
+/// functions.
+///
+/// Arrays are more complex since there isn't a builtin function that handles array
+/// types.
+pub fn translate_projected_variable(
+    state: &mut State,
+    r#type: &database::Type,
+    exp: sql::ast::Expression,
+) -> sql::ast::Expression {
     match r#type {
-        database::Type::CompositeType(_type_name) => {
-            let exp = sql::ast::Expression::BinaryOperation {
-                left: Box::new(variables_reference),
-                operator: sql::ast::BinaryOperator("->".to_string()),
-                right: Box::new(sql::ast::Expression::Value(sql::ast::Value::String(
-                    variable.clone(),
-                ))),
+        database::Type::CompositeType(_type_name) => sql::ast::Expression::FunctionCall {
+            function: sql::ast::Function::JsonbPopulateRecord,
+            args: vec![
+                sql::ast::Expression::Cast {
+                    expression: Box::new(sql::ast::Expression::Value(sql::ast::Value::Null)),
+                    r#type: type_to_ast_scalar_type(r#type),
+                },
+                exp,
+            ],
+        },
+        // We translate projection of array types into the following sql:
+        // ```
+        // ( SELECT
+        //     array_agg(
+        //       jsonb_populate_record(cast(null as <type>), "array"."element")
+        //     ) AS "element"
+        //   FROM
+        //    jsonb_array_elements((<variable_table> -> <label>)) AS "array"("element")
+        // )
+        // ```
+        database::Type::ArrayType(type_name) => {
+            let array_table = state.make_table_alias("array".to_string());
+            let element_column = sql::ast::ColumnAlias {
+                name: "element".to_string(),
             };
-            sql::ast::Expression::FunctionCall {
-                function: sql::ast::Function::JsonbPopulateRecord,
-                args: vec![
-                    sql::ast::Expression::Cast {
-                        expression: Box::new(sql::ast::Expression::Value(sql::ast::Value::Null)),
-                        r#type: type_to_ast_scalar_type(r#type),
-                    },
-                    exp,
-                ],
-            }
-        }
-        _ => {
-            let exp = sql::ast::Expression::BinaryOperation {
-                left: Box::new(variables_reference),
-                operator: sql::ast::BinaryOperator("->>".to_string()),
-                right: Box::new(sql::ast::Expression::Value(sql::ast::Value::String(
-                    variable.clone(),
-                ))),
+
+            let from_arr = sql::ast::From::JsonbArrayElements {
+                expression: exp,
+                alias: array_table.clone(),
+                column: element_column.clone(),
             };
-            sql::ast::Expression::Cast {
-                expression: Box::new(exp),
-                r#type: type_to_ast_scalar_type(r#type),
-            }
+
+            let element_expression =
+                sql::ast::Expression::ColumnReference(ColumnReference::AliasedColumn {
+                    table: sql::ast::TableReference::AliasedTable(array_table.clone()),
+                    column: element_column.clone(),
+                });
+
+            let converted_element_exp =
+                translate_projected_variable(state, type_name, element_expression);
+
+            let mut result_select = simple_select(vec![(
+                element_column.clone(),
+                sql::ast::Expression::FunctionCall {
+                    function: sql::ast::Function::Unknown("array_agg".to_string()),
+                    args: vec![converted_element_exp],
+                },
+            )]);
+
+            result_select.from = Some(from_arr);
+
+            sql::ast::Expression::CorrelatedSubSelect(Box::new(result_select))
         }
+        database::Type::ScalarType(_) => sql::ast::Expression::Cast {
+            expression: Box::new(sql::ast::Expression::BinaryOperation {
+                left: Box::new(exp),
+                operator: sql::ast::BinaryOperator("#>>".to_string()),
+                right: Box::new(sql::ast::Expression::Cast {
+                    expression: Box::new(sql::ast::Expression::Value(sql::ast::Value::Array(
+                        vec![],
+                    ))),
+                    r#type: sql::ast::ScalarType("text[]".to_string()),
+                }),
+            }),
+            r#type: type_to_ast_scalar_type(r#type),
+        },
     }
 }

crates/query-engine/translation/tests/goldenfiles/select_array_variable/request.json

@@ -0,0 +1,27 @@
+{
+  "collection": "count_elements",
+  "query": {
+    "fields": {
+      "result": {
+        "type": "column",
+        "column": "result",
+        "arguments": {}
+      }
+    }
+  },
+  "arguments": {
+    "array_argument": {
+      "type": "variable",
+      "name": "variable_array_argument"
+    }
+  },
+  "collection_relationships": {},
+  "variables": [
+    {
+      "variable_array_argument": ["one", "two"]
+    },
+    {
+      "variable_array_argument": ["one", "two", "three"]
+    }
+  ]
+}

crates/query-engine/translation/tests/goldenfiles/select_array_variable/tables.json

@@ -0,0 +1,27 @@
+{
+  "nativeQueries": {
+    "count_elements": {
+      "sql": "SELECT array_length({{array_argument}}, 1) as result",
+      "columns": {
+        "result": {
+          "name": "result",
+          "type": {
+            "scalarType": "int4"
+          },
+          "nullable": "nullable",
+          "description": null
+        }
+      },
+      "arguments": {
+        "array_argument": {
+          "name": "array_argument",
+          "type": {
+            "arrayType": { "scalarType": "text" }
+          },
+          "nullable": "nullable"
+        }
+      },
+      "description": "A native query used to test support array-valued variables"
+    }
+  }
+}

crates/query-engine/translation/tests/goldenfiles/select_array_variable_nested_types/request.json

@@ -0,0 +1,107 @@
+{
+  "collection": "summarize_organizations",
+  "query": {
+    "fields": {
+      "result": {
+        "type": "column",
+        "column": "result",
+        "arguments": {}
+      }
+    }
+  },
+  "arguments": {
+    "organizations": {
+      "type": "variable",
+      "name": "variable_organizations"
+    }
+  },
+  "collection_relationships": {},
+  "variables": [
+    {
+      "variable_organizations": []
+    },
+    {
+      "variable_organizations": [
+        {
+          "name": "Federation of United Solipsists",
+          "committees": [
+            {
+              "name": "Positively existing entities",
+              "members": [
+                {
+                  "first_name": "Yo",
+                  "last_name": "El mismo"
+                }
+              ]
+            }
+          ]
+        }
+      ]
+    },
+    {
+      "variable_organizations": [
+        {
+          "name": "RC Model Airplane Enthusiasts",
+          "committees": [
+            {
+              "name": "Founders",
+              "members": [
+                {
+                  "first_name": "Orville",
+                  "last_name": "Wright"
+                },
+                {
+                  "first_name": "Wilbur",
+                  "last_name": "Wright"
+                }
+              ]
+            },
+            {
+              "name": "Parts supply management",
+              "members": [
+                {
+                  "first_name": "Orville",
+                  "last_name": "Wright"
+                },
+                {
+                  "first_name": "Wilbur",
+                  "last_name": "Wright"
+                },
+                {
+                  "first_name": "Guybrush",
+                  "last_name": "Threepwood"
+                }
+              ]
+            }
+          ]
+        },
+        {
+          "name": "Argonauts' Alumni Association",
+          "committees": [
+            {
+              "name": "Crew",
+              "members": [
+                {
+                  "first_name": "Jason",
+                  "last_name": "(The)"
+                },
+                {
+                  "first_name": "Heracles",
+                  "last_name": "(The)"
+                },
+                {
+                  "first_name": "Castor",
+                  "last_name": "(The)"
+                },
+                {
+                  "first_name": "Polydeuces",
+                  "last_name": "(The)"
+                }
+              ]
+            }
+          ]
+        }
+      ]
+    }
+  ]
+}