pyobvector 0.2.7__py3-none-any.whl → 0.2.8__py3-none-any.whl

pyobvector/client/ob_vec_json_table_client.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import re
-from typing import Dict, List, Optional, Any
+from typing import Dict, List, Optional, Any, Union
 
 from sqlalchemy import Column, Integer, String, JSON, Engine, select, text, func, CursorResult
 from sqlalchemy.dialects.mysql import TINYINT
@@ -163,7 +163,12 @@ class ObVecJsonTableClient(ObVecClient):
     def refresh_metadata(self) -> None:
         self.jmetadata.reflect(self.engine)
 
-    def perform_json_table_sql(self, sql: str, select_with_data_id: bool = False) -> Optional[CursorResult]:
+    def perform_json_table_sql(
+        self,
+        sql: str,
+        select_with_data_id: bool = False,
+        opt_user_id: Optional[str] = None,
+    ) -> Union[Optional[CursorResult], int]:
         """Perform common SQL that operates on JSON Table."""
         ast = parse_one(sql, dialect="oceanbase")
         if isinstance(ast, exp.Create):
@@ -176,16 +181,16 @@ class ObVecJsonTableClient(ObVecClient):
             self._handle_alter_json_table(ast)
             return None
         elif isinstance(ast, exp.Insert):
-            self._handle_jtable_dml_insert(ast)
-            return None
+            row_count = self._handle_jtable_dml_insert(ast, opt_user_id)
+            return row_count
         elif isinstance(ast, exp.Update):
-            self._handle_jtable_dml_update(ast)
-            return None
+            row_count = self._handle_jtable_dml_update(ast, opt_user_id)
+            return row_count
         elif isinstance(ast, exp.Delete):
-            self._handle_jtable_dml_delete(ast)
-            return None
+            row_count = self._handle_jtable_dml_delete(ast, opt_user_id)
+            return row_count
         elif isinstance(ast, exp.Select):
-            return self._handle_jtable_dml_select(ast, select_with_data_id)
+            return self._handle_jtable_dml_select(ast, select_with_data_id, opt_user_id)
         else:
             raise ValueError(f"{type(ast)} not supported")
         
@@ -642,8 +647,10 @@ class ObVecJsonTableClient(ObVecClient):
         finally:
             session.close()
 
-    def _handle_jtable_dml_insert(self, ast: Expression):
-        if self.user_id is None:
+    def _handle_jtable_dml_insert(self, ast: Expression, opt_user_id: Optional[str] = None):
+        real_user_id = opt_user_id or self.user_id
+
+        if real_user_id is None:
             raise ValueError(f"inserting is disabled when user_id is None")
 
         if isinstance(ast.this, exp.Schema):
@@ -673,6 +680,7 @@ class ObVecJsonTableClient(ObVecClient):
             raise ValueError(f"Invalid ast type {ast.this}")
 
         session = self.session()
+        n_new_records = 0
         for tuple in ast.expression.expressions:
             expr_list = tuple.expressions
             if len(expr_list) != len(insert_col_names):
@@ -691,21 +699,26 @@ class ObVecJsonTableClient(ObVecClient):
             logger.debug(f"================= [INSERT] =============== {kv}")
 
             session.add(ObVecJsonTableClient.JsonTableDataTBL(
-                user_id = self.user_id,
+                user_id = real_user_id,
                 admin_id = self.admin_id,
                 jtable_name = table_name,
                 jdata = kv,
             ))
+            n_new_records += 1
         
         try:
             session.commit()
         except Exception as e:
             session.rollback()
             logger.error(f"Error occurred: {e}")
+            n_new_records = 0
         finally:
             session.close()
+            return n_new_records
+
+    def _handle_jtable_dml_update(self, ast: Expression, opt_user_id: Optional[str] = None):
+        real_user_id = opt_user_id or self.user_id
 
