Skip to content

Commit 879066a

Browse files
snowersnower
committed
Optimization: You can use <pk> in the as name to specify the outputer primary key
1 parent 713101e commit 879066a

3 files changed

Lines changed: 59 additions & 4 deletions

File tree

‎examples/insert_types/update.sql‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,4 +26,13 @@ update `cdata` set id=id, name='胡萝卜' where id > 5 or name='花生';
2626
select `id`, `name`, `create_time` from `cdata`;
2727

2828
update `cdata` set id=id, name='青椒' where (id + 1) > 4 and concat(name, '11')='花菜11';
29+
select `id`, `name`, `create_time` from `cdata`;
30+
31+
insert into `cdata<U>` select `create_time`, `id<pk>`, `name` from `ndata` where `id`=1;
32+
select `id`, `name`, `create_time` from `cdata`;
33+
34+
insert into `cdata<U>` select `name`, 1 as `id<pk>` from `ndata` where `id`=2;
35+
select `id`, `name`, `create_time` from `cdata`;
36+
37+
insert into `cdata<U>` select '豆角' as `name`, 1 as `id<pk>`;
2938
select `id`, `name`, `create_time` from `cdata`;

‎syncanysql/compiler.py‎

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -390,7 +390,8 @@ def parse(union_expression):
390390

391391
def compile_select(self, expression, config, arguments):
392392
primary_table = {"db": None, "name": None, "table_name": None, "table_alias": None, "seted_primary_keys": False,
393-
"loader_primary_keys": [], "outputer_primary_keys": [], "columns": {}, "subquery": None, "select_columns": {}}
393+
"loader_primary_keys": [], "outputer_primary_keys": [], "alias_primary_keys": [],
394+
"columns": {}, "subquery": None, "select_columns": {}}
394395

395396
from_expression = expression.args.get("from")
396397
if not from_expression:
@@ -460,9 +461,11 @@ def compile_select(self, expression, config, arguments):
460461
if self.is_column(select_expression, config, arguments):
461462
column_expression = select_expression
462463
elif isinstance(select_expression, sqlglot_expressions.Alias):
463-
column_alias = select_expression.args["alias"].name if select_expression.args.get("alias") else None
464-
if column_alias and column_alias in self.mapping:
465-
column_alias = self.mapping[column_alias]
464+
column_alias = self.parse_column_alias(select_expression, config, arguments)
465+
if column_alias is not None:
466+
if column_alias["typing_options"] and "pk" in column_alias["typing_options"]:
467+
primary_table["alias_primary_keys"].append(column_alias["column_alias"])
468+
column_alias = column_alias["column_alias"]
466469
if self.is_const(select_expression.args["this"], config, arguments):
467470
const_info = self.parse_const(select_expression.args["this"], config, arguments)
468471
config["schema"][column_alias] = self.compile_const(select_expression.args["this"], config, arguments, const_info)
@@ -570,6 +573,8 @@ def compile_select(self, expression, config, arguments):
570573
if name in config["aggregate"]["schema"]:
571574
continue
572575
config["aggregate"]["distinct_keys"].append(copy.deepcopy(column))
576+
if primary_table["alias_primary_keys"]:
577+
primary_table["outputer_primary_keys"] = primary_table["alias_primary_keys"]
573578
if not from_expression and not primary_table["outputer_primary_keys"] and isinstance(config["schema"], dict):
574579
for column_alias in config["schema"]:
575580
if not column_alias.isidentifier():
@@ -2762,6 +2767,26 @@ def parse_const(self, expression, config, arguments):
27622767
"expression": expression,
27632768
}
27642769

2770+
def parse_column_alias(self, expression, config, arguments):
2771+
if not expression.args.get("alias"):
2772+
return None
2773+
origin_alias = expression.args["alias"].name
2774+
if origin_alias and origin_alias in self.mapping:
2775+
column_alias = self.mapping[origin_alias]
2776+
else:
2777+
column_alias = origin_alias
2778+
try:
2779+
start_index, end_index = column_alias.index("<"), column_alias.rindex(">")
2780+
typing_options = column_alias[start_index + 1: end_index].split(",")
2781+
column_alias = column_alias[:start_index]
2782+
except ValueError:
2783+
typing_options = []
2784+
return {
2785+
"origin_alias": origin_alias,
2786+
"column_alias": column_alias,
2787+
"typing_options": typing_options,
2788+
}
2789+
27652790
def parse_cast_typing_filter(self, expression, config, arguments):
27662791
to_type = expression.args["to"].args["this"]
27672792
if to_type in sqlglot_expressions.DataType.FLOAT_TYPES:

‎tests/test_example_insert_types.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,24 @@ def test_update(self):
103103
{'create_time': '2023-03-12 10:12:34', 'id': 5, 'name': '玉米'},
104104
{'create_time': '2023-03-12 10:12:34', 'id': 6, 'name': '胡萝卜'},
105105
{'create_time': '2023-03-12 10:12:34', 'id': 8, 'name': '胡萝卜'}], "data error")
106+
107+
self.assert_data(32, [{'create_time': '2023-03-12 10:12:34', 'id': 1, 'name': '白菜'},
108+
{'create_time': '2023-03-12 10:12:34', 'id': 2, 'name': '玉米'},
109+
{'create_time': '2023-03-12 10:12:34', 'id': 4, 'name': '青椒'},
110+
{'create_time': '2023-03-12 10:12:34', 'id': 5, 'name': '玉米'},
111+
{'create_time': '2023-03-12 10:12:34', 'id': 6, 'name': '胡萝卜'},
112+
{'create_time': '2023-03-12 10:12:34', 'id': 8, 'name': '胡萝卜'}], "data error")
113+
114+
self.assert_data(35, [{'create_time': '2023-03-12 10:12:34', 'id': 1, 'name': '青菜'},
115+
{'create_time': '2023-03-12 10:12:34', 'id': 2, 'name': '玉米'},
116+
{'create_time': '2023-03-12 10:12:34', 'id': 4, 'name': '青椒'},
117+
{'create_time': '2023-03-12 10:12:34', 'id': 5, 'name': '玉米'},
118+
{'create_time': '2023-03-12 10:12:34', 'id': 6, 'name': '胡萝卜'},
119+
{'create_time': '2023-03-12 10:12:34', 'id': 8, 'name': '胡萝卜'}], "data error")
120+
121+
self.assert_data(38, [{'create_time': '2023-03-12 10:12:34', 'id': 1, 'name': '豆角'},
122+
{'create_time': '2023-03-12 10:12:34', 'id': 2, 'name': '玉米'},
123+
{'create_time': '2023-03-12 10:12:34', 'id': 4, 'name': '青椒'},
124+
{'create_time': '2023-03-12 10:12:34', 'id': 5, 'name': '玉米'},
125+
{'create_time': '2023-03-12 10:12:34', 'id': 6, 'name': '胡萝卜'},
126+
{'create_time': '2023-03-12 10:12:34', 'id': 8, 'name': '胡萝卜'}], "data error")

0 commit comments

Comments
 (0)