From 04e342bb36067349a6a52ad7c6d8fb5d167382d5 Mon Sep 17 00:00:00 2001 From: Philip Lykke Carlsen Date: Mon, 22 Jan 2024 15:05:31 +0100 Subject: [PATCH] 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 --- crates/query-engine/sql/src/sql/ast.rs | 6 + crates/query-engine/sql/src/sql/convert.rs | 20 +++ .../src/translation/query/filtering.rs | 1 + .../src/translation/query/native_queries.rs | 10 ++ .../src/translation/query/values.rs | 132 +++++++++++++----- .../select_array_variable/request.json | 27 ++++ .../select_array_variable/tables.json | 27 ++++ .../request.json | 107 ++++++++++++++ .../tables.json | 86 ++++++++++++ .../tests__select_array_variable.snap | 43 ++++++ ...s__select_array_variable_nested_types.snap | 66 +++++++++ .../query-engine/translation/tests/tests.rs | 12 ++ .../databases-tests/src/aurora/query_tests.rs | 12 ++ ...y_tests__basic__select_array_variable.snap | 20 +++ ...c__select_array_variable_nested_types.snap | 26 ++++ .../databases-tests/src/citus/query_tests.rs | 12 ++ ...tests__explain__select_where_variable.snap | 13 +- ...y_tests__basic__select_array_variable.snap | 20 +++ ...c__select_array_variable_nested_types.snap | 26 ++++ ...schema_tests__schema_test__get_schema.snap | 108 ++++++++++++++ .../src/cockroach/query_tests.rs | 13 ++ ...tests__explain__select_where_variable.snap | 13 +- ...y_tests__basic__select_array_variable.snap | 20 +++ ...schema_tests__schema_test__get_schema.snap | 35 +++++ .../src/postgres/query_tests.rs | 12 ++ ...in__native_queries__embedded_variable.snap | 45 +++--- ...ts__explain__select_where_no_variable.snap | 13 +- ...tests__explain__select_where_variable.snap | 13 +- ...y_tests__basic__select_array_variable.snap | 20 +++ ...c__select_array_variable_nested_types.snap | 26 ++++ ...schema_tests__schema_test__get_schema.snap | 108 ++++++++++++++ .../src/yugabyte/query_tests.rs | 12 ++ ...y_tests__basic__select_array_variable.snap | 20 +++ ...c__select_array_variable_nested_types.snap | 26 ++++ .../goldenfiles/select_array_variable.json | 27 ++++ .../select_array_variable_nested_types.json | 107 ++++++++++++++ .../v2-chinook-deployment-template.json | 96 +++++++++++++ static/citus/v2-chinook-deployment.json | 96 +++++++++++++ static/cockroach/v2-chinook-deployment.json | 26 ++++ static/composite-types-complex.sql | 13 ++ static/postgres/v2-chinook-deployment.json | 96 +++++++++++++ static/yugabyte/v2-chinook-deployment.json | 96 +++++++++++++ 42 files changed, 1604 insertions(+), 103 deletions(-) create mode 100644 crates/query-engine/translation/tests/goldenfiles/select_array_variable/request.json create mode 100644 crates/query-engine/translation/tests/goldenfiles/select_array_variable/tables.json create mode 100644 crates/query-engine/translation/tests/goldenfiles/select_array_variable_nested_types/request.json create mode 100644 crates/query-engine/translation/tests/goldenfiles/select_array_variable_nested_types/tables.json create mode 100644 crates/query-engine/translation/tests/snapshots/tests__select_array_variable.snap create mode 100644 crates/query-engine/translation/tests/snapshots/tests__select_array_variable_nested_types.snap create mode 100644 crates/tests/databases-tests/src/aurora/snapshots/databases_tests__aurora__query_tests__basic__select_array_variable.snap create mode 100644 crates/tests/databases-tests/src/aurora/snapshots/databases_tests__aurora__query_tests__basic__select_array_variable_nested_types.snap create mode 100644 crates/tests/databases-tests/src/citus/snapshots/databases_tests__citus__query_tests__basic__select_array_variable.snap create mode 100644 crates/tests/databases-tests/src/citus/snapshots/databases_tests__citus__query_tests__basic__select_array_variable_nested_types.snap create mode 100644 crates/tests/databases-tests/src/cockroach/snapshots/databases_tests__cockroach__query_tests__basic__select_array_variable.snap create mode 100644 crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__query_tests__basic__select_array_variable.snap create mode 100644 crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__query_tests__basic__select_array_variable_nested_types.snap create mode 100644 crates/tests/databases-tests/src/yugabyte/snapshots/databases_tests__yugabyte__query_tests__basic__select_array_variable.snap create mode 100644 crates/tests/databases-tests/src/yugabyte/snapshots/databases_tests__yugabyte__query_tests__basic__select_array_variable_nested_types.snap create mode 100644 crates/tests/tests-common/goldenfiles/select_array_variable.json create mode 100644 crates/tests/tests-common/goldenfiles/select_array_variable_nested_types.json diff --git a/crates/query-engine/sql/src/sql/ast.rs b/crates/query-engine/sql/src/sql/ast.rs index 6e45ecdf7..f28574ecb 100644 --- a/crates/query-engine/sql/src/sql/ast.rs +++ b/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), + CorrelatedSubSelect(Box