Skip to content

Commit 2659780

Browse files
authored
support online basic cases (#7)
Signed-off-by: shanhaikang.shk <shanhaikang.shk@oceanbase.com>
1 parent 4959ed0 commit 2659780

File tree

2 files changed

+130
-14
lines changed

2 files changed

+130
-14
lines changed

pyobvector/client/ob_vec_json_table_client.py

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
from sqlalchemy import Column, Integer, String, JSON, Engine, select, text, func, CursorResult
77
from sqlalchemy.dialects.mysql import TINYINT
88
from sqlalchemy.orm import declarative_base, sessionmaker, Session
9-
from sqlglot import parse_one, exp, Expression
9+
from sqlglot import parse_one, exp, Expression, to_identifier
10+
from sqlglot.expressions import Concat
11+
1012

1113
from .ob_vec_client import ObVecClient
1214
from ..json_table import (
@@ -35,7 +37,7 @@ class ObVecJsonTableClient(ObVecClient):
3537
class JsonTableMetaTBL(Base):
3638
__tablename__ = JSON_TABLE_META_TABLE_NAME
3739

38-
user_id = Column(Integer, primary_key=True)
40+
user_id = Column(String(128), primary_key=True, autoincrement=False)
3941
jtable_name = Column(String(512), primary_key=True)
4042
jcol_id = Column(Integer, primary_key=True)
4143
jcol_name = Column(String(512), primary_key=True)
@@ -47,13 +49,13 @@ class JsonTableMetaTBL(Base):
4749
class JsonTableDataTBL(Base):
4850
__tablename__ = JSON_TABLE_DATA_TABLE_NAME
4951

50-
user_id = Column(Integer, primary_key=True)
52+
user_id = Column(String(128), primary_key=True, autoincrement=False)
5153
jtable_name = Column(String(512), primary_key=True)
5254
jdata_id = Column(Integer, primary_key=True, autoincrement=True, nullable=False)
5355
jdata = Column(JSON)
5456

5557
class JsonTableMetadata:
56-
def __init__(self, user_id: int):
58+
def __init__(self, user_id: str):
5759
self.user_id = user_id
5860
self.meta_cache: Dict[str, List] = {}
5961

@@ -78,7 +80,7 @@ def _parse_col_type(cls, col_type: str):
7880
if col_type == 'DECIMAL':
7981
factory = JsonTableDecimalFactory(10, 0)
8082
else:
81-
decimal_pattern = r'DECIMAL\((\d+),\s*(\d+)\)'
83+
decimal_pattern = r'DECIMAL\s*\((\d+),\s*(\d+)\)'
8284
decimal_matches = re.findall(decimal_pattern, col_type)
8385
x, y = decimal_matches[0]
8486
factory = JsonTableDecimalFactory(int(x), int(y))
@@ -119,7 +121,7 @@ def reflect(self, engine: Engine):
119121

120122
def __init__(
121123
self,
122-
user_id: int,
124+
user_id: str,
123125
uri: str = "127.0.0.1:2881",
124126
user: str = "root@test",
125127
password: str = "",
@@ -238,7 +240,9 @@ def _handle_create_json_table(self, ast: Expression):
238240
elif isinstance(cons.kind, exp.NotNullColumnConstraint):
239241
col_nullable = False
240242
else:
241-
raise ValueError(f"{cons.kind} constriaint is not supported.")
243+
pass
244+
# raise ValueError(f"{cons.kind} constriaint is not supported.")
245+
# TODO support json index
242246

243247
if col_has_default and (col_default_val is not None):
244248
# check default value is valid
@@ -689,6 +693,20 @@ def _handle_jtable_dml_update(self, ast: Expression):
689693
if not self._check_col_exists(table_name, col_name):
690694
raise ValueError(f"Column {col_name} does not exists")
691695
col_expr = expr.expression
696+
new_node = parse_one(f"JSON_VALUE({JSON_TABLE_DATA_TABLE_NAME}.jdata, '$.{col_name}')")
697+
for column in col_expr.find_all(exp.Column):
698+
parent_node = column.parent
699+
if isinstance(parent_node.args[column.arg_key], list):
700+
new_expr_list = []
701+
for node in parent_node.args[column.arg_key]:
702+
if isinstance(node, exp.Column):
703+
new_expr_list.append(new_node)
704+
else:
705+
new_expr_list.append(node)
706+
parent_node.args[column.arg_key] = new_expr_list
707+
elif isinstance(parent_node.args[column.arg_key], exp.Column):
708+
parent_node.args[column.arg_key] = new_node
709+
logger.info(str(col_expr))
692710
path_settings.append(f"'$.{col_name}', {str(col_expr)}")
693711

694712
where_clause = None
@@ -700,7 +718,7 @@ def _handle_jtable_dml_update(self, ast: Expression):
700718
column.parent.args['this'] = parse_one(
701719
f"JSON_VALUE({JSON_TABLE_DATA_TABLE_NAME}.jdata, '$.{where_col_name}')"
702720
)
703-
where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = {self.user_id} AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
721+
where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{self.user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
704722

705723
if where_clause:
706724
update_sql = f"UPDATE {JSON_TABLE_DATA_TABLE_NAME} SET jdata = JSON_REPLACE({JSON_TABLE_DATA_TABLE_NAME}.jdata, {', '.join(path_settings)}) WHERE {where_clause}"
@@ -724,7 +742,7 @@ def _handle_jtable_dml_delete(self, ast: Expression):
724742
column.parent.args['this'] = parse_one(
725743
f"JSON_VALUE({JSON_TABLE_DATA_TABLE_NAME}.jdata, '$.{where_col_name}')"
726744
)
727-
where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = {self.user_id} AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
745+
where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{self.user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
728746

729747
if where_clause:
730748
delete_sql = f"DELETE FROM {JSON_TABLE_DATA_TABLE_NAME} WHERE {where_clause}"
@@ -746,7 +764,7 @@ def _handle_jtable_dml_select(self, ast: Expression):
746764
if not self._check_table_exists(table_name):
747765
raise ValueError(f"Table {table_name} does not exists")
748766

749-
ast.args['from'].args['this'].args['this'].args['this'] = JSON_TABLE_DATA_TABLE_NAME
767+
ast.args['from'].args['this'].args['this'] = to_identifier(name=JSON_TABLE_DATA_TABLE_NAME, quoted=False)
750768

751769
col_meta = self.jmetadata.meta_cache[table_name]
752770
json_table_meta_str = []
@@ -794,7 +812,7 @@ def _handle_jtable_dml_select(self, ast: Expression):
794812
else:
795813
ast.args['joins'] = [join_node]
796814

797-
extra_filter_str = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = {self.user_id} AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}'"
815+
extra_filter_str = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{self.user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}'"
798816
if 'where' in ast.args.keys():
799817
filter_str = str(ast.args['where'].args['this'])
800818
new_filter_str = f"{extra_filter_str} AND ({filter_str})"

tests/test_json_table.py

Lines changed: 101 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@ def get_all_rows(res):
2424

2525
class ObVecJsonTableTest(unittest.TestCase):
2626
def setUp(self) -> None:
27-
self.root_client = ObVecJsonTableClient(user_id=0)
28-
self.client = ObVecJsonTableClient(user_id=1, user="jtuser@test")
27+
self.root_client = ObVecJsonTableClient(user_id='0')
28+
self.client = ObVecJsonTableClient(user_id='e5a69db04c5ea54adf324907d4b8f364', user="jtuser@test")
2929

3030
def test_create_and_alter_jtable(self):
3131
self.root_client._reset()
@@ -34,7 +34,7 @@ def test_create_and_alter_jtable(self):
3434
self.client.perform_json_table_sql(
3535
"create table `t2` (c1 int NOT NULL DEFAULT 10, c2 varchar(30) DEFAULT 'ca', c3 varchar not null, c4 decimal(10, 2));"
3636
)
37-
tmp_client = ObVecJsonTableClient(user_id=1)
37+
tmp_client = ObVecJsonTableClient(user_id='e5a69db04c5ea54adf324907d4b8f364')
3838
self.assertEqual(sub_dict(tmp_client.jmetadata.meta_cache['t2'], keys_to_check),
3939
[
4040
{'jcol_id': 16, 'jcol_name': 'c1', 'jcol_type': 'INT', 'jcol_nullable': False, 'jcol_has_default': True, 'jcol_default': '10'},
@@ -364,3 +364,101 @@ def test_col_name_conflict(self):
364364
(2, 'bob'),
365365
]
366366
)
367+
368+
def test_timestamp_datatype(self):
369+
self.root_client._reset()
370+
self.client.refresh_metadata()
371+
self.client.perform_json_table_sql(
372+
"create table `t1` (c1 int DEFAULT NULL, c2 TIMESTAMP);"
373+
)
374+
375+
self.client.perform_json_table_sql(
376+
"insert into t1 values (1, CURRENT_DATE - INTERVAL '1' MONTH);"
377+
)
378+
379+
self.client.perform_json_table_sql(
380+
"select * from t1"
381+
)
382+
383+
def test_online_cases(self):
384+
self.root_client._reset()
385+
self.client.refresh_metadata()
386+
keys_to_check = ['jcol_id', 'jcol_name', 'jcol_type', 'jcol_nullable', 'jcol_has_default', 'jcol_default']
387+
self.client.perform_json_table_sql(
388+
"CREATE TABLE `table_unit_test`(id INT, field0 VARCHAR(1024), field1 INT, field2 DECIMAL(10,2), field3 timestamp NOT NULL default CURRENT_TIMESTAMP, field4 varchar(1024));"
389+
)
390+
logger.info(sub_dict(self.client.jmetadata.meta_cache['table_unit_test'], keys_to_check))
391+
self.assertEqual(sub_dict(self.client.jmetadata.meta_cache['table_unit_test'], keys_to_check),
392+
[
393+
{'jcol_id': 16, 'jcol_name': 'id', 'jcol_type': 'INT', 'jcol_nullable': True, 'jcol_has_default': False, 'jcol_default': None},
394+
{'jcol_id': 17, 'jcol_name': 'field0', 'jcol_type': 'VARCHAR(1024)', 'jcol_nullable': True, 'jcol_has_default': False, 'jcol_default': None},
395+
{'jcol_id': 18, 'jcol_name': 'field1', 'jcol_type': 'INT', 'jcol_nullable': True, 'jcol_has_default': False, 'jcol_default': None},
396+
{'jcol_id': 19, 'jcol_name': 'field2', 'jcol_type': 'DECIMAL(10,2)', 'jcol_nullable': True, 'jcol_has_default': False, 'jcol_default': None},
397+
{'jcol_id': 20, 'jcol_name': 'field3', 'jcol_type': 'TIMESTAMP', 'jcol_nullable': False, 'jcol_has_default': True, 'jcol_default': 'CURRENT_TIMESTAMP()'},
398+
{'jcol_id': 21, 'jcol_name': 'field4', 'jcol_type': 'VARCHAR(1024)', 'jcol_nullable': True, 'jcol_has_default': False, 'jcol_default': None},
399+
]
400+
)
401+
self.client.perform_json_table_sql(
402+
"INSERT INTO `table_unit_test` (id, field0, field1, field2, field4) VALUES (1, '汽车保养', 1, 1000.00, '4S 店')"
403+
)
404+
self.client.perform_json_table_sql(
405+
"INSERT INTO `table_unit_test` (id, field0, field1, field2, field3, field4) VALUES (2, '奶茶', 2, 15.00, CURRENT_TIMESTAMP(), '商场'), (3, '书', 2, 40.00, NOW() - INTERVAL '1' DAY, '商场'), (4, '手机', 2, 6000.00, '2025-03-09', '商场')"
406+
)
407+
408+
res = self.client.perform_json_table_sql(
409+
"SELECT field0 AS 消费内容, field1 AS 消费类型, field2 AS 消费金额, field4 AS 消费地点 FROM `table_unit_test` LIMIT 2"
410+
)
411+
self.assertEqual(
412+
get_all_rows(res),
413+
[
414+
('汽车保养', 1, Decimal('1000.00'), '4S 店'),
415+
('奶茶', 2, Decimal('15.00'), '商场')
416+
]
417+
)
418+
419+
res = self.client.perform_json_table_sql(
420+
"SELECT field2 FROM `table_unit_test` WHERE field0 like '%汽车%' AND field4 like '%店%' ORDER BY field2 DESC LIMIT 2"
421+
)
422+
self.assertEqual(
423+
get_all_rows(res),
424+
[(Decimal('1000.00'),)]
425+
)
426+
427+
res = self.client.perform_json_table_sql(
428+
"SELECT field0, field1, field2 FROM `table_unit_test` WHERE DATE(field3)='2025-03-09' ORDER BY field2 DESC LIMIT 2"
429+
)
430+
# logger.info(get_all_rows(res))
431+
self.assertEqual(
432+
get_all_rows(res),
433+
[('手机', 2, Decimal('6000.00'))]
434+
)
435+
436+
res = self.client.perform_json_table_sql(
437+
"SELECT field0, field1, field2, field3 FROM `table_unit_test` WHERE field4 like '%商场%' AND field3 <= CAST('2025-03-10' AS DATE) LIMIT 20"
438+
)
439+
self.assertEqual(
440+
get_all_rows(res),
441+
[('手机', 2, Decimal('6000.00'), datetime.datetime(2025, 3, 9, 0, 0))]
442+
)
443+
444+
self.client.perform_json_table_sql(
445+
"UPDATE `table_unit_test` SET field3 = '2025-03-14' WHERE field0 = '汽车保养'"
446+
)
447+
res = self.client.perform_json_table_sql(
448+
"SELECT field0, field1, field2, field3 FROM `table_unit_test` WHERE DATE(field3) = '2025-03-14' AND field0 = '汽车保养'"
449+
)
450+
self.assertEqual(
451+
get_all_rows(res),
452+
[('汽车保养', 1, Decimal('1000.00'), datetime.datetime(2025, 3, 14, 0, 0))]
453+
)
454+
455+
self.client.perform_json_table_sql(
456+
"UPDATE `table_unit_test` SET field0 = CONCAT(field0, '代办') WHERE DATE(field3) = '2025-03-14'"
457+
)
458+
res = self.client.perform_json_table_sql(
459+
"SELECT field0, field1, field2, field3 FROM `table_unit_test` WHERE DATE(field3) = '2025-03-14'"
460+
)
461+
self.assertEqual(
462+
get_all_rows(res),
463+
[('汽车保养代办', 1, Decimal('1000.00'), datetime.datetime(2025, 3, 14, 0, 0))]
464+
)

0 commit comments

Comments
 (0)