pyobvector 0.2.7__tar.gz → 0.2.9__tar.gz

{pyobvector-0.2.7 → pyobvector-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyobvector
-Version: 0.2.7
+Version: 0.2.9
 Summary: A python SDK for OceanBase Vector Store, based on SQLAlchemy, compatible with Milvus API.
 Author: shanhaikang.shk
 Author-email: shanhaikang.shk@oceanbase.com
@@ -15,7 +15,7 @@ Requires-Dist: aiomysql (>=0.2.0,<0.3.0)
 Requires-Dist: numpy (>=1.17.0,<2.0.0)
 Requires-Dist: pydantic (>=2.7.0,<3)
 Requires-Dist: pymysql (>=1.1.1,<2.0.0)
-Requires-Dist: sqlalchemy (>=1.4,<2.0.36)
+Requires-Dist: sqlalchemy (>=1.4,<=3)
 Requires-Dist: sqlglot (>=26.0.1,<27.0.0)
 Description-Content-Type: text/markdown
 
@@ -36,7 +36,7 @@ poetry install
 - install with pip:
 
 ```shell
-pip install pyobvector==0.2.7
+pip install pyobvector==0.2.9
 ```
 
 ## Build Doc

{pyobvector-0.2.7 → pyobvector-0.2.9}/README.md

@@ -15,7 +15,7 @@ poetry install
 - install with pip:
 
 ```shell
-pip install pyobvector==0.2.7
+pip install pyobvector==0.2.9
 ```
 
 ## Build Doc

{pyobvector-0.2.7 → pyobvector-0.2.9}/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
@@ -144,15 +144,16 @@ class ObVecJsonTableClient(ObVecClient):
         if self.user_id is None:
             return True
         
-        session = self.session()
-        distinct_admin_ids = (
-            session.query(ObVecJsonTableClient.JsonTableDataTBL.admin_id)
-            .filter_by(user_id = self.user_id)
-            .distinct()
-            .all()
-        )
-        admin_ids = [admin_id[0] for admin_id in distinct_admin_ids]
-        return (len(admin_ids) == 0) or (len(admin_ids) == 1 and admin_ids[0] == self.admin_id)
+        with self.session() as session:
+            with session.begin():
+                distinct_admin_ids = (
+                    session.query(ObVecJsonTableClient.JsonTableDataTBL.admin_id)
+                    .filter_by(user_id = self.user_id)
+                    .distinct()
+                    .all()
+                )
+                admin_ids = [admin_id[0] for admin_id in distinct_admin_ids]
+                return (len(admin_ids) == 0) or (len(admin_ids) == 1 and admin_ids[0] == self.admin_id)
 
     def _reset(self):
         # Only for test
@@ -163,7 +164,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 +182,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 +648,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 +681,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 +700,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 +747,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 +755,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 +766,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 +785,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 +796,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 +806,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 +881,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 → pyobvector-0.2.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyobvector"
-version = "0.2.7"
+version = "0.2.9"
 description = "A python SDK for OceanBase Vector Store, based on SQLAlchemy, compatible with Milvus API."
 authors = ["shanhaikang.shk <shanhaikang.shk@oceanbase.com>"]
 readme = "README.md"
@@ -8,7 +8,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
 numpy = ">=1.17.0,<2.0.0"
-sqlalchemy = ">=1.4,<2.0.36"
+sqlalchemy = ">=1.4,<=3"
 pymysql = "^1.1.1"
 aiomysql = "^0.2.0"
 sqlglot = "^26.0.1"