-    def _handle_jtable_dml_update(self, ast: Expression):
         table_name = ast.this.this.this
         if not self._check_table_exists(table_name):
             raise ValueError(f"Table {table_name} does not exists")
@@ -733,7 +746,7 @@ class ObVecJsonTableClient(ObVecClient):
             path_settings.append(f"'$.{col_name}', {str(col_expr)}")
 
         where_clause = None        
-        if 'where' in ast.args.keys():
+        if 'where' in ast.args.keys() and ast.args['where']:
             for column in ast.args['where'].find_all(exp.Column):
                 where_col_name = column.this.this
                 if not self._check_col_exists(table_name, where_col_name):
@@ -741,8 +754,8 @@ class ObVecJsonTableClient(ObVecClient):
                 column.parent.args['this'] = parse_one(
                     f"JSON_VALUE({JSON_TABLE_DATA_TABLE_NAME}.jdata, '$.{where_col_name}')"
                 )
-            if self.user_id:
-                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)})"
+            if real_user_id:
+                where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{real_user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
             else:
                 where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
         
@@ -752,15 +765,18 @@ class ObVecJsonTableClient(ObVecClient):
             update_sql = f"UPDATE {JSON_TABLE_DATA_TABLE_NAME} SET jdata = JSON_REPLACE({JSON_TABLE_DATA_TABLE_NAME}.jdata, {', '.join(path_settings)})"
 
         logger.debug(f"===================== do update: {update_sql}")
-        self.perform_raw_text_sql(update_sql)
+        res = self.perform_raw_text_sql(update_sql)
+        return res.rowcount
+
+    def _handle_jtable_dml_delete(self, ast: Expression, opt_user_id: Optional[str] = None):
+        real_user_id = opt_user_id or self.user_id
 
-    def _handle_jtable_dml_delete(self, ast: Expression):
         table_name = ast.this.this.this
         if not self._check_table_exists(table_name):
             raise ValueError(f"Table {table_name} does not exists")
         
         where_clause = None
-        if 'where' in ast.args.keys():
+        if 'where' in ast.args.keys() and ast.args['where']:
             for column in ast.args['where'].find_all(exp.Column):
                 where_col_name = column.this.this
                 if not self._check_col_exists(table_name, where_col_name):
@@ -768,8 +784,8 @@ class ObVecJsonTableClient(ObVecClient):
                 column.parent.args['this'] = parse_one(
                     f"JSON_VALUE({JSON_TABLE_DATA_TABLE_NAME}.jdata, '$.{where_col_name}')"
                 )
-            if self.user_id:
-                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)})"
+            if real_user_id:
+                where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{real_user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
             else:
                 where_clause = f"{JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}' AND ({str(ast.args['where'].this)})"
         
@@ -779,7 +795,8 @@ class ObVecJsonTableClient(ObVecClient):
             delete_sql = f"DELETE FROM {JSON_TABLE_DATA_TABLE_NAME}"
 
         logger.debug(f"===================== do delete: {delete_sql}")
-        self.perform_raw_text_sql(delete_sql)
+        res = self.perform_raw_text_sql(delete_sql)
+        return res.rowcount
 
     def _get_full_datatype(self, jdata_type: str):
         if jdata_type.upper() == "VARCHAR":
@@ -788,7 +805,14 @@ class ObVecJsonTableClient(ObVecClient):
             return "DECIMAL(10, 0)"
         return jdata_type
 
-    def _handle_jtable_dml_select(self, ast: Expression, select_with_data_id: bool = False):
+    def _handle_jtable_dml_select(
+        self,
+        ast: Expression,
+        select_with_data_id: bool = False,
+        opt_user_id: Optional[str] = None
+    ):
+        real_user_id = opt_user_id or self.user_id
+
         table_name = ast.args['from'].this.this.this
         if not self._check_table_exists(table_name):
             raise ValueError(f"Table {table_name} does not exists")
@@ -856,8 +880,8 @@ class ObVecJsonTableClient(ObVecClient):
         else:
             ast.args['joins'] = [join_node]
 
-        if self.user_id:
-            extra_filter_str = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{self.user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}'"
+        if real_user_id:
+            extra_filter_str = f"{JSON_TABLE_DATA_TABLE_NAME}.user_id = '{real_user_id}' AND {JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}'"
         else:
             extra_filter_str = f"{JSON_TABLE_DATA_TABLE_NAME}.jtable_name = '{table_name}'"
         if 'where' in ast.args.keys():

{pyobvector-0.2.7.dist-info → pyobvector-0.2.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyobvector
-Version: 0.2.7
+Version: 0.2.8
 Summary: A python SDK for OceanBase Vector Store, based on SQLAlchemy, compatible with Milvus API.
 Author: shanhaikang.shk
 Author-email: shanhaikang.shk@oceanbase.com
@@ -36,7 +36,7 @@ poetry install
 - install with pip:
 
 ```shell
-pip install pyobvector==0.2.7
+pip install pyobvector==0.2.8
 ```
 
 ## Build Doc

{pyobvector-0.2.7.dist-info → pyobvector-0.2.8.dist-info}/RECORD

@@ -7,7 +7,7 @@ pyobvector/client/fts_index_param.py,sha256=hMCjA3Aecnt0uQQT6UQGTIIqdPk1M4gX4-zR
 pyobvector/client/index_param.py,sha256=3gXi66Ey1PO9x5_61CrH7DmPb496kviBQI5NT7nfbGc,6309
 pyobvector/client/milvus_like_client.py,sha256=CpPo6mkGE8iNFpKGBFof3h7E1VTzy1DAPGlFM9F_s8g,26373
 pyobvector/client/ob_vec_client.py,sha256=Yt2nG0w4268hg7DE0tqkGaytGsY-jqojX8hGTQjmsKg,29390
-pyobvector/client/ob_vec_json_table_client.py,sha256=vQQoRGv8Dl2ZCnOcd4aC2g1Vc_15zwsHYw9CgHkdTYs,37807
+pyobvector/client/ob_vec_json_table_client.py,sha256=l9zlryiyuk7wubksWJjmbUuNWErQZwM69K2qu4K8yVE,38621
 pyobvector/client/partitions.py,sha256=Bxwr5yVNlXwZc7SXBC03NeqL9giy4Fe6S2qZdHD8xGw,15621
 pyobvector/client/schema_type.py,sha256=ICCSriOhk-P7Q1PhK0D0XQMTd3ZDFenCrJMXp6hRQdw,1579
 pyobvector/json_table/__init__.py,sha256=X5MmK3f10oyJleUUFZJFeunMEfzmf6P1f_7094b-FZc,554
@@ -29,7 +29,7 @@ pyobvector/schema/vector_index.py,sha256=aNtrEBUclc4s6QuqCZpu3Hj3OdjyhbWgtLiJzo6
 pyobvector/util/__init__.py,sha256=D9EgRDlcMSDhY3uI__vnCl45Or75dOXMWSval5P5fqs,251
 pyobvector/util/ob_version.py,sha256=ZIySam8q_MCiwctAiAHPB4GdAzGQiXEo1wVkc9IOTDU,1539
 pyobvector/util/vector.py,sha256=xyM-NuOyd78K7P3kinqyWvLIzEbf9c-4TKn_QVF7qgw,2265
-pyobvector-0.2.7.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-pyobvector-0.2.7.dist-info/METADATA,sha256=fXmr0B-1B9Q-fqT4K6PAHiHakyqMBHv88lmySvv2t4U,6661
-pyobvector-0.2.7.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-pyobvector-0.2.7.dist-info/RECORD,,
+pyobvector-0.2.8.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+pyobvector-0.2.8.dist-info/METADATA,sha256=PVTCENc27h0nh_UwNAF4M1ewnyHZlRGijs9wR3Gkv1o,6661
+pyobvector-0.2.8.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+pyobvector-0.2.8.dist-info/RECORD,,