From dc293964bf982845a4ac2a77502df2dd61a14cd8 Mon Sep 17 00:00:00 2001 From: Onur Tirtir Date: Mon, 7 Feb 2022 16:17:38 +0300 Subject: [PATCH] Expand target entries with merged array/jsonb subscripting ops into multiple ones When re-writing query tree, postgres combines multiple subscripting operators referencing to the same attribute into a single `TargetEntry` by linking `SubscriptingRef` objects to each other via their `refexpr` fields. (See `rewriteTargetListIU` & `process_matched_tle` functions.) However, ruleutils function `get_update_query_targetlist_def` doesn't know how to properly deparse such a `TargetEntry`. As a result, we were only taking the last set-by-subscript operation into account when generating the shard query for such an `UPDATE` command. In postgres, this doesn't cause any problems (e.g.: when generating definition of a rule based object) since the query-rewrite transformations aren't performed on the query tree that `get_update_query_targetlist_def` is expected to process. For this reason; with this commit, before processing the target entry list in our ruleutils based deparser, we first expand such target entries into multiple ones. To detect such `SubscriptingRef` objects, we also need to investigate `FieldStore` and `CoerceToDomain` objects as postgres functions processing `SubscriptingRef` objects do --although they do so for different purposes. However, given that Citus already doesn't allow `INSERT/UPDATE` via `FieldStore`, we only do that for `CoerceToDomain` objects. --- .../distributed/deparser/citus_ruleutils.c | 127 ++++++++++++++++++ .../distributed/deparser/ruleutils_12.c | 2 + .../distributed/deparser/ruleutils_13.c | 2 + .../distributed/deparser/ruleutils_14.c | 2 + .../planner/multi_router_planner.c | 7 +- src/backend/distributed/utils/listutils.c | 17 +++ src/include/distributed/citus_ruleutils.h | 1 + src/include/distributed/listutils.h | 1 + .../expected/multi_deparse_shard_query.out | 4 - .../regress/expected/multi_test_helpers.out | 4 + src/test/regress/expected/pg14.out | 77 +++++++++++ src/test/regress/expected/subscripting_op.out | 120 +++++++++++++++++ src/test/regress/multi_schedule | 2 +- .../regress/sql/multi_deparse_shard_query.sql | 5 - src/test/regress/sql/multi_test_helpers.sql | 5 + src/test/regress/sql/pg14.sql | 59 ++++++++ src/test/regress/sql/subscripting_op.sql | 97 +++++++++++++ 17 files changed, 521 insertions(+), 11 deletions(-) create mode 100644 src/test/regress/expected/subscripting_op.out create mode 100644 src/test/regress/sql/subscripting_op.sql diff --git a/src/backend/distributed/deparser/citus_ruleutils.c b/src/backend/distributed/deparser/citus_ruleutils.c index a2002851d5f..e128e19f338 100644 --- a/src/backend/distributed/deparser/citus_ruleutils.c +++ b/src/backend/distributed/deparser/citus_ruleutils.c @@ -80,6 +80,7 @@ static void deparse_index_columns(StringInfo buffer, List *indexParameterList, static void AppendStorageParametersToString(StringInfo stringBuffer, List *optionList); static void simple_quote_literal(StringInfo buf, const char *val); +static SubscriptingRef * TargetEntryExprFindSubsRef(Expr *expr); static char * flatten_reloptions(Oid relid); static void AddVacuumParams(ReindexStmt *reindexStmt, StringInfo buffer); @@ -1353,3 +1354,129 @@ RoleSpecString(RoleSpec *spec, bool withQuoteIdentifier) } } } + + +/* + * ExpandMergedSubscriptingRefEntries takes a list of target entries and expands + * each one that references a SubscriptingRef node that indicates multiple (field) + * updates on the same attribute, which is applicable for array/json types atm. + */ +List * +ExpandMergedSubscriptingRefEntries(List *targetEntryList) +{ + List *newTargetEntryList = NIL; + + TargetEntry *targetEntry = NULL; + foreach_ptr(targetEntry, targetEntryList) + { + List *expandedTargetEntries = NIL; + + Expr *expr = targetEntry->expr; + while (expr) + { + SubscriptingRef *subsRef = TargetEntryExprFindSubsRef(expr); + if (!subsRef) + { + break; + } + + /* + * Remove refexpr from the SubscriptingRef that we are about to + * wrap in a new TargetEntry and save it for the next one. + */ + Expr *refexpr = subsRef->refexpr; + subsRef->refexpr = NULL; + + /* + * Wrap the Expr that holds SubscriptingRef (directly or indirectly) + * in a new TargetEntry; note that it doesn't have a refexpr anymore. + */ + TargetEntry *newTargetEntry = copyObject(targetEntry); + newTargetEntry->expr = expr; + expandedTargetEntries = lappend(expandedTargetEntries, newTargetEntry); + + /* now inspect the refexpr that SubscriptingRef at hand were holding */ + expr = refexpr; + } + + if (expandedTargetEntries == NIL) + { + /* return original entry since it doesn't hold a SubscriptingRef node */ + newTargetEntryList = lappend(newTargetEntryList, targetEntry); + } + else + { + /* + * Need to concat expanded target list entries in reverse order + * to preserve ordering of the original target entry list. + */ + newTargetEntryList = list_concat(newTargetEntryList, + list_reverse(expandedTargetEntries)); + } + } + + return newTargetEntryList; +} + + +/* + * TargetEntryExprFindSubsRef searches given Expr --assuming that it is part + * of a target list entry-- to see if it directly (i.e.: itself) or indirectly + * (e.g.: behind some level of coercions) holds a SubscriptingRef node. + * + * Returns the original SubscriptingRef node on success or NULL otherwise. + * + * Note that it wouldn't add much value to use expression_tree_walker here + * since we are only interested in a subset of the fields of a few certain + * node types. + */ +static SubscriptingRef * +TargetEntryExprFindSubsRef(Expr *expr) +{ + Node *node = (Node *) expr; + while (node) + { + if (IsA(node, FieldStore)) + { + /* + * ModifyPartialQuerySupported doesn't allow INSERT/UPDATE via + * FieldStore. If we decide supporting such commands, then we + * should take the first element of "newvals" list into account + * here. This is because, to support such commands, we will need + * to expand merged FieldStore into separate target entries too. + * + * For this reason, this block is not reachable atm and need to + * uncomment the following if we decide supporting such commands. + * + * """ + * FieldStore *fieldStore = (FieldStore *) node; + * node = (Node *) linitial(fieldStore->newvals); + * """ + */ + ereport(ERROR, (errmsg("unexpectedly got FieldStore object when " + "generating shard query"))); + } + else if (IsA(node, CoerceToDomain)) + { + CoerceToDomain *coerceToDomain = (CoerceToDomain *) node; + if (coerceToDomain->coercionformat != COERCE_IMPLICIT_CAST) + { + /* not an implicit coercion, cannot reach to a SubscriptingRef */ + break; + } + + node = (Node *) coerceToDomain->arg; + } + else if (IsA(node, SubscriptingRef)) + { + return (SubscriptingRef *) node; + } + else + { + /* got a node that we are not interested in */ + break; + } + } + + return NULL; +} diff --git a/src/backend/distributed/deparser/ruleutils_12.c b/src/backend/distributed/deparser/ruleutils_12.c index 513678df116..29d860bd214 100644 --- a/src/backend/distributed/deparser/ruleutils_12.c +++ b/src/backend/distributed/deparser/ruleutils_12.c @@ -3324,6 +3324,8 @@ get_update_query_targetlist_def(Query *query, List *targetList, SubLink *cur_ma_sublink; List *ma_sublinks; + targetList = ExpandMergedSubscriptingRefEntries(targetList); + /* * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks * into a list. We expect them to appear, in ID order, in resjunk tlist diff --git a/src/backend/distributed/deparser/ruleutils_13.c b/src/backend/distributed/deparser/ruleutils_13.c index 5c1749c5f68..3ce4e4e66d6 100644 --- a/src/backend/distributed/deparser/ruleutils_13.c +++ b/src/backend/distributed/deparser/ruleutils_13.c @@ -3263,6 +3263,8 @@ get_update_query_targetlist_def(Query *query, List *targetList, SubLink *cur_ma_sublink; List *ma_sublinks; + targetList = ExpandMergedSubscriptingRefEntries(targetList); + /* * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks * into a list. We expect them to appear, in ID order, in resjunk tlist diff --git a/src/backend/distributed/deparser/ruleutils_14.c b/src/backend/distributed/deparser/ruleutils_14.c index bde8e1b2355..e660884fb9c 100644 --- a/src/backend/distributed/deparser/ruleutils_14.c +++ b/src/backend/distributed/deparser/ruleutils_14.c @@ -3439,6 +3439,8 @@ get_update_query_targetlist_def(Query *query, List *targetList, SubLink *cur_ma_sublink; List *ma_sublinks; + targetList = ExpandMergedSubscriptingRefEntries(targetList); + /* * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks * into a list. We expect them to appear, in ID order, in resjunk tlist diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 60d3facc222..bff5a42431a 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -704,7 +704,12 @@ ModifyPartialQuerySupported(Query *queryTree, bool multiShardQuery, if (FindNodeMatchingCheckFunction((Node *) targetEntry->expr, NodeIsFieldStore)) { - /* DELETE cannot do field indirection already */ + /* + * DELETE cannot do field indirection already, so assert here. + * + * NB: See TargetEntryExprFindSubsRef if you decide removing + * this error check. + */ Assert(commandType == CMD_UPDATE || commandType == CMD_INSERT); return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, "inserting or modifying composite type fields is not " diff --git a/src/backend/distributed/utils/listutils.c b/src/backend/distributed/utils/listutils.c index 836a4bff649..de5f215a772 100644 --- a/src/backend/distributed/utils/listutils.c +++ b/src/backend/distributed/utils/listutils.c @@ -264,3 +264,20 @@ GenerateListFromElement(void *listElement, int listLength) return list; } + + +/* + * list_reverse returns a new list by reverting order of the elements within + * given list. + */ +List * +list_reverse(const List *list) +{ + List *newList = NIL; + for (int i = list_length(list) - 1; i >= 0; i--) + { + newList = lappend(newList, list_nth(list, i)); + } + + return newList; +} diff --git a/src/include/distributed/citus_ruleutils.h b/src/include/distributed/citus_ruleutils.h index 2396214151a..e8b712d6700 100644 --- a/src/include/distributed/citus_ruleutils.h +++ b/src/include/distributed/citus_ruleutils.h @@ -45,6 +45,7 @@ extern char * pg_get_indexclusterdef_string(Oid indexRelationId); extern bool contain_nextval_expression_walker(Node *node, void *context); extern char * pg_get_replica_identity_command(Oid tableRelationId); extern const char * RoleSpecString(RoleSpec *spec, bool withQuoteIdentifier); +extern List * ExpandMergedSubscriptingRefEntries(List *targetEntryList); /* Function declarations for version dependent PostgreSQL ruleutils functions */ extern void pg_get_query_def(Query *query, StringInfo buffer); diff --git a/src/include/distributed/listutils.h b/src/include/distributed/listutils.h index a808be43ff7..e257b7692f6 100644 --- a/src/include/distributed/listutils.h +++ b/src/include/distributed/listutils.h @@ -125,5 +125,6 @@ extern List * ListTake(List *pointerList, int size); extern void * safe_list_nth(const List *list, int index); extern List * GeneratePositiveIntSequenceList(int upTo); extern List * GenerateListFromElement(void *listElement, int listLength); +extern List * list_reverse(const List *list); #endif /* CITUS_LISTUTILS_H */ diff --git a/src/test/regress/expected/multi_deparse_shard_query.out b/src/test/regress/expected/multi_deparse_shard_query.out index 32a97eafa3b..6f3039826bd 100644 --- a/src/test/regress/expected/multi_deparse_shard_query.out +++ b/src/test/regress/expected/multi_deparse_shard_query.out @@ -3,10 +3,6 @@ -- SET citus.next_shard_id TO 13100000; SET citus.shard_replication_factor TO 1; -CREATE FUNCTION deparse_shard_query_test(text) - RETURNS VOID - AS 'citus' - LANGUAGE C STRICT; -- create the first table CREATE TABLE raw_events_1 (tenant_id bigint, diff --git a/src/test/regress/expected/multi_test_helpers.out b/src/test/regress/expected/multi_test_helpers.out index f46e6bcc8ca..5a47bf40984 100644 --- a/src/test/regress/expected/multi_test_helpers.out +++ b/src/test/regress/expected/multi_test_helpers.out @@ -124,3 +124,7 @@ BEGIN END LOOP; END; $$ LANGUAGE plpgsql; +CREATE OR REPLACE FUNCTION deparse_shard_query_test(text) + RETURNS VOID + AS 'citus' + LANGUAGE C STRICT; diff --git a/src/test/regress/expected/pg14.out b/src/test/regress/expected/pg14.out index cbdd70e0551..f315fccdc0d 100644 --- a/src/test/regress/expected/pg14.out +++ b/src/test/regress/expected/pg14.out @@ -1270,5 +1270,82 @@ SELECT count(*) FROM 1 (1 row) +CREATE TABLE jsonb_subscript_update (id INT, data JSONB); +SELECT create_distributed_table('jsonb_subscript_update', 'id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}'); +UPDATE jsonb_subscript_update +SET data['b'] = updated_vals.b::TEXT::jsonb, + data['c'] = updated_vals.c::TEXT::jsonb, + data['d'] = updated_vals.d::TEXT::jsonb +FROM ( + SELECT id, + data['a'] AS a, + data['a']::NUMERIC + 1 AS b, + data['a']::NUMERIC + 2 AS c, + data['a']::NUMERIC + 3 AS d + FROM jsonb_subscript_update +) updated_vals +WHERE jsonb_subscript_update.id = updated_vals.id; +SELECT * FROM jsonb_subscript_update ORDER BY 1,2; + id | data +--------------------------------------------------------------------- + 1 | {"a": 1, "b": 2, "c": 3, "d": 4} + 2 | {"a": 2, "b": 3, "c": 4, "d": 5} +(2 rows) + +TRUNCATE jsonb_subscript_update; +INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}'), (4, '{"a": 4, "b": 10}'); +ALTER TABLE jsonb_subscript_update ADD CONSTRAINT pkey PRIMARY KEY (id, data); +INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}') +ON CONFLICT (id, data) +DO UPDATE SET data['d']=(jsonb_subscript_update.data['a']::INT*100)::TEXT::JSONB, + data['b']=(jsonb_subscript_update.data['a']::INT*-100)::TEXT::JSONB; +SELECT * FROM jsonb_subscript_update ORDER BY 1,2; + id | data +--------------------------------------------------------------------- + 1 | {"a": 1, "b": -100, "d": 100} + 2 | {"a": 2, "b": -200, "d": 200} + 4 | {"a": 4, "b": 10} +(3 rows) + +CREATE TABLE nested_obj_update(id INT, data JSONB, text_col TEXT); +SELECT create_distributed_table('nested_obj_update', 'id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO nested_obj_update VALUES + (1, '{"a": [1,2,3], "b": [4,5,6], "c": [7,8,9], "d": [1,2,1,2]}', '4'), + (2, '{"a": [10,20,30], "b": [41,51,61], "c": [72,82,92], "d": [11,21,11,21]}', '6'); +UPDATE nested_obj_update +SET data['a'][0] = (updated_vals.b * 1)::TEXT::JSONB, + data['b'][2] = (updated_vals.c * 2)::TEXT::JSONB, + data['c'][0] = (updated_vals.d * 3)::TEXT::JSONB, + text_col = (nested_obj_update.id*1000)::TEXT, + data['a'][0] = (text_col::INT * data['a'][0]::INT)::TEXT::JSONB, + data['d'][6] = (nested_obj_update.id*1)::TEXT::JSONB, + data['d'][4] = (nested_obj_update.id*2)::TEXT::JSONB +FROM ( + SELECT id, + data['a'][0] AS a, + data['b'][0]::NUMERIC + 1 AS b, + data['c'][0]::NUMERIC + 2 AS c, + data['c'][1]::NUMERIC + 3 AS d + FROM nested_obj_update +) updated_vals +WHERE nested_obj_update.id = updated_vals.id; +SELECT * FROM nested_obj_update ORDER BY 1,2,3; + id | data | text_col +--------------------------------------------------------------------- + 1 | {"a": [4, 2, 3], "b": [4, 5, 18], "c": [33, 8, 9], "d": [1, 2, 1, 2, 2, null, 1]} | 1000 + 2 | {"a": [60, 20, 30], "b": [41, 51, 148], "c": [255, 82, 92], "d": [11, 21, 11, 21, 4, null, 2]} | 2000 +(2 rows) + set client_min_messages to error; drop schema pg14 cascade; diff --git a/src/test/regress/expected/subscripting_op.out b/src/test/regress/expected/subscripting_op.out new file mode 100644 index 00000000000..1a96108ae65 --- /dev/null +++ b/src/test/regress/expected/subscripting_op.out @@ -0,0 +1,120 @@ +\set VERBOSITY terse +SET citus.next_shard_id TO 1520000; +CREATE SCHEMA subscripting_op; +SET search_path TO subscripting_op; +CREATE TABLE arr_subs_update(id INT, arr INT[], text_col TEXT, int_col_1 INT, int_col_2 INT); +SELECT create_distributed_table('arr_subs_update', 'id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO arr_subs_update VALUES (1, '{1,2,3}', 'foo', 50, 60), (2, '{4,5,6}', 'bar', 60, 70); +-- test if we can properly expand target list entries when there are dropped columns +ALTER TABLE arr_subs_update DROP COLUMN int_col_1; +UPDATE arr_subs_update +SET arr[1] = updated_vals.b, + arr[3] = updated_vals.d, + int_col_2 = 400, + arr[2] = updated_vals.c +FROM ( + SELECT id, + arr[0] AS a, + arr[1]::NUMERIC + 1 AS b, + arr[2]::NUMERIC + 2 AS c, + arr[3]::NUMERIC + 3 AS d + FROM arr_subs_update +) updated_vals +WHERE arr_subs_update.id = updated_vals.id; +SELECT * FROM arr_subs_update ORDER BY 1,2,3,4; + id | arr | text_col | int_col_2 +--------------------------------------------------------------------- + 1 | {2,4,6} | foo | 400 + 2 | {5,7,9} | bar | 400 +(2 rows) + +TRUNCATE arr_subs_update; +INSERT INTO arr_subs_update VALUES (1, '{1,2,3}', 'foo', 60), (2, '{4,5,6}', 'bar', 70); +ALTER TABLE arr_subs_update ADD CONSTRAINT pkey PRIMARY KEY (id, arr); +INSERT INTO arr_subs_update VALUES (1, '{1,2,3}') +ON CONFLICT (id, arr) +DO UPDATE SET arr[0]=100, arr[1]=200, arr[5]=500; +SELECT * FROM arr_subs_update ORDER BY 1,2,3,4; + id | arr | text_col | int_col_2 +--------------------------------------------------------------------- + 1 | [0:5]={100,200,2,3,NULL,500} | foo | 60 + 2 | {4,5,6} | bar | 70 +(2 rows) + +CREATE DOMAIN single_int_dom AS int[] CHECK (VALUE[1] != 0); +CREATE DOMAIN dummy_dom AS single_int_dom CHECK (VALUE[2] != 5); +-- Citus doesn't propagate DOMAIN objects +SELECT run_command_on_workers( +$$ +CREATE DOMAIN subscripting_op.single_int_dom AS INT[] CHECK (VALUE[1] != 0); +CREATE DOMAIN subscripting_op.dummy_dom AS subscripting_op.single_int_dom CHECK (VALUE[2] != 5); +$$); + run_command_on_workers +--------------------------------------------------------------------- + (localhost,57637,t,"CREATE DOMAIN") + (localhost,57638,t,"CREATE DOMAIN") +(2 rows) + +CREATE TABLE dummy_dom_test (id int, dummy_dom_col dummy_dom); +SELECT create_distributed_table('dummy_dom_test', 'id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO dummy_dom_test VALUES (1, '{1,2,3}'), (2, '{6,7,8}'); +UPDATE dummy_dom_test +SET dummy_dom_col[2] = 50, + dummy_dom_col[1] = 60; +SELECT * FROM dummy_dom_test ORDER BY 1,2; + id | dummy_dom_col +--------------------------------------------------------------------- + 1 | {60,50,3} + 2 | {60,50,8} +(2 rows) + +CREATE TYPE two_ints as (if1 int, if2 int[]); +CREATE DOMAIN two_ints_dom AS two_ints CHECK ((VALUE).if1 > 0); +-- Citus doesn't propagate DOMAIN objects +SELECT run_command_on_workers( +$$ +CREATE DOMAIN subscripting_op.two_ints_dom AS subscripting_op.two_ints CHECK ((VALUE).if1 > 0); +$$); + run_command_on_workers +--------------------------------------------------------------------- + (localhost,57637,t,"CREATE DOMAIN") + (localhost,57638,t,"CREATE DOMAIN") +(2 rows) + +CREATE TABLE two_ints_dom_indirection_test (id int, two_ints_dom_col two_ints_dom); +SELECT create_distributed_table('two_ints_dom_indirection_test', 'id'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO two_ints_dom_indirection_test VALUES (1, '(5, "{1,2,3}")'), (2, '(50, "{10,20,30}")'); +-- Citus planner already doesn't allow doing field indirection (e.g.: +-- insert/update .) and we have an extra guard against +-- that in deparser for future implementations; so here we test that by using +-- deparse_shard_query_test() as well. +-- i) planner would throw an error +UPDATE two_ints_dom_indirection_test +SET two_ints_dom_col.if2[1] = 50, + two_ints_dom_col.if2[3] = 60; +ERROR: inserting or modifying composite type fields is not supported +-- ii) deparser would throw an error +SELECT public.deparse_shard_query_test( +$$ +UPDATE two_ints_dom_indirection_test +SET two_ints_dom_col.if2[1] = 50, + two_ints_dom_col.if2[3] = 60; +$$); +ERROR: unexpectedly got FieldStore object when generating shard query +SET client_min_messages TO WARNING; +DROP SCHEMA subscripting_op CASCADE; diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule index ee6a52da5ab..b80d2f722b4 100644 --- a/src/test/regress/multi_schedule +++ b/src/test/regress/multi_schedule @@ -78,7 +78,7 @@ test: custom_aggregate_support aggregate_support tdigest_aggregate_support test: multi_average_expression multi_working_columns multi_having_pushdown having_subquery test: multi_array_agg multi_limit_clause multi_orderby_limit_pushdown test: multi_jsonb_agg multi_jsonb_object_agg multi_json_agg multi_json_object_agg bool_agg ch_bench_having chbenchmark_all_queries expression_reference_join anonymous_columns -test: ch_bench_subquery_repartition +test: ch_bench_subquery_repartition subscripting_op test: multi_agg_type_conversion multi_count_type_conversion recursive_relation_planning_restriction_pushdown test: multi_partition_pruning single_hash_repartition_join test: multi_join_pruning multi_hash_pruning intermediate_result_pruning diff --git a/src/test/regress/sql/multi_deparse_shard_query.sql b/src/test/regress/sql/multi_deparse_shard_query.sql index f90f7eb4fba..2ddce7a3420 100644 --- a/src/test/regress/sql/multi_deparse_shard_query.sql +++ b/src/test/regress/sql/multi_deparse_shard_query.sql @@ -6,11 +6,6 @@ SET citus.next_shard_id TO 13100000; SET citus.shard_replication_factor TO 1; -CREATE FUNCTION deparse_shard_query_test(text) - RETURNS VOID - AS 'citus' - LANGUAGE C STRICT; - -- create the first table CREATE TABLE raw_events_1 (tenant_id bigint, diff --git a/src/test/regress/sql/multi_test_helpers.sql b/src/test/regress/sql/multi_test_helpers.sql index 85afac3e0fa..e89eefdc4ae 100644 --- a/src/test/regress/sql/multi_test_helpers.sql +++ b/src/test/regress/sql/multi_test_helpers.sql @@ -135,3 +135,8 @@ BEGIN END; $$ LANGUAGE plpgsql; +CREATE OR REPLACE FUNCTION deparse_shard_query_test(text) + RETURNS VOID + AS 'citus' + LANGUAGE C STRICT; + diff --git a/src/test/regress/sql/pg14.sql b/src/test/regress/sql/pg14.sql index 8bc422f5de7..1e625f5cc48 100644 --- a/src/test/regress/sql/pg14.sql +++ b/src/test/regress/sql/pg14.sql @@ -640,5 +640,64 @@ SELECT count(*) FROM UNION select row(pg_proc.pronargs, pg_proc.proargtypes, pg_proc.prosrc, pg_proc.proowner)::text from pg_proc where proname = 'proc_with_out_param') as test; +CREATE TABLE jsonb_subscript_update (id INT, data JSONB); +SELECT create_distributed_table('jsonb_subscript_update', 'id'); + +INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}'); + +UPDATE jsonb_subscript_update +SET data['b'] = updated_vals.b::TEXT::jsonb, + data['c'] = updated_vals.c::TEXT::jsonb, + data['d'] = updated_vals.d::TEXT::jsonb +FROM ( + SELECT id, + data['a'] AS a, + data['a']::NUMERIC + 1 AS b, + data['a']::NUMERIC + 2 AS c, + data['a']::NUMERIC + 3 AS d + FROM jsonb_subscript_update +) updated_vals +WHERE jsonb_subscript_update.id = updated_vals.id; + +SELECT * FROM jsonb_subscript_update ORDER BY 1,2; + +TRUNCATE jsonb_subscript_update; +INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}'), (4, '{"a": 4, "b": 10}'); + +ALTER TABLE jsonb_subscript_update ADD CONSTRAINT pkey PRIMARY KEY (id, data); + +INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}') +ON CONFLICT (id, data) +DO UPDATE SET data['d']=(jsonb_subscript_update.data['a']::INT*100)::TEXT::JSONB, + data['b']=(jsonb_subscript_update.data['a']::INT*-100)::TEXT::JSONB; + +SELECT * FROM jsonb_subscript_update ORDER BY 1,2; + +CREATE TABLE nested_obj_update(id INT, data JSONB, text_col TEXT); +SELECT create_distributed_table('nested_obj_update', 'id'); +INSERT INTO nested_obj_update VALUES + (1, '{"a": [1,2,3], "b": [4,5,6], "c": [7,8,9], "d": [1,2,1,2]}', '4'), + (2, '{"a": [10,20,30], "b": [41,51,61], "c": [72,82,92], "d": [11,21,11,21]}', '6'); + +UPDATE nested_obj_update +SET data['a'][0] = (updated_vals.b * 1)::TEXT::JSONB, + data['b'][2] = (updated_vals.c * 2)::TEXT::JSONB, + data['c'][0] = (updated_vals.d * 3)::TEXT::JSONB, + text_col = (nested_obj_update.id*1000)::TEXT, + data['a'][0] = (text_col::INT * data['a'][0]::INT)::TEXT::JSONB, + data['d'][6] = (nested_obj_update.id*1)::TEXT::JSONB, + data['d'][4] = (nested_obj_update.id*2)::TEXT::JSONB +FROM ( + SELECT id, + data['a'][0] AS a, + data['b'][0]::NUMERIC + 1 AS b, + data['c'][0]::NUMERIC + 2 AS c, + data['c'][1]::NUMERIC + 3 AS d + FROM nested_obj_update +) updated_vals +WHERE nested_obj_update.id = updated_vals.id; + +SELECT * FROM nested_obj_update ORDER BY 1,2,3; + set client_min_messages to error; drop schema pg14 cascade; diff --git a/src/test/regress/sql/subscripting_op.sql b/src/test/regress/sql/subscripting_op.sql new file mode 100644 index 00000000000..3e32a60d270 --- /dev/null +++ b/src/test/regress/sql/subscripting_op.sql @@ -0,0 +1,97 @@ +\set VERBOSITY terse + +SET citus.next_shard_id TO 1520000; + +CREATE SCHEMA subscripting_op; +SET search_path TO subscripting_op; + +CREATE TABLE arr_subs_update(id INT, arr INT[], text_col TEXT, int_col_1 INT, int_col_2 INT); +SELECT create_distributed_table('arr_subs_update', 'id'); +INSERT INTO arr_subs_update VALUES (1, '{1,2,3}', 'foo', 50, 60), (2, '{4,5,6}', 'bar', 60, 70); + +-- test if we can properly expand target list entries when there are dropped columns +ALTER TABLE arr_subs_update DROP COLUMN int_col_1; + +UPDATE arr_subs_update +SET arr[1] = updated_vals.b, + arr[3] = updated_vals.d, + int_col_2 = 400, + arr[2] = updated_vals.c +FROM ( + SELECT id, + arr[0] AS a, + arr[1]::NUMERIC + 1 AS b, + arr[2]::NUMERIC + 2 AS c, + arr[3]::NUMERIC + 3 AS d + FROM arr_subs_update +) updated_vals +WHERE arr_subs_update.id = updated_vals.id; + +SELECT * FROM arr_subs_update ORDER BY 1,2,3,4; + +TRUNCATE arr_subs_update; +INSERT INTO arr_subs_update VALUES (1, '{1,2,3}', 'foo', 60), (2, '{4,5,6}', 'bar', 70); + +ALTER TABLE arr_subs_update ADD CONSTRAINT pkey PRIMARY KEY (id, arr); + +INSERT INTO arr_subs_update VALUES (1, '{1,2,3}') +ON CONFLICT (id, arr) +DO UPDATE SET arr[0]=100, arr[1]=200, arr[5]=500; + +SELECT * FROM arr_subs_update ORDER BY 1,2,3,4; + +CREATE DOMAIN single_int_dom AS int[] CHECK (VALUE[1] != 0); +CREATE DOMAIN dummy_dom AS single_int_dom CHECK (VALUE[2] != 5); + +-- Citus doesn't propagate DOMAIN objects +SELECT run_command_on_workers( +$$ +CREATE DOMAIN subscripting_op.single_int_dom AS INT[] CHECK (VALUE[1] != 0); +CREATE DOMAIN subscripting_op.dummy_dom AS subscripting_op.single_int_dom CHECK (VALUE[2] != 5); +$$); + +CREATE TABLE dummy_dom_test (id int, dummy_dom_col dummy_dom); +SELECT create_distributed_table('dummy_dom_test', 'id'); + +INSERT INTO dummy_dom_test VALUES (1, '{1,2,3}'), (2, '{6,7,8}'); + +UPDATE dummy_dom_test +SET dummy_dom_col[2] = 50, + dummy_dom_col[1] = 60; + +SELECT * FROM dummy_dom_test ORDER BY 1,2; + +CREATE TYPE two_ints as (if1 int, if2 int[]); +CREATE DOMAIN two_ints_dom AS two_ints CHECK ((VALUE).if1 > 0); + +-- Citus doesn't propagate DOMAIN objects +SELECT run_command_on_workers( +$$ +CREATE DOMAIN subscripting_op.two_ints_dom AS subscripting_op.two_ints CHECK ((VALUE).if1 > 0); +$$); + +CREATE TABLE two_ints_dom_indirection_test (id int, two_ints_dom_col two_ints_dom); +SELECT create_distributed_table('two_ints_dom_indirection_test', 'id'); + +INSERT INTO two_ints_dom_indirection_test VALUES (1, '(5, "{1,2,3}")'), (2, '(50, "{10,20,30}")'); + +-- Citus planner already doesn't allow doing field indirection (e.g.: +-- insert/update .) and we have an extra guard against +-- that in deparser for future implementations; so here we test that by using +-- deparse_shard_query_test() as well. + +-- i) planner would throw an error +UPDATE two_ints_dom_indirection_test +SET two_ints_dom_col.if2[1] = 50, + two_ints_dom_col.if2[3] = 60; + +-- ii) deparser would throw an error +SELECT public.deparse_shard_query_test( +$$ +UPDATE two_ints_dom_indirection_test +SET two_ints_dom_col.if2[1] = 50, + two_ints_dom_col.if2[3] = 60; +$$); + +SET client_min_messages TO WARNING; +DROP SCHEMA subscripting_op CASCADE;