diff --git a/.github/workflows/publish-pyobsql-pypi.yml b/.github/workflows/publish-pyobsql-pypi.yml new file mode 100644 index 00000000..15d6db1a --- /dev/null +++ b/.github/workflows/publish-pyobsql-pypi.yml @@ -0,0 +1,132 @@ +name: Publish PyObsql to PyPI + +on: + workflow_dispatch: + inputs: + version: + description: 'Package version to publish (e.g., 0.1.0). Leave empty to use version from pyproject.toml' + required: false + type: string + publish_to_test_pypi: + description: 'Publish to Test PyPI instead of PyPI' + required: false + type: boolean + default: false + update_version: + description: 'Update version in pyproject.toml before publishing' + required: false + type: boolean + default: true + +permissions: + contents: read + +jobs: + build: + runs-on: ubuntu-latest + permissions: + contents: read + outputs: + publish_version: ${{ steps.verify_version.outputs.version }} + publish_to_test_pypi: ${{ inputs.publish_to_test_pypi }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Check current version + working-directory: pyobsql-oceanbase-plugin + id: current_version + run: | + CURRENT_VERSION=$(grep '^version = ' pyproject.toml | sed 's/version = "\(.*\)"/\1/') + echo "version=$CURRENT_VERSION" >> $GITHUB_OUTPUT + echo "Current version in pyproject.toml: $CURRENT_VERSION" + + - name: Update version in pyproject.toml + if: inputs.update_version == true && inputs.version != '' + working-directory: pyobsql-oceanbase-plugin + run: | + # Update version in pyproject.toml + sed -i "s/^version = \".*\"/version = \"${{ inputs.version }}\"/" pyproject.toml + echo "✅ Updated version to ${{ inputs.version }}" + cat pyproject.toml | grep "^version" + + - name: Verify version + working-directory: pyobsql-oceanbase-plugin + id: verify_version + run: | + if [ "${{ inputs.version }}" != "" ] && [ "${{ inputs.update_version }}" == "true" ]; then + PUBLISH_VERSION="${{ inputs.version }}" + else + PUBLISH_VERSION="${{ steps.current_version.outputs.version }}" + fi + echo "Publishing version: $PUBLISH_VERSION" + echo "version=$PUBLISH_VERSION" >> $GITHUB_OUTPUT + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + working-directory: pyobsql-oceanbase-plugin + run: | + # Remove LICENSE file to prevent hatchling from auto-detecting it + rm -f LICENSE LICENSE.txt + python -m build + + - name: Check package + working-directory: pyobsql-oceanbase-plugin + run: | + twine check dist/* + + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: pypi-distributions + path: pyobsql-oceanbase-plugin/dist/* + + publish: + runs-on: ubuntu-latest + needs: build + permissions: + contents: read + + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + name: pypi-distributions + path: dist/ + + - name: Publish to Test PyPI + if: needs.build.outputs.publish_to_test_pypi == 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + run: | + pip install twine + twine upload --repository-url https://test.pypi.org/legacy/ dist/* + + - name: Publish to PyPI + if: needs.build.outputs.publish_to_test_pypi != 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + pip install twine + twine upload dist/* + + - name: Display published package info + run: | + echo "✅ Package published successfully!" + if [ "${{ needs.build.outputs.publish_to_test_pypi }}" == "true" ]; then + echo "📍 Published to: Test PyPI (https://test.pypi.org/project/pyobsql/)" + else + echo "📍 Published to: PyPI (https://pypi.org/project/pyobsql/)" + fi + echo "📦 Version: ${{ needs.build.outputs.publish_version }}" diff --git a/README.md b/README.md index 6a37faa6..d1f29b66 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ OceanBase is a high-performance database compatible with both MySQL and Oracle p | [OceanBase SQLAlchemy Plugin](./oceanbase-sqlalchemy-plugin/README.md) | Python ORM | SQLAlchemy dialect for OceanBase Oracle mode, compatible with SQLAlchemy 1.3+ and 2.0+ | | [OceanBase Dify Plugin](./dify-plugin-oceanbase/README.md) | AI Applications | Enables secure SQL query execution on OceanBase databases through Dify applications | | [LangGraph Checkpoint OceanBase Plugin](./langgraph-checkpoint-oceanbase-plugin/README.md) | LangGraph CheckpointSaver | Implementation of LangGraph CheckpointSaver that uses OceanBase MySQL mode | +| [PyObsql OceanBase Plugin](./pyobsql-oceanbase-plugin/README.md) | Python SDK | A Python SDK for OceanBase SQL with JSON Table support and SQLAlchemy dialect extensions | --- @@ -88,6 +89,14 @@ OceanBase is a high-performance database compatible with both MySQL and Oracle p --- +### ✅ PyObsql OceanBase Plugin + +- **Function**: A Python SDK for OceanBase SQL, providing extended SQLAlchemy dialect support, JSON Table operations, and advanced data types (VECTOR, SPARSE_VECTOR, ARRAY, POINT). +- **Use Case**: Python applications that need to interact with OceanBase databases using SQLAlchemy with OceanBase-specific features. +- **Documentation**: [PyObsql OceanBase Plugin](./pyobsql-oceanbase-plugin/README.md) + +--- + ## 📚 Full Documentation Links | Plugin Name | Documentation Link | @@ -100,6 +109,7 @@ OceanBase is a high-performance database compatible with both MySQL and Oracle p | OceanBase SQLAlchemy Plugin | [OceanBase SQLAlchemy Plugin](./oceanbase-sqlalchemy-plugin/README.md) | | OceanBase Dify Plugin | [OceanBase Dify Plugin](./dify-plugin-oceanbase/README.md) | | LangGraph Checkpoint OceanBase Plugin | [LangGraph Checkpoint OceanBase Plugin](./langgraph-checkpoint-oceanbase-plugin/README.md) | +| PyObsql OceanBase Plugin | [PyObsql OceanBase Plugin](./pyobsql-oceanbase-plugin/README.md) | --- diff --git a/README_CN.md b/README_CN.md index 0041729c..c73a70a9 100644 --- a/README_CN.md +++ b/README_CN.md @@ -20,7 +20,7 @@ OceanBase 是一款兼容 MySQL 和 Oracle 协议的高性能数据库。本仓 | [OceanBase SQLAlchemy 插件](./oceanbase-sqlalchemy-plugin/README.md) | Python ORM | SQLAlchemy 方言，支持 OceanBase Oracle 模式，兼容 SQLAlchemy 1.3+ 和 2.0+ | | [LangGraph Checkpoint OceanBase 插件](./langgraph-checkpoint-oceanbase-plugin/README.md) | 保存 LangGraph 的 checkpoint | 使用 OceanBase MySQL 模式实现了 LangGraph CheckpointSaver | | [OceanBase Dify 插件](./dify-plugin-oceanbase/README_CN.md) | AI 应用 | 通过 Dify 应用程序在 OceanBase 数据库上安全执行 SQL 查询 | -| [LangGraph Checkpoint OceanBase 插件](./langgraph-checkpoint-oceanbase-plugin/README_CN.md) | 保存 LangGraph 的 checkpoint | 使用 OceanBase MySQL 模式实现了 LangGraph CheckpointSaver | +| [PyObsql OceanBase 插件](./pyobsql-oceanbase-plugin/README.md) | Python SDK | 支持 JSON Table、SQLAlchemy 方言扩展和高级数据类型的 OceanBase Python SDK | --- @@ -91,11 +91,11 @@ OceanBase 是一款兼容 MySQL 和 Oracle 协议的高性能数据库。本仓 --- -### ✅ LangGraph Checkpoint OceanBase 插件 +### ✅ PyObsql OceanBase 插件 -- **功能**：使用 OceanBase MySQL 模式实现了 LangGraph CheckpointSaver。 -- **适用场景**：使用 OceanBase 作为 LangGraph 的 Checkpointer。 -- **详细文档**：[LangGraph Checkpoint OceanBase 插件](./langgraph-checkpoint-oceanbase-plugin/README_CN.md) +- **功能**：OceanBase SQL 的 Python SDK，提供扩展的 SQLAlchemy 方言支持、JSON Table 操作和高级数据类型（VECTOR、SPARSE_VECTOR、ARRAY、POINT）。 +- **适用场景**：需要使用 SQLAlchemy 与 OceanBase 数据库交互并利用 OceanBase 特定功能的 Python 应用程序。 +- **详细文档**：[PyObsql OceanBase 插件](./pyobsql-oceanbase-plugin/README.md) --- @@ -111,7 +111,7 @@ OceanBase 是一款兼容 MySQL 和 Oracle 协议的高性能数据库。本仓 | OceanBase SQLAlchemy 插件 | [OceanBase SQLAlchemy 插件](./oceanbase-sqlalchemy-plugin/README.md) | | LangGraph Checkpoint OceanBase 插件 | [LangGraph Checkpoint OceanBase 插件](./langgraph-checkpoint-oceanbase-plugin/README.md) | | OceanBase Dify 插件 | [OceanBase Dify 插件](./dify-plugin-oceanbase/README_CN.md) | -| LangGraph Checkpoint OceanBase 插件 | [LangGraph Checkpoint OceanBase 插件](./langgraph-checkpoint-oceanbase-plugin/README_CN.md) | +| PyObsql OceanBase 插件 | [PyObsql OceanBase 插件](./pyobsql-oceanbase-plugin/README.md) | --- diff --git a/pyobsql-oceanbase-plugin/.github/workflows/publish-pypi.yml b/pyobsql-oceanbase-plugin/.github/workflows/publish-pypi.yml new file mode 100644 index 00000000..3213341c --- /dev/null +++ b/pyobsql-oceanbase-plugin/.github/workflows/publish-pypi.yml @@ -0,0 +1,97 @@ +name: Publish PyObsql to PyPI + +on: + workflow_dispatch: + inputs: + version: + description: 'Package version to publish (e.g., 0.1.0)' + required: true + type: string + publish_to_test_pypi: + description: 'Publish to Test PyPI instead of PyPI' + required: false + type: boolean + default: false + +permissions: + contents: read + id-token: write # Required for PyPI publishing with trusted publishing + +jobs: + build: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Update version in pyproject.toml + if: inputs.version != '' + run: | + cd pyobsql-oceanbase-plugin + # Update version in pyproject.toml + sed -i "s/^version = \".*\"/version = \"${{ inputs.version }}\"/" pyproject.toml + echo "Updated version to ${{ inputs.version }}" + cat pyproject.toml | grep "^version" + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + run: | + cd pyobsql-oceanbase-plugin + python -m build + + - name: Check package + run: | + cd pyobsql-oceanbase-plugin + twine check dist/* + + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: pypi-distributions + path: pyobsql-oceanbase-plugin/dist/* + + publish: + runs-on: ubuntu-latest + needs: build + permissions: + id-token: write # Required for PyPI trusted publishing + contents: read + + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + name: pypi-distributions + path: dist/ + + - name: Publish to Test PyPI + if: inputs.publish_to_test_pypi == true + uses: pypa/gh-action-pypi-publish@release/v1 + with: + repository-url: https://test.pypi.org/legacy/ + packages-dir: dist/ + + - name: Publish to PyPI + if: inputs.publish_to_test_pypi == false + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: dist/ + + - name: Display published package info + run: | + echo "✅ Package published successfully!" + if [ "${{ inputs.publish_to_test_pypi }}" == "true" ]; then + echo "📍 Published to: Test PyPI (https://test.pypi.org/project/pyobsql/)" + else + echo "📍 Published to: PyPI (https://pypi.org/project/pyobsql/)" + fi + echo "📦 Version: ${{ inputs.version }}" diff --git a/pyobsql-oceanbase-plugin/.gitignore b/pyobsql-oceanbase-plugin/.gitignore new file mode 100644 index 00000000..10b50f9b --- /dev/null +++ b/pyobsql-oceanbase-plugin/.gitignore @@ -0,0 +1,45 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Virtual environments +venv/ +ENV/ +env/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# Testing +.pytest_cache/ +.coverage +htmlcov/ + +# Documentation +docs/_build/ + + + + diff --git a/pyobsql-oceanbase-plugin/LICENSE b/pyobsql-oceanbase-plugin/LICENSE new file mode 100644 index 00000000..6f5c5bb9 --- /dev/null +++ b/pyobsql-oceanbase-plugin/LICENSE @@ -0,0 +1,21 @@ +Apache License +Version 2.0, January 2004 +http://www.apache.org/licenses/ + +Copyright 2024 OceanBase + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + + + + diff --git a/pyobsql-oceanbase-plugin/QUICKSTART.md b/pyobsql-oceanbase-plugin/QUICKSTART.md new file mode 100644 index 00000000..9d37df94 --- /dev/null +++ b/pyobsql-oceanbase-plugin/QUICKSTART.md @@ -0,0 +1,493 @@ +# pyobsql Quick Start Guide + +Get started with pyobsql in 5 minutes! This guide will walk you through the basics of using pyobsql to interact with OceanBase databases. + +## Prerequisites + +- Python 3.9 or higher +- OceanBase database (version 4.4.1 or later) +- Database connection credentials + +## Installation + +Install pyobsql using pip: + +```bash +pip install pyobsql +``` + +Or install from source: + +```bash +git clone https://github.com/oceanbase/ecology-plugins.git +cd ecology-plugins/pyobsql-oceanbase-plugin +pip install -e . +``` + +## Step 1: Connect to OceanBase + +First, import and create a client connection: + +```python +from pyobsql.client import ObClient + +client = ObClient( + uri="127.0.0.1:2881", # Host:Port + user="root@test", # Username@Tenant + password="your_password", # Password + db_name="test" # Database name +) + +print(f"Connected! OceanBase version: {client.ob_version}") +``` + +**Connection Parameters:** +- `uri`: Database host and port (format: `host:port`) +- `user`: Username with tenant (format: `username@tenant`) +- `password`: Database password +- `db_name`: Name of the database to connect to + +## Step 2: Create Your First Table + +Define table columns and create a table: + +```python +from sqlalchemy import Column, Integer, String, JSON + +columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('age', Integer), + Column('email', String(255)), + Column('metadata', JSON) +] + +client.create_table('users', columns=columns) +print("Table 'users' created successfully!") +``` + +## Step 3: Insert Data + +Insert a single record: + +```python +client.insert('users', { + 'id': 1, + 'name': 'Alice', + 'age': 30, + 'email': 'alice@example.com', + 'metadata': {'department': 'Engineering', 'role': 'Developer'} +}) +``` + +Insert multiple records at once: + +```python +users_data = [ + {'id': 2, 'name': 'Bob', 'age': 25, 'email': 'bob@example.com', 'metadata': {}}, + {'id': 3, 'name': 'Charlie', 'age': 35, 'email': 'charlie@example.com', 'metadata': {}} +] + +client.insert('users', users_data) +``` + +## Step 4: Query Data + +Query all records: + +```python +result = client.get('users') +for row in result: + print(f"ID: {row.id}, Name: {row.name}, Email: {row.email}") +``` + +Query by primary key: + +```python +result = client.get('users', ids=1) +user = list(result)[0] +print(f"Found user: {user.name}") +``` + +Query with conditions: + +```python +from sqlalchemy import Table + +table = Table('users', client.metadata_obj, autoload_with=client.engine) + +# Find users older than 25 +result = client.get( + 'users', + where_clause=[table.c.age > 25], + n_limits=10 +) + +for row in result: + print(f"{row.name} is {row.age} years old") +``` + +## Step 5: Update Data + +Update records using SQLAlchemy: + +```python +from sqlalchemy import Table, update + +table = Table('users', client.metadata_obj, autoload_with=client.engine) + +with client.engine.connect() as conn: + with conn.begin(): + update_stmt = update(table).where(table.c.id == 1).values( + age=31, + metadata={'department': 'Engineering', 'role': 'Senior Developer'} + ) + conn.execute(update_stmt) + +print("User updated successfully!") +``` + +## Step 6: Delete Data + +Delete by primary key: + +```python +client.delete('users', ids=1) +``` + +Delete multiple records: + +```python +client.delete('users', ids=[2, 3]) +``` + +Delete with conditions: + +```python +table = Table('users', client.metadata_obj, autoload_with=client.engine) +client.delete('users', where_clause=[table.c.age < 18]) +``` + +## Working with Advanced Data Types + +### VECTOR Type (for embeddings) + +```python +from sqlalchemy import Column, Integer, String +from pyobsql.schema import VECTOR + +columns = [ + Column('id', Integer, primary_key=True), + Column('text', String(500)), + Column('embedding', VECTOR(128)) # 128-dimensional vector +] + +client.create_table('documents', columns=columns) + +# Insert vector data +embedding = [0.1] * 128 # Your embedding vector +client.insert('documents', { + 'id': 1, + 'text': 'Sample document', + 'embedding': embedding +}) +``` + +### ARRAY Type + +```python +from pyobsql.schema import ARRAY + +columns = [ + Column('id', Integer, primary_key=True), + Column('tags', ARRAY(String(50))), # String array + Column('scores', ARRAY(Integer)) # Integer array +] + +client.create_table('items', columns=columns) + +client.insert('items', { + 'id': 1, + 'tags': ['python', 'database', 'oceanbase'], + 'scores': [95, 87, 92] +}) +``` + +### SPARSE_VECTOR Type + +```python +from pyobsql.schema import SPARSE_VECTOR + +columns = [ + Column('id', Integer, primary_key=True), + Column('sparse_vec', SPARSE_VECTOR) +] + +client.create_table('sparse_data', columns=columns) + +# Sparse vector as dictionary: {index: value} +sparse_vector = {1: 0.5, 5: 0.8, 10: 0.3} +client.insert('sparse_data', { + 'id': 1, + 'sparse_vec': sparse_vector +}) +``` + +### JSON Type + +```python +from sqlalchemy import JSON + +columns = [ + Column('id', Integer, primary_key=True), + Column('metadata', JSON) +] + +client.create_table('products', columns=columns) + +client.insert('products', { + 'id': 1, + 'metadata': { + 'price': 99.99, + 'category': 'Electronics', + 'specs': {'color': 'black', 'weight': '1.5kg'} + } +}) +``` + +## Working with JSON Table + +pyobsql provides special data types for JSON Table operations: + +```python +from pyobsql import ( + JsonTableBool, + JsonTableInt, + JsonTableTimestamp, + JsonTableVarcharFactory, + JsonTableDecimalFactory, + val2json +) + +# Create JSON Table types +bool_type = JsonTableBool(val=True) +int_type = JsonTableInt(val=42) +varchar_factory = JsonTableVarcharFactory(length=255) +varchar_type = varchar_factory.get_json_table_varchar_type()(val="test") + +# Convert to JSON +json_value = val2json(bool_type.val) +``` + +## Using json_value Function + +Extract values from JSON columns: + +```python +from pyobsql import json_value +from sqlalchemy import Table, select + +table = Table('products', client.metadata_obj, autoload_with=client.engine) + +stmt = select( + table.c.id, + json_value(table.c.metadata, '$.price', 'DECIMAL(10,2)').label('price') +).where(table.c.id == 1) + +with client.engine.connect() as conn: + result = conn.execute(stmt) + for row in result: + print(f"Product {row.id} price: {row.price}") +``` + +## Creating Partitioned Tables + +Create a table with Range partitioning: + +```python +from pyobsql.client.partitions import ObRangePartition, RangeListPartInfo + +partition = ObRangePartition( + is_range_columns=False, + range_part_infos=[ + RangeListPartInfo('p0', 100), + RangeListPartInfo('p1', 200), + RangeListPartInfo('p2', 'MAXVALUE') + ], + range_expr='id' +) + +columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('data', String(500)) +] + +client.create_table('partitioned_data', columns=columns, partitions=partition) + +# Insert to specific partition +client.insert( + 'partitioned_data', + {'id': 50, 'name': 'test', 'data': 'sample'}, + partition_name='p0' +) +``` + +## Upsert Operations (Insert or Replace) + +Use `upsert` to insert or replace records: + +```python +# If id=1 exists, replace it; otherwise insert +client.upsert('users', { + 'id': 1, + 'name': 'Alice Updated', + 'age': 31, + 'email': 'alice.new@example.com', + 'metadata': {'status': 'updated'} +}) +``` + +## Using SQLAlchemy Queries + +For complex queries, use SQLAlchemy directly: + +```python +from sqlalchemy import Table, select, func + +table = Table('users', client.metadata_obj, autoload_with=client.engine) + +stmt = select( + table.c.id, + table.c.name, + func.json_extract(table.c.metadata, '$.department').label('department') +).where( + table.c.age > 25 +).order_by( + table.c.id.desc() +).limit(10) + +with client.engine.connect() as conn: + result = conn.execute(stmt) + for row in result: + print(f"{row.name} - {row.department}") +``` + +## Complete Example + +Here's a complete working example: + +```python +from pyobsql.client import ObClient +from sqlalchemy import Column, Integer, String, JSON + +# 1. Connect +client = ObClient( + uri="127.0.0.1:2881", + user="root@test", + password="password", + db_name="test" +) + +# 2. Create table +columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('email', String(255)), + Column('metadata', JSON) +] +client.create_table('employees', columns=columns) + +# 3. Insert data +client.insert('employees', { + 'id': 1, + 'name': 'John Doe', + 'email': 'john@example.com', + 'metadata': {'department': 'Sales', 'level': 'Senior'} +}) + +# 4. Query data +result = client.get('employees', ids=1) +employee = list(result)[0] +print(f"Employee: {employee.name}, Email: {employee.email}") + +# 5. Update data +from sqlalchemy import Table, update +table = Table('employees', client.metadata_obj, autoload_with=client.engine) +with client.engine.connect() as conn: + with conn.begin(): + update_stmt = update(table).where(table.c.id == 1).values( + metadata={'department': 'Sales', 'level': 'Manager'} + ) + conn.execute(update_stmt) + +# 6. Clean up (optional) +client.drop_table_if_exist('employees') +``` + +## Common Operations Reference + +### Check if table exists +```python +if client.check_table_exists('users'): + print("Table exists!") +``` + +### Drop table +```python +client.drop_table_if_exist('users') +``` + +### Refresh metadata +```python +# Refresh all tables +client.refresh_metadata() + +# Refresh specific tables +client.refresh_metadata(tables=['users', 'products']) +``` + +### Create indexes +```python +from sqlalchemy import Index + +indexes = [ + Index('idx_name', 'name'), + Index('idx_email', 'email') +] + +client.create_table('users', columns=columns, indexes=indexes) +``` + +## Next Steps + +- Read the [full README](README.md) for detailed documentation +- Explore advanced features like partitioning and vector search +- Check out the [test examples](tests/) for more use cases +- Join the community for support and updates + +## Troubleshooting + +### Connection Issues +- Verify your database is running and accessible +- Check firewall settings +- Ensure credentials are correct (username@tenant format) + +### Import Errors +- Make sure pyobsql is installed: `pip list | grep pyobsql` +- Check Python version: `python --version` (requires 3.9+) + +### Type Errors +- Ensure you're using the correct data types from `pyobsql.schema` +- Check OceanBase version compatibility (some features require 4.4.1+) + +## Need Help? + +- **Documentation**: See [README.md](README.md) +- **Issues**: Report on [GitHub Issues](https://github.com/oceanbase/ecology-plugins/issues) +- **Examples**: Check the `tests/` directory for more examples + +--- + +**Happy coding with pyobsql! 🚀** + diff --git a/pyobsql-oceanbase-plugin/README.md b/pyobsql-oceanbase-plugin/README.md new file mode 100644 index 00000000..ad319176 --- /dev/null +++ b/pyobsql-oceanbase-plugin/README.md @@ -0,0 +1,573 @@ +# pyobsql + +A python SDK for OceanBase SQL, including JSON Table support and SQLAlchemy dialect extensions. + +## Installation + +### Install from Source + +```shell +git clone https://github.com/oceanbase/ecology-plugins.git +cd ecology-plugins/pyobsql-oceanbase-plugin +pip install -e . +``` + +### Install from PyPI + +```shell +pip install pyobsql +``` + +## Features + +`pyobsql` provides the following features: + +- **OceanBase SQL Dialect Parsing**: Extended SQLGlot support for OceanBase-specific SQL syntax +- **SQLAlchemy Integration**: Provides synchronous and asynchronous OceanBase dialect support +- **Extended Data Types**: Supports VECTOR, SPARSE_VECTOR, ARRAY, POINT and other OceanBase-specific types +- **JSON Table Support**: Virtual data types and utility functions for handling JSON tables +- **Table Structure Reflection**: Automatically parses OceanBase table structures +- **Partition Support**: Supports various partition strategies including Range, Hash, Key, List, etc. + +## Detailed Usage Guide + +### 1. Connect to Database + +```python +from pyobsql.client import ObClient + +client = ObClient( + uri="127.0.0.1:2881", + user="root@test", + password="password", + db_name="test" +) +``` + +### 2. Create Tables + +#### 2.1 Create Basic Table + +```python +from sqlalchemy import Column, Integer, String, JSON +from pyobsql.schema import VECTOR, SPARSE_VECTOR, ARRAY, POINT + +columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('embedding', VECTOR(128)), + Column('sparse_vec', SPARSE_VECTOR), + Column('tags', ARRAY(String(50))), + Column('location', POINT(srid=4326)), + Column('metadata', JSON) +] + +client.create_table( + table_name='my_table', + columns=columns +) +``` + +#### 2.2 Create Partitioned Table + +```python +from pyobsql.client.partitions import ObRangePartition, RangeListPartInfo + +range_partition = ObRangePartition( + is_range_columns=False, + range_part_infos=[ + RangeListPartInfo('p0', 100), + RangeListPartInfo('p1', 200), + RangeListPartInfo('p2', 'MAXVALUE') + ], + range_expr='id' +) + +columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('embedding', VECTOR(128)) +] + +client.create_table( + table_name='partitioned_table', + columns=columns, + partitions=range_partition +) +``` + +#### 2.3 Create Table with Indexes + +```python +from sqlalchemy import Index + +indexes = [ + Index('idx_name', 'name'), + Index('idx_embedding', 'embedding', postgresql_using='hnsw') +] + +client.create_table( + table_name='indexed_table', + columns=columns, + indexes=indexes +) +``` + +### 3. Insert Data + +#### 3.1 Insert Single Record + +```python +client.insert( + table_name='my_table', + data={ + 'id': 1, + 'name': 'example', + 'embedding': [0.1, 0.2, 0.3, ...], + 'sparse_vec': {1: 0.5, 5: 0.8, 10: 0.3}, + 'tags': ['tag1', 'tag2', 'tag3'], + 'location': (116.3974, 39.9093), + 'metadata': {'key': 'value'} + } +) +``` + +#### 3.2 Batch Insert Data + +```python +data_list = [ + { + 'id': i, + 'name': f'item_{i}', + 'embedding': [0.1 * i, 0.2 * i, 0.3 * i, ...], + 'tags': [f'tag_{i}'], + 'metadata': {'index': i} + } + for i in range(100) +] + +client.insert( + table_name='my_table', + data=data_list +) +``` + +#### 3.3 Insert to Specified Partition + +```python +client.insert( + table_name='partitioned_table', + data={'id': 50, 'name': 'test', 'embedding': [0.1, 0.2, ...]}, + partition_name='p0' +) +``` + +#### 3.4 Use REPLACE INTO (Insert or Replace) + +```python +from pyobsql.schema import ReplaceStmt +from sqlalchemy import Table + +table = Table('my_table', client.metadata_obj, autoload_with=client.engine) +with client.engine.connect() as conn: + with conn.begin(): + stmt = ReplaceStmt(table).values({ + 'id': 1, + 'name': 'updated_name', + 'embedding': [0.5, 0.6, ...] + }) + conn.execute(stmt) + +client.upsert( + table_name='my_table', + data={'id': 1, 'name': 'updated_name', 'embedding': [0.5, 0.6, ...]} +) +``` + +### 4. Update Data + +```python +from sqlalchemy import Table + +table = Table('my_table', client.metadata_obj, autoload_with=client.engine) + +client.update( + table_name='my_table', + values_clause=[table.c.name == 'new_name'], + where_clause=[table.c.id == 1] +) + +client.update( + table_name='my_table', + values_clause=[ + table.c.name == 'updated_name', + table.c.metadata == {'status': 'updated'} + ], + where_clause=[table.c.id.in_([1, 2, 3])] +) + +client.update( + table_name='partitioned_table', + values_clause=[table.c.name == 'new_name'], + where_clause=[table.c.id == 50], + partition_name='p0' +) +``` + +### 5. Delete Data + +```python +client.delete( + table_name='my_table', + ids=1 +) + +client.delete( + table_name='my_table', + ids=[1, 2, 3] +) + +from sqlalchemy import Table + +table = Table('my_table', client.metadata_obj, autoload_with=client.engine) +client.delete( + table_name='my_table', + where_clause=[table.c.name == 'old_name'] +) + +client.delete( + table_name='partitioned_table', + ids=50, + partition_name='p0' +) +``` + +### 6. Query Data + +#### 6.1 Basic Queries + +```python +from sqlalchemy import Table, select + +table = Table('my_table', client.metadata_obj, autoload_with=client.engine) + +result = client.get(table_name='my_table') +for row in result: + print(row) + +result = client.get( + table_name='my_table', + ids=1 +) + +result = client.get( + table_name='my_table', + ids=[1, 2, 3] +) + +result = client.get( + table_name='my_table', + where_clause=[table.c.name == 'example'] +) + +result = client.get( + table_name='my_table', + output_column_name=['id', 'name', 'embedding'] +) + +result = client.get( + table_name='my_table', + n_limits=10 +) + +result = client.get( + table_name='partitioned_table', + partition_names=['p0', 'p1'] +) +``` + +#### 6.2 Use SQLAlchemy Native Queries + +```python +from sqlalchemy import Table, select, func, text + +table = Table('my_table', client.metadata_obj, autoload_with=client.engine) + +stmt = select( + table.c.id, + table.c.name, + func.json_extract(table.c.metadata, '$.key').label('extracted_key') +).where( + table.c.id > 10 +).order_by( + table.c.id.desc() +).limit(10) + +with client.engine.connect() as conn: + result = conn.execute(stmt) + for row in result: + print(row) +``` + +### 7. JSON Table Support + +#### 7.1 JSON Table Virtual Data Types + +```python +from pyobsql import ( + JsonTableBool, + JsonTableInt, + JsonTableTimestamp, + JsonTableVarcharFactory, + JsonTableDecimalFactory, + val2json +) + +bool_type = JsonTableBool(val=True) +int_type = JsonTableInt(val=42) +varchar_factory = JsonTableVarcharFactory(length=255) +varchar_type = varchar_factory.get_json_table_varchar_type()(val="test") +decimal_factory = JsonTableDecimalFactory(precision=10, scale=2) +decimal_type = decimal_factory.get_json_table_decimal_type()(val=123.45) + +json_value = val2json(bool_type) +``` + +#### 7.2 Use json_value Function + +```python +from pyobsql import json_value +from sqlalchemy import Table, select + +table = Table('my_table', client.metadata_obj, autoload_with=client.engine) + +stmt = select( + table.c.id, + json_value(table.c.metadata, '$.key', 'VARCHAR(100)').label('extracted_value') +).where( + table.c.id == 1 +) + +with client.engine.connect() as conn: + result = conn.execute(stmt) + for row in result: + print(row) +``` + +### 8. SQL Parsing (SQLGlot) + +#### 8.1 Parse OceanBase SQL + +```python +from pyobsql import OceanBase +from sqlglot import parse_one, transpile + +sql = "ALTER TABLE t2 CHANGE COLUMN c2 changed_col INT" +ast = parse_one(sql, dialect=OceanBase) +print(ast) + +sql = "ALTER TABLE t1 MODIFY COLUMN c1 VARCHAR(100) NOT NULL" +ast = parse_one(sql, dialect=OceanBase) + +sql = "ALTER TABLE t1 DROP COLUMN c1" +ast = parse_one(sql, dialect=OceanBase) + +sql = "SELECT * FROM table1" +mysql_sql = transpile(sql, read=OceanBase, write="mysql")[0] +``` + +### 9. Data Type Details + +#### 9.1 VECTOR (Vector Type) + +```python +from pyobsql.schema import VECTOR +from pyobsql.util import Vector + +column = Column('embedding', VECTOR(128)) + +vector_data = [0.1, 0.2, 0.3, ...] +vector_obj = Vector(vector_data) + +client.insert( + table_name='vector_table', + data={'id': 1, 'embedding': vector_data} +) +``` + +#### 9.2 SPARSE_VECTOR (Sparse Vector Type) + +```python +from pyobsql.schema import SPARSE_VECTOR +from pyobsql.util import SparseVector + +column = Column('sparse_vec', SPARSE_VECTOR) + +sparse_data = {1: 0.5, 5: 0.8, 10: 0.3} + +client.insert( + table_name='sparse_table', + data={'id': 1, 'sparse_vec': sparse_data} +) +``` + +#### 9.3 ARRAY (Array Type) + +```python +from pyobsql.schema import ARRAY +from sqlalchemy import String, Integer + +tags_column = Column('tags', ARRAY(String(50))) +scores_column = Column('scores', ARRAY(Integer)) +nested_array = Column('matrix', ARRAY(ARRAY(Integer))) + +client.insert( + table_name='array_table', + data={ + 'id': 1, + 'tags': ['tag1', 'tag2', 'tag3'], + 'scores': [100, 200, 300], + 'matrix': [[1, 2], [3, 4]] + } +) +``` + +#### 9.4 POINT (Geographic Coordinate Point Type) + +```python +from pyobsql.schema import POINT + +location_column = Column('location', POINT(srid=4326)) + +client.insert( + table_name='location_table', + data={ + 'id': 1, + 'location': (116.3974, 39.9093) + } +) + +from pyobsql.schema import ST_GeomFromText, st_distance, st_dwithin + +table = Table('location_table', client.metadata_obj, autoload_with=client.engine) +stmt = select( + table.c.id, + st_distance( + table.c.location, + ST_GeomFromText('POINT(116.3974 39.9093)', 4326) + ).label('distance') +).where( + st_dwithin( + table.c.location, + ST_GeomFromText('POINT(116.3974 39.9093)', 4326), + 1000 + ) +) +``` + +### 10. Table Structure Management + +#### 10.1 Drop Table + +```python +client.drop_table_if_exist('my_table') +``` + +#### 10.2 Drop Index + +```python +client.drop_index(table_name='my_table', index_name='idx_name') +``` + +#### 10.3 Refresh Metadata + +```python +client.refresh_metadata() + +client.refresh_metadata(tables=['my_table', 'other_table']) +``` + +### 11. Async Operations (Optional) + +```python +from pyobsql.schema import AsyncOceanBaseDialect +from sqlalchemy.ext.asyncio import create_async_engine + +engine = create_async_engine( + "mysql+aiomysql://user:password@127.0.0.1:2881/dbname", + dialect=AsyncOceanBaseDialect() +) + +async with engine.connect() as conn: + result = await conn.execute(select(table)) + rows = result.fetchall() +``` + +## Complete Example + +```python +from pyobsql.client import ObClient +from pyobsql.client.partitions import ObRangePartition, RangeListPartInfo +from sqlalchemy import Column, Integer, String, JSON, Table +from pyobsql.schema import VECTOR, ARRAY + +client = ObClient( + uri="127.0.0.1:2881", + user="root@test", + password="password", + db_name="test" +) + +partition = ObRangePartition( + is_range_columns=False, + range_part_infos=[ + RangeListPartInfo('p0', 100), + RangeListPartInfo('p1', 'MAXVALUE') + ], + range_expr='id' +) + +columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('embedding', VECTOR(128)), + Column('tags', ARRAY(String(50))), + Column('metadata', JSON) +] + +client.create_table('products', columns=columns, partitions=partition) + +data = [ + { + 'id': i, + 'name': f'product_{i}', + 'embedding': [0.1 * i] * 128, + 'tags': [f'tag_{i}', f'category_{i % 5}'], + 'metadata': {'price': i * 10, 'stock': 100 - i} + } + for i in range(50) +] +client.insert('products', data) + +table = Table('products', client.metadata_obj, autoload_with=client.engine) +result = client.get( + table_name='products', + where_clause=[table.c.id < 10], + output_column_name=['id', 'name', 'tags'] +) + +for row in result: + print(f"ID: {row.id}, Name: {row.name}, Tags: {row.tags}") + +client.update( + table_name='products', + values_clause=[table.c.metadata == {'price': 999, 'stock': 50}], + where_clause=[table.c.id == 1] +) + +client.delete(table_name='products', ids=[1, 2, 3]) +``` + +## License + +Apache-2.0 \ No newline at end of file diff --git a/pyobsql-oceanbase-plugin/pyobsql/__init__.py b/pyobsql-oceanbase-plugin/pyobsql/__init__.py new file mode 100644 index 00000000..1b32ae7f --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/__init__.py @@ -0,0 +1,36 @@ +"""A python SDK for OceanBase SQL, including JSON Table support and SQLAlchemy dialect. + +`pyobsql` provides: +* OceanBase SQL dialect for SQLAlchemy +* JSON Table support with virtual data types +* SQL parsing and execution utilities +""" +from .json_table import ( + OceanBase, + ChangeColumn, + JType, + JsonTableDataType, + JsonTableBool, + JsonTableTimestamp, + JsonTableVarcharFactory, + JsonTableDecimalFactory, + JsonTableInt, + val2json, + json_value +) + +__all__ = [ + "OceanBase", + "ChangeColumn", + "JType", + "JsonTableDataType", + "JsonTableBool", + "JsonTableTimestamp", + "JsonTableVarcharFactory", + "JsonTableDecimalFactory", + "JsonTableInt", + "val2json", + "json_value", +] + +__version__ = "0.1.0" diff --git a/pyobsql-oceanbase-plugin/pyobsql/client/__init__.py b/pyobsql-oceanbase-plugin/pyobsql/client/__init__.py new file mode 100644 index 00000000..ef9ba496 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/client/__init__.py @@ -0,0 +1,44 @@ +"""OceanBase SQL Client module. + +* ObClient OceanBase SQL client +* FtsParser Text Parser Type for Full Text Search +* FtsIndexParam Full Text Search index parameter +* ObPartition Partition strategy base class +* PartType Partition type enum +""" +from .ob_client import ObClient +from .fts_index_param import FtsParser, FtsIndexParam +from .partitions import ( + ObPartition, + PartType, + ObRangePartition, + ObSubRangePartition, + ObListPartition, + ObSubListPartition, + ObHashPartition, + ObSubHashPartition, + ObKeyPartition, + ObSubKeyPartition, + RangeListPartInfo, +) + +__all__ = [ + "ObClient", + "FtsParser", + "FtsIndexParam", + "ObPartition", + "PartType", + "ObRangePartition", + "ObSubRangePartition", + "ObListPartition", + "ObSubListPartition", + "ObHashPartition", + "ObSubHashPartition", + "ObKeyPartition", + "ObSubKeyPartition", + "RangeListPartInfo", +] + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/client/enum.py b/pyobsql-oceanbase-plugin/pyobsql/client/enum.py new file mode 100644 index 00000000..577b9a65 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/client/enum.py @@ -0,0 +1,7 @@ +"""Common module for int type enumerate.""" +from enum import Enum + + +class IntEnum(int, Enum): + """Int type enumerate definition.""" + \ No newline at end of file diff --git a/pyobsql-oceanbase-plugin/pyobsql/client/exceptions.py b/pyobsql-oceanbase-plugin/pyobsql/client/exceptions.py new file mode 100644 index 00000000..495f1485 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/client/exceptions.py @@ -0,0 +1,115 @@ +"""Exception for MilvusLikeClient.""" +from .enum import IntEnum + + +class ErrorCode(IntEnum): + """Error codes for MilvusLikeClient.""" + SUCCESS = 0 + UNEXPECTED_ERROR = 1 + INVALID_ARGUMENT = 2 + NOT_SUPPORTED = 3 + COLLECTION_NOT_FOUND = 100 + INDEX_NOT_FOUND = 700 + + +class ObException(Exception): + """Base class for MilvusLikeClient exception.""" + def __init__( + self, + code: int = ErrorCode.UNEXPECTED_ERROR, + message: str = "", + ) -> None: + super().__init__() + self._code = code + self._message = message + + @property + def code(self): + """Get error code.""" + return self._code + + @property + def message(self): + """Get error message.""" + return self._message + + def __str__(self) -> str: + return f"<{type(self).__name__}: (code={self.code}, message={self.message})>" + + +class PartitionFieldException(ObException): + """Raise when partition field invalid""" + + +class PrimaryKeyException(ObException): + """Raise when primary key are invalid""" + + +class VectorFieldParamException(ObException): + """Raise when Vector Field parameters are invalid""" + + +class VarcharFieldParamException(ObException): + """Raise when Varchar Field parameters are invalid""" + + +class ArrayFieldParamException(ObException): + """Raise when Array Field parameters are invalid""" + + +class CollectionStatusException(ObException): + """Raise when collection status is invalid""" + + +class VectorMetricTypeException(ObException): + """Raise when vector metric type is invalid""" + + +class MilvusCompatibilityException(ObException): + """Raise when compatibility conflict with milvus""" + + +class ClusterVersionException(ObException): + """Raise when cluster version is not valid""" + + +class ExceptionsMessage: + """Exception Messages definition.""" + PartitionExprNotExists = "Partition expression string does not exist." + PartitionMultiField = "Multi-Partition Field is not supported." + PartitionLevelMoreThanTwo = "Partition Level should less than or equal to 2." + PartitionRangeCutNotIncreasing = ( + "Range cut list should be monotonically increasing." + ) + PartitionRangeExprMissing = ( + "Range expression is necessary when partition type is Range" + ) + PartitionRangeColNameListMissing = ( + "Column name list is necessary when parititon type is RangeColumns" + ) + PartitionListExprMissing = ( + "List expression is necessary when partition type is List" + ) + PartitionListColNameListMissing = ( + "Column name list is necessary when parititon type is ListColumns" + ) + PartitionHashNameListAndPartCntMissing = ( + "One of hash_part_name_list and part_count must be set when partition type is Hash" + ) + PartitionKeyNameListAndPartCntMissing = ( + "One of key_part_name_list and part_count must be set when partition type is Key" + ) + PrimaryFieldType = "Param primary_field must be int or str type." + VectorFieldMissingDimParam = "Param 'dim' must be set for vector field." + VarcharFieldMissingLengthParam = "Param 'max_length' must be set for varchar field." + ArrayFieldMissingElementType = "Param 'element_type' must be set for array field." + ArrayFieldInvalidElementType = ( + "Param 'element_type' can not be array/vector/varchar." + ) + CollectionNotExists = "Collection does not exist." + MetricTypeParamTypeInvalid = "MetricType param type should be string." + MetricTypeValueInvalid = "MetricType should be 'l2'/'ip'/'neg_ip'/'cosine' in ann search." + UsingInIDsWhenMultiPrimaryKey = "Using 'ids' when table has multi primary key." + ClusterVersionIsLow = ( + "OceanBase %s feature is not supported because cluster version is below %s." + ) diff --git a/pyobsql-oceanbase-plugin/pyobsql/client/fts_index_param.py b/pyobsql-oceanbase-plugin/pyobsql/client/fts_index_param.py new file mode 100644 index 00000000..72bc9eeb --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/client/fts_index_param.py @@ -0,0 +1,73 @@ +"""A module to specify fts index parameters""" +from enum import Enum +from typing import List, Optional, Union + +class FtsParser(Enum): + """Built-in full-text search parser types supported by OceanBase""" + IK = 0 + NGRAM = 1 + NGRAM2 = 2 # NGRAM2 parser (supported from V4.3.5 BP2+) + BASIC_ENGLISH = 3 # Basic English parser + JIEBA = 4 # jieba parser + + +class FtsIndexParam: + """Full-text search index parameter. + + Args: + index_name: Index name + field_names: List of field names to create full-text index on + parser_type: Parser type, can be FtsParser enum or string (for custom parsers) + If None, uses default Space parser + """ + def __init__( + self, + index_name: str, + field_names: List[str], + parser_type: Optional[Union[FtsParser, str]] = None, + ): + self.index_name = index_name + self.field_names = field_names + self.parser_type = parser_type + + def param_str(self) -> Optional[str]: + """Convert parser type to string format for SQL.""" + if self.parser_type is None: + return None # Default Space parser, no need to specify + + if isinstance(self.parser_type, str): + # Custom parser name (e.g., "thai_ftparser") + return self.parser_type.lower() + + if isinstance(self.parser_type, FtsParser): + if self.parser_type == FtsParser.IK: + return "ik" + if self.parser_type == FtsParser.NGRAM: + return "ngram" + if self.parser_type == FtsParser.NGRAM2: + return "ngram2" + if self.parser_type == FtsParser.BASIC_ENGLISH: + return "beng" + if self.parser_type == FtsParser.JIEBA: + return "jieba" + # Raise exception for unrecognized FtsParser enum values + raise ValueError(f"Unrecognized FtsParser enum value: {self.parser_type}") + + return None + + def __iter__(self): + yield "index_name", self.index_name + yield "field_names", self.field_names + if self.parser_type: + yield "parser_type", self.parser_type + + def __str__(self): + return str(dict(self)) + + def __eq__(self, other: object) -> bool: + if isinstance(other, self.__class__): + return dict(self) == dict(other) + + if isinstance(other, dict): + return dict(self) == other + return False diff --git a/pyobsql-oceanbase-plugin/pyobsql/client/ob_client.py b/pyobsql-oceanbase-plugin/pyobsql/client/ob_client.py new file mode 100644 index 00000000..b5999700 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/client/ob_client.py @@ -0,0 +1,456 @@ +"""The OceanBase Client""" +import logging +from typing import List, Optional, Dict, Union +from urllib.parse import quote + +import sqlalchemy.sql.functions as func_mod +from sqlalchemy import ( + create_engine, + MetaData, + Table, + Column, + Index, + select, + delete, + update, + insert, + text, + inspect, + and_, +) +from sqlalchemy.dialects import registry +from sqlalchemy.exc import NoSuchTableError + +from .partitions import ObPartition +from ..schema import ( + ObTable, + ST_GeomFromText, + st_distance, + st_dwithin, + st_astext, + ReplaceStmt, +) +from ..util import ObVersion + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class ObClient: + """The OceanBase Client""" + + def __init__( + self, + uri: str = "127.0.0.1:2881", + user: str = "root@test", + password: str = "", + db_name: str = "test", + **kwargs, + ): + registry.register("mysql.oceanbase", "pyobsql.schema.dialect", "OceanBaseDialect") + + setattr(func_mod, "ST_GeomFromText", ST_GeomFromText) + setattr(func_mod, "st_distance", st_distance) + setattr(func_mod, "st_dwithin", st_dwithin) + setattr(func_mod, "st_astext", st_astext) + + user = quote(user, safe="") + password = quote(password, safe="") + + connection_str = ( + f"mysql+oceanbase://{user}:{password}@{uri}/{db_name}?charset=utf8mb4" + ) + self.engine = create_engine(connection_str, **kwargs) + self.metadata_obj = MetaData() + self.metadata_obj.reflect(bind=self.engine) + + with self.engine.connect() as conn: + with conn.begin(): + res = conn.execute(text("SELECT OB_VERSION() FROM DUAL")) + version = [r[0] for r in res][0] + self.ob_version = ObVersion.from_db_version_string(version) + + def refresh_metadata(self, tables: Optional[list[str]] = None): + """Reload metadata from the database. + + Args: + tables (Optional[list[str]]): names of the tables to refresh. If None, refresh all tables. + """ + if tables is not None: + for table_name in tables: + if table_name in self.metadata_obj.tables: + self.metadata_obj.remove(Table(table_name, self.metadata_obj)) + self.metadata_obj.reflect(bind=self.engine, only=tables, extend_existing=True) + else: + self.metadata_obj.clear() + self.metadata_obj.reflect(bind=self.engine, extend_existing=True) + + def _is_seekdb(self) -> bool: + """Check if the database is SeekDB by querying version. + + Returns: + bool: True if database is SeekDB, False otherwise + """ + is_seekdb = False + try: + if hasattr(self, '_is_seekdb_cached'): + return self._is_seekdb_cached + with self.engine.connect() as conn: + result = conn.execute(text("SELECT VERSION()")) + version_str = [r[0] for r in result][0] + is_seekdb = "SeekDB" in version_str + self._is_seekdb_cached = is_seekdb + logger.debug(f"Version query result: {version_str}, is_seekdb: {is_seekdb}") + except Exception as e: + logger.warning(f"Failed to query version: {e}") + return is_seekdb + + def _insert_partition_hint_for_query_sql(self, sql: str, partition_hint: str): + from_index = sql.find("FROM") + assert from_index != -1 + first_space_after_from = sql.find(" ", from_index + len("FROM") + 1) + if first_space_after_from == -1: + return sql + " " + partition_hint + return ( + sql[:first_space_after_from] + + " " + + partition_hint + + sql[first_space_after_from:] + ) + + def check_table_exists(self, table_name: str): + """Check if table exists. + + Args: + table_name (string): table name + + Returns: + bool: True if table exists, False otherwise + """ + inspector = inspect(self.engine) + return inspector.has_table(table_name) + + def create_table( + self, + table_name: str, + columns: List[Column], + indexes: Optional[List[Index]] = None, + partitions: Optional[ObPartition] = None, + **kwargs, + ): + """Create a table. + + Args: + table_name (string): table name + columns (List[Column]): column schema + indexes (Optional[List[Index]): optional index schema + partitions (Optional[ObPartition]): optional partition strategy + **kwargs: additional keyword arguments + """ + kwargs.setdefault("extend_existing", True) + with self.engine.connect() as conn: + with conn.begin(): + if indexes is not None: + table = ObTable( + table_name, + self.metadata_obj, + *columns, + *indexes, + **kwargs, + ) + else: + table = ObTable( + table_name, + self.metadata_obj, + *columns, + **kwargs, + ) + table.create(self.engine, checkfirst=True) + # do partition + if partitions is not None: + conn.execute( + text(f"ALTER TABLE `{table_name}` {partitions.do_compile()}") + ) + + def drop_table_if_exist(self, table_name: str): + """Drop table if exists.""" + try: + table = Table(table_name, self.metadata_obj, autoload_with=self.engine) + except NoSuchTableError: + return + with self.engine.connect() as conn: + with conn.begin(): + table.drop(self.engine, checkfirst=True) + self.metadata_obj.remove(table) + + def drop_index(self, table_name: str, index_name: str): + """drop index on specified table. + + If the index not exists, SQL ERROR 1091 will raise. + """ + with self.engine.connect() as conn: + with conn.begin(): + conn.execute(text(f"DROP INDEX `{index_name}` ON `{table_name}`")) + + def insert( + self, + table_name: str, + data: Union[Dict, List[Dict]], + partition_name: Optional[str] = "", + ): + """Insert data into table. + + Args: + table_name (string): table name + data (Union[Dict, List[Dict]]): data that will be inserted + partition_name (Optional[str]): limit the query to certain partition + """ + if isinstance(data, Dict): + data = [data] + + if len(data) == 0: + return + + table = Table(table_name, self.metadata_obj, autoload_with=self.engine) + + with self.engine.connect() as conn: + with conn.begin(): + if partition_name is None or partition_name == "": + conn.execute(insert(table).values(data)) + else: + conn.execute( + insert(table) + .with_hint(f"PARTITION({partition_name})") + .values(data) + ) + + def upsert( + self, + table_name: str, + data: Union[Dict, List[Dict]], + partition_name: Optional[str] = "", + ): + """Update data in table. If primary key is duplicated, replace it. + + Args: + table_name (string): table name + data (Union[Dict, List[Dict]]): data that will be upserted + partition_name (Optional[str]): limit the query to certain partition + """ + if isinstance(data, Dict): + data = [data] + + if len(data) == 0: + return + + table = Table(table_name, self.metadata_obj, autoload_with=self.engine) + + with self.engine.connect() as conn: + with conn.begin(): + upsert_stmt = ( + ReplaceStmt(table).with_hint(f"PARTITION({partition_name})") + if partition_name is not None and partition_name != "" + else ReplaceStmt(table) + ) + upsert_stmt = upsert_stmt.values(data) + conn.execute(upsert_stmt) + + def update( + self, + table_name: str, + values_clause, + where_clause=None, + partition_name: Optional[str] = "", + ): + """Update data in table. + + Args: + table_name (string): table name + values_clause: update values clause + where_clause: update with filter + partition_name (Optional[str]): limit the query to certain partition + """ + table = Table(table_name, self.metadata_obj, autoload_with=self.engine) + + with self.engine.connect() as conn: + with conn.begin(): + update_stmt = ( + update(table).with_hint(f"PARTITION({partition_name})") + if partition_name is not None and partition_name != "" + else update(table) + ) + if where_clause is not None: + update_stmt = update_stmt.where(*where_clause).values( + *values_clause + ) + else: + update_stmt = update_stmt.values(*values_clause) + conn.execute(update_stmt) + + def delete( + self, + table_name: str, + ids: Optional[Union[list, str, int]] = None, + where_clause=None, + partition_name: Optional[str] = "", + ): + """Delete data in table. + + Args: + table_name (string): table name + ids (Optional[Union[list, str, int]]): ids of data to delete + where_clause: delete with filter + partition_name (Optional[str]): limit the query to certain partition + """ + table = Table(table_name, self.metadata_obj, autoload_with=self.engine) + where_in_clause = None + if ids is not None: + primary_keys = table.primary_key + pkey_names = [column.name for column in primary_keys] + if len(pkey_names) == 1: + if isinstance(ids, list): + where_in_clause = table.c[pkey_names[0]].in_(ids) + elif isinstance(ids, (str, int)): + where_in_clause = table.c[pkey_names[0]].in_([ids]) + else: + raise TypeError("'ids' is not a list/str/int") + + with self.engine.connect() as conn: + with conn.begin(): + delete_stmt = ( + delete(table).with_hint(f"PARTITION({partition_name})") + if partition_name is not None and partition_name != "" + else delete(table) + ) + if where_in_clause is None and where_clause is None: + conn.execute(delete_stmt) + elif where_in_clause is not None and where_clause is None: + conn.execute(delete_stmt.where(where_in_clause)) + elif where_in_clause is None and where_clause is not None: + conn.execute(delete_stmt.where(*where_clause)) + else: + conn.execute( + delete_stmt.where(and_(where_in_clause, *where_clause)) + ) + + def get( + self, + table_name: str, + ids: Optional[Union[list, str, int]] = None, + where_clause=None, + output_column_name: Optional[List[str]] = None, + partition_names: Optional[List[str]] = None, + n_limits: Optional[int] = None, + ): + """Get records with specified primary field `ids`. + + Args: + table_name (string): table name + ids (Optional[Union[list, str, int]]): specified primary field values + where_clause: SQL filter + output_column_name (Optional[List[str]]): output fields name + partition_names (Optional[List[str]]): limit the query to certain partitions + n_limits (Optional[int]): limit the number of results + + Returns: + Result object from SQLAlchemy execution + """ + table = Table(table_name, self.metadata_obj, autoload_with=self.engine) + if output_column_name is not None: + columns = [table.c[column_name] for column_name in output_column_name] + stmt = select(*columns) + else: + stmt = select(table) + primary_keys = table.primary_key + pkey_names = [column.name for column in primary_keys] + where_in_clause = None + if ids is not None and len(pkey_names) == 1: + if isinstance(ids, list): + where_in_clause = table.c[pkey_names[0]].in_(ids) + elif isinstance(ids, (str, int)): + where_in_clause = table.c[pkey_names[0]].in_([ids]) + else: + raise TypeError("'ids' is not a list/str/int") + + if where_in_clause is not None and where_clause is None: + stmt = stmt.where(where_in_clause) + elif where_in_clause is None and where_clause is not None: + stmt = stmt.where(*where_clause) + elif where_in_clause is not None and where_clause is not None: + stmt = stmt.where(and_(where_in_clause, *where_clause)) + + if n_limits is not None: + stmt = stmt.limit(n_limits) + + with self.engine.connect() as conn: + with conn.begin(): + if partition_names is None: + execute_res = conn.execute(stmt) + else: + stmt_str = str(stmt.compile( + dialect=self.engine.dialect, + compile_kwargs={"literal_binds": True} + )) + stmt_str = self._insert_partition_hint_for_query_sql( + stmt_str, f"PARTITION({', '.join(partition_names)})" + ) + logging.debug(stmt_str) + execute_res = conn.execute(text(stmt_str)) + return execute_res + + def perform_raw_text_sql( + self, + text_sql: str, + ): + """Execute raw text SQL.""" + with self.engine.connect() as conn: + with conn.begin(): + return conn.execute(text(text_sql)) + + def add_columns( + self, + table_name: str, + columns: list[Column], + ): + """Add multiple columns to an existing table. + + Args: + table_name (string): table name + columns (list[Column]): list of SQLAlchemy Column objects representing the new columns + """ + compiler = self.engine.dialect.ddl_compiler(self.engine.dialect, None) + column_specs = [compiler.get_column_specification(column) for column in columns] + columns_ddl = ", ".join(f"ADD COLUMN {spec}" for spec in column_specs) + + with self.engine.connect() as conn: + with conn.begin(): + conn.execute( + text(f"ALTER TABLE `{table_name}` {columns_ddl}") + ) + + self.refresh_metadata([table_name]) + + def drop_columns( + self, + table_name: str, + column_names: list[str], + ): + """Drop multiple columns from an existing table. + + Args: + table_name (string): table name + column_names (list[str]): names of the columns to drop + """ + columns_ddl = ", ".join(f"DROP COLUMN `{name}`" for name in column_names) + + with self.engine.connect() as conn: + with conn.begin(): + conn.execute( + text(f"ALTER TABLE `{table_name}` {columns_ddl}") + ) + + self.refresh_metadata([table_name]) + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/client/partitions.py b/pyobsql-oceanbase-plugin/pyobsql/client/partitions.py new file mode 100644 index 00000000..c3eed133 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/client/partitions.py @@ -0,0 +1,407 @@ +"""A module to do compilation of OceanBase Parition Clause.""" +from typing import List, Optional, Union +import logging +from dataclasses import dataclass +from .enum import IntEnum +from .exceptions import * + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class PartType(IntEnum): + """Partition type of table or collection for both ObVecClient and MilvusLikeClient""" + Range = 0 + Hash = 1 + Key = 2 + List = 3 + RangeColumns = 4 + ListColumns = 5 + + +class ObPartition: + """Base class of all kind of Partition strategy + + Attributes: + part_type (PartType) : type of partition strategy + sub_partition (ObPartition) : subpartition strategy + is_sub (bool) : this partition strategy is a subpartition or not + """ + def __init__(self, part_type: PartType): + self.part_type = part_type + self.sub_partition = None + self.is_sub = False + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + raise NotImplementedError() + + def add_subpartition(self, sub_part): + """Add subpartition strategy to current partition. + + Args: + sub_part (ObPartition) : subpartition strategy + """ + if self.is_sub: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionLevelMoreThanTwo, + ) + + if sub_part is None: + return + + if not sub_part.is_sub: + raise ValueError("not a subparition") + + self.sub_partition = sub_part + + +@dataclass +class RangeListPartInfo: + """Range/RangeColumns/List/ListColumns partition info for each partition. + + Attributes: + part_name (string) : partition name + part_upper_bound_expr (Union[List, str, int]) : + For example, using `[1,2]`/`'DEFAULT'` as default case/`7` when create + List/ListColumns partition. + Using 100 / `MAXVALUE` when create Range/RangeColumns partition. + """ + part_name: str + part_upper_bound_expr: Union[List, str, int] + + def get_part_expr_str(self): + """Parse part_upper_bound_expr to text SQL.""" + if isinstance(self.part_upper_bound_expr, List): + return ",".join([str(v) for v in self.part_upper_bound_expr]) + if isinstance(self.part_upper_bound_expr, str): + return self.part_upper_bound_expr + if isinstance(self.part_upper_bound_expr, int): + return str(self.part_upper_bound_expr) + raise ValueError("Invalid datatype") + + +class ObRangePartition(ObPartition): + """Range/RangeColumns partition strategy.""" + def __init__( + self, + is_range_columns: bool, + range_part_infos: List[RangeListPartInfo], + range_expr: Optional[str] = None, + col_name_list: Optional[List[str]] = None, + ): + super().__init__(PartType.RangeColumns if is_range_columns else PartType.Range) + self.range_part_infos = range_part_infos + self.range_expr = range_expr + self.col_name_list = col_name_list + + if not is_range_columns and range_expr is None: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionRangeExprMissing, + ) + + if is_range_columns and col_name_list is None: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionRangeColNameListMissing, + ) + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"PARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_type == PartType.Range: + assert self.range_expr is not None + if self.sub_partition is None: + return f"RANGE ({self.range_expr}) ({self._parse_range_part_list()})" + return f"RANGE ({self.range_expr}) {self.sub_partition.do_compile()} " \ + f"({self._parse_range_part_list()})" + assert self.col_name_list is not None + if self.sub_partition is None: + return f"RANGE COLUMNS ({','.join(self.col_name_list)}) " \ + f"({self._parse_range_part_list()})" + return f"RANGE COLUMNS ({','.join(self.col_name_list)}) " \ + f"{self.sub_partition.do_compile()} ({self._parse_range_part_list()})" + + def _parse_range_part_list(self) -> str: + range_partitions_complied = [ + f"PARTITION {range_part_info.part_name} VALUES LESS THAN " \ + f"({range_part_info.get_part_expr_str()})" + for range_part_info in self.range_part_infos + ] + return ",".join(range_partitions_complied) + + +class ObSubRangePartition(ObRangePartition): + """Range/RangeColumns subpartition strategy.""" + def __init__( + self, + is_range_columns: bool, + range_part_infos: List[RangeListPartInfo], + range_expr: Optional[str] = None, + col_name_list: Optional[List[str]] = None, + ): + super().__init__(is_range_columns, range_part_infos, range_expr, col_name_list) + self.is_sub = True + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"SUBPARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_type == PartType.Range: + assert self.range_expr is not None + assert self.sub_partition is None + return f"RANGE ({self.range_expr}) SUBPARTITION TEMPLATE " \ + f"({self._parse_range_part_list()})" + assert self.col_name_list is not None + assert self.sub_partition is None + return f"RANGE COLUMNS ({','.join(self.col_name_list)}) SUBPARTITION TEMPLATE " \ + f"({self._parse_range_part_list()})" + + def _parse_range_part_list(self) -> str: + range_partitions_complied = [ + f"SUBPARTITION {range_part_info.part_name} VALUES LESS THAN " \ + f"({range_part_info.get_part_expr_str()})" + for range_part_info in self.range_part_infos + ] + return ",".join(range_partitions_complied) + + +class ObListPartition(ObPartition): + """List/ListColumns partition strategy.""" + def __init__( + self, + is_list_columns: bool, + list_part_infos: List[RangeListPartInfo], + list_expr: Optional[str] = None, + col_name_list: Optional[List[str]] = None, + ): + super().__init__(PartType.ListColumns if is_list_columns else PartType.List) + self.list_part_infos = list_part_infos + self.list_expr = list_expr + self.col_name_list = col_name_list + + if not is_list_columns and list_expr is None: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionListExprMissing, + ) + + if is_list_columns and col_name_list is None: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionListColNameListMissing, + ) + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"PARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_type == PartType.List: + assert self.list_expr is not None + if self.sub_partition is None: + return f"LIST ({self.list_expr}) ({self._parse_list_part_list()})" + return f"LIST ({self.list_expr}) {self.sub_partition.do_compile()} " \ + f"({self._parse_list_part_list()})" + assert self.col_name_list is not None + if self.sub_partition is None: + return f"LIST COLUMNS ({','.join(self.col_name_list)}) " \ + f"({self._parse_list_part_list()})" + return f"LIST COLUMNS ({','.join(self.col_name_list)}) " \ + f"{self.sub_partition.do_compile()} ({self._parse_list_part_list()})" + + def _parse_list_part_list(self) -> str: + list_partitions_complied = [ + f"PARTITION {list_part_info.part_name} VALUES IN ({list_part_info.get_part_expr_str()})" + for list_part_info in self.list_part_infos + ] + return ",".join(list_partitions_complied) + + +class ObSubListPartition(ObListPartition): + """List/ListColumns subpartition strategy.""" + def __init__( + self, + is_list_columns: bool, + list_part_infos: List[RangeListPartInfo], + list_expr: Optional[str] = None, + col_name_list: Optional[List[str]] = None, + ): + super().__init__(is_list_columns, list_part_infos, list_expr, col_name_list) + self.is_sub = True + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"SUBPARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_type == PartType.List: + assert self.list_expr is not None + assert self.sub_partition is None + return f"LIST ({self.list_expr}) SUBPARTITION TEMPLATE ({self._parse_list_part_list()})" + assert self.col_name_list is not None + assert self.sub_partition is None + return f"LIST COLUMNS ({','.join(self.col_name_list)}) SUBPARTITION TEMPLATE " \ + f"({self._parse_list_part_list()})" + + def _parse_list_part_list(self) -> str: + list_partitions_complied = [ + f"SUBPARTITION {list_part_info.part_name} VALUES IN " \ + f"({list_part_info.get_part_expr_str()})" + for list_part_info in self.list_part_infos + ] + return ",".join(list_partitions_complied) + + +class ObHashPartition(ObPartition): + """Hash partition strategy.""" + def __init__( + self, + hash_expr: str, + hash_part_name_list: List[str] = None, + part_count: Optional[int] = None, + ): + super().__init__(PartType.Hash) + self.hash_expr = hash_expr + self.hash_part_name_list = hash_part_name_list + self.part_count = part_count + + if self.hash_part_name_list is None and self.part_count is None: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionHashNameListAndPartCntMissing, + ) + + if self.part_count is not None and self.hash_part_name_list is not None: + logging.warning( + "part_count & hash_part_name_list are both set, " \ + "hash_part_name_list will be override by part_count" + ) + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"PARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_count is not None: + if self.sub_partition is None: + return f"HASH ({self.hash_expr}) PARTITIONS {self.part_count}" + return f"HASH ({self.hash_expr}) {self.sub_partition.do_compile()} " \ + f"PARTITIONS {self.part_count}" + assert self.hash_part_name_list is not None + if self.sub_partition is None: + return f"HASH ({self.hash_expr}) ({self._parse_hash_part_list()})" + return f"HASH ({self.hash_expr}) {self.sub_partition.do_compile()} " \ + f"({self._parse_hash_part_list()})" + + def _parse_hash_part_list(self): + return ",".join([f"PARTITION {name}" for name in self.hash_part_name_list]) + + +class ObSubHashPartition(ObHashPartition): + """Hash subpartition strategy.""" + def __init__( + self, + hash_expr: str, + hash_part_name_list: List[str] = None, + part_count: Optional[int] = None, + ): + super().__init__(hash_expr, hash_part_name_list, part_count) + self.is_sub = True + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"SUBPARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_count is not None: + assert self.sub_partition is None + return f"HASH ({self.hash_expr}) SUBPARTITIONS {self.part_count}" + assert self.hash_part_name_list is not None + assert self.sub_partition is None + return f"HASH ({self.hash_expr}) SUBPARTITION TEMPLATE ({self._parse_hash_part_list()})" + + def _parse_hash_part_list(self): + return ",".join([f"SUBPARTITION {name}" for name in self.hash_part_name_list]) + + +class ObKeyPartition(ObPartition): + """Key partition strategy.""" + def __init__( + self, + col_name_list: List[str], + key_part_name_list: List[str] = None, + part_count: Optional[int] = None, + ): + super().__init__(PartType.Key) + self.col_name_list = col_name_list + self.key_part_name_list = key_part_name_list + self.part_count = part_count + + if self.key_part_name_list is None and self.part_count is None: + raise PartitionFieldException( + code=ErrorCode.INVALID_ARGUMENT, + message=ExceptionsMessage.PartitionKeyNameListAndPartCntMissing, + ) + + if self.part_count is not None and self.key_part_name_list is not None: + logging.warning( + "part_count & key_part_name_list are both set, " \ + "key_part_name_list will be override by part_count" + ) + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"PARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_count is not None: + if self.sub_partition is None: + return ( + f"KEY ({','.join(self.col_name_list)}) PARTITIONS {self.part_count}" + ) + return f"KEY ({','.join(self.col_name_list)}) {self.sub_partition.do_compile()} " \ + f"PARTITIONS {self.part_count}" + assert self.key_part_name_list is not None + if self.sub_partition is None: + return f"KEY ({','.join(self.col_name_list)}) ({self._parse_key_part_list()})" + return f"KEY ({','.join(self.col_name_list)}) {self.sub_partition.do_compile()} " \ + f"({self._parse_key_part_list()})" + + def _parse_key_part_list(self): + return ",".join([f"PARTITION {name}" for name in self.key_part_name_list]) + + +class ObSubKeyPartition(ObKeyPartition): + """Key subpartition strategy.""" + def __init__( + self, + col_name_list: List[str], + key_part_name_list: List[str] = None, + part_count: Optional[int] = None, + ): + super().__init__(col_name_list, key_part_name_list, part_count) + self.is_sub = True + + def do_compile(self) -> str: + """Compile partition strategy to text SQL.""" + return f"SUBPARTITION BY {self._compile_helper()}" + + def _compile_helper(self) -> str: + if self.part_count is not None: + assert self.sub_partition is None + return ( + f"KEY ({','.join(self.col_name_list)}) SUBPARTITIONS {self.part_count}" + ) + assert self.key_part_name_list is not None + assert self.sub_partition is None + return f"KEY ({','.join(self.col_name_list)}) SUBPARTITION TEMPLATE " \ + f"({self._parse_key_part_list()})" + + def _parse_key_part_list(self): + return ",".join([f"SUBPARTITION {name}" for name in self.key_part_name_list]) diff --git a/pyobsql-oceanbase-plugin/pyobsql/json_table/__init__.py b/pyobsql-oceanbase-plugin/pyobsql/json_table/__init__.py new file mode 100644 index 00000000..9d718f8a --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/json_table/__init__.py @@ -0,0 +1,25 @@ +from .oceanbase_dialect import OceanBase, ChangeColumn +from .virtual_data_type import ( + JType, + JsonTableDataType, + JsonTableBool, + JsonTableTimestamp, + JsonTableVarcharFactory, + JsonTableDecimalFactory, + JsonTableInt, + val2json, +) +from .json_value_returning_func import json_value + +__all__ = [ + "OceanBase", "ChangeColumn", + "JType", + "JsonTableDataType", + "JsonTableBool", + "JsonTableTimestamp", + "JsonTableVarcharFactory", + "JsonTableDecimalFactory", + "JsonTableInt", + "val2json", + "json_value" +] diff --git a/pyobsql-oceanbase-plugin/pyobsql/json_table/json_value_returning_func.py b/pyobsql-oceanbase-plugin/pyobsql/json_table/json_value_returning_func.py new file mode 100644 index 00000000..18f7f9f6 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/json_table/json_value_returning_func.py @@ -0,0 +1,51 @@ +import logging +import re +from typing import Tuple + +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.functions import FunctionElement +from sqlalchemy import BINARY, Float, Boolean, Text + +logger = logging.getLogger(__name__) + +class json_value(FunctionElement): + type = Text() + inherit_cache = True + + def __init__(self, *args): + super().__init__() + self.args = args + +@compiles(json_value) +def compile_json_value(element, compiler, **kwargs): + args = [] + if len(element.args) != 3: + raise ValueError("Number of args for json_value should be 3") + args.append(compiler.process(element.args[0])) + if not (isinstance(element.args[1], str) and isinstance(element.args[2], str)): + raise ValueError("Invalid args for json_value") + + if element.args[2].startswith('TINYINT'): + returning_type = "SIGNED" + elif element.args[2].startswith('TIMESTAMP'): + returning_type = "DATETIME" + elif element.args[2].startswith('INT'): + returning_type = "SIGNED" + elif element.args[2].startswith('VARCHAR'): + if element.args[2] == 'VARCHAR': + returning_type = "CHAR(255)" + else: + varchar_pattern = r'VARCHAR\((\d+)\)' + varchar_matches = re.findall(varchar_pattern, element.args[2]) + returning_type = f"CHAR({int(varchar_matches[0])})" + elif element.args[2].startswith('DECIMAL'): + if element.args[2] == 'DECIMAL': + returning_type = "DECIMAL(10, 0)" + else: + decimal_pattern = r'DECIMAL\((\d+),\s*(\d+)\)' + decimal_matches = re.findall(decimal_pattern, element.args[2]) + x, y = decimal_matches[0] + returning_type = f"DECIMAL({x}, {y})" + args.append(f"'{element.args[1]}' RETURNING {returning_type}") + args = ", ".join(args) + return f"json_value({args})" diff --git a/pyobsql-oceanbase-plugin/pyobsql/json_table/oceanbase_dialect.py b/pyobsql-oceanbase-plugin/pyobsql/json_table/oceanbase_dialect.py new file mode 100644 index 00000000..0aa4f4c9 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/json_table/oceanbase_dialect.py @@ -0,0 +1,116 @@ +import typing as t +from sqlglot import parser, exp, Expression +from sqlglot.dialects.mysql import MySQL +from sqlglot.tokens import TokenType + +class ChangeColumn(Expression): + arg_types = { + "this": True, + "origin_col_name": True, + "dtype": True, + } + + @property + def origin_col_name(self) -> str: + origin_col_name = self.args.get("origin_col_name") + return origin_col_name + + @property + def dtype(self) -> Expression: + dtype = self.args.get("dtype") + return dtype + +class OceanBase(MySQL): + class Parser(MySQL.Parser): + ALTER_PARSERS = { + **parser.Parser.ALTER_PARSERS, + "MODIFY": lambda self: self._parse_alter_table_alter(), + "CHANGE": lambda self: self._parse_change_table_column(), + } + + def _parse_alter_table_alter(self) -> t.Optional[exp.Expression]: + if self._match_texts(self.ALTER_ALTER_PARSERS): + return self.ALTER_ALTER_PARSERS[self._prev.text.upper()](self) + + self._match(TokenType.COLUMN) + column = self._parse_field_def() + + if self._match_pair(TokenType.DROP, TokenType.DEFAULT): + return self.expression(exp.AlterColumn, this=column, drop=True) + if self._match_pair(TokenType.SET, TokenType.DEFAULT): + return self.expression(exp.AlterColumn, this=column, default=self._parse_assignment()) + if self._match(TokenType.COMMENT): + return self.expression(exp.AlterColumn, this=column, comment=self._parse_string()) + if self._match_text_seq("DROP", "NOT", "NULL"): + return self.expression( + exp.AlterColumn, + this=column, + drop=True, + allow_null=True, + ) + if self._match_text_seq("SET", "NOT", "NULL"): + return self.expression( + exp.AlterColumn, + this=column, + allow_null=False, + ) + self._match_text_seq("SET", "DATA") + self._match_text_seq("TYPE") + return self.expression( + exp.AlterColumn, + this=column, + dtype=self._parse_types(), + collate=self._match(TokenType.COLLATE) and self._parse_term(), + using=self._match(TokenType.USING) and self._parse_assignment(), + ) + + def _parse_drop(self, exists: bool = False) -> t.Union[exp.Drop, exp.Command]: + temporary = self._match(TokenType.TEMPORARY) + materialized = self._match_text_seq("MATERIALIZED") + + kind = self._match_set(self.CREATABLES) and self._prev.text.upper() + if not kind: + kind = "COLUMN" + + concurrently = self._match_text_seq("CONCURRENTLY") + if_exists = exists or self._parse_exists() + + if kind == "COLUMN": + this = self._parse_column() + else: + this = self._parse_table_parts( + schema=True, is_db_reference=self._prev.token_type == TokenType.SCHEMA + ) + + cluster = self._parse_on_property() if self._match(TokenType.ON) else None + + if self._match(TokenType.L_PAREN, advance=False): + expressions = self._parse_wrapped_csv(self._parse_types) + else: + expressions = None + + return self.expression( + exp.Drop, + exists=if_exists, + this=this, + expressions=expressions, + kind=self.dialect.CREATABLE_KIND_MAPPING.get(kind) or kind, + temporary=temporary, + materialized=materialized, + cascade=self._match_text_seq("CASCADE"), + constraints=self._match_text_seq("CONSTRAINTS"), + purge=self._match_text_seq("PURGE"), + cluster=cluster, + concurrently=concurrently, + ) + + def _parse_change_table_column(self) -> t.Optional[exp.Expression]: + self._match(TokenType.COLUMN) + origin_col = self._parse_field(any_token=True) + column = self._parse_field() + return self.expression( + ChangeColumn, + this=column, + origin_col_name=origin_col, + dtype=self._parse_types(), + ) diff --git a/pyobsql-oceanbase-plugin/pyobsql/json_table/virtual_data_type.py b/pyobsql-oceanbase-plugin/pyobsql/json_table/virtual_data_type.py new file mode 100644 index 00000000..57348238 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/json_table/virtual_data_type.py @@ -0,0 +1,114 @@ +from datetime import datetime +from decimal import Decimal, InvalidOperation, ROUND_DOWN +from enum import Enum +from typing import Optional +from typing_extensions import Annotated + +from pydantic import BaseModel, Field, AfterValidator, create_model + + +class IntEnum(int, Enum): + """Int type enumerate definition.""" + +class JType(IntEnum): + J_BOOL = 1 + J_TIMESTAMP = 2 + J_VARCHAR = 3 + J_DECIMAL = 4 + J_INT = 5 + +class JsonTableDataType(BaseModel): + type: JType + +class JsonTableBool(JsonTableDataType): + type: JType = Field(default=JType.J_BOOL) + val: Optional[bool] + +class JsonTableTimestamp(JsonTableDataType): + type: JType = Field(default=JType.J_TIMESTAMP) + val: Optional[datetime] + +def check_varchar_len_with_length(length: int): + def check_varchar_len(x: Optional[str]): + if x is None: + return None + if len(x) > length: + raise ValueError(f'{x} is longer than {length}') + return x + + return check_varchar_len + +class JsonTableVarcharFactory: + def __init__(self, length: int): + self.length = length + + def get_json_table_varchar_type(self): + model_name = f"JsonTableVarchar{self.length}" + fields = { + 'type': (JType, JType.J_VARCHAR), + 'val': (Annotated[Optional[str], AfterValidator(check_varchar_len_with_length(self.length))], ...) + } + return create_model( + model_name, + __base__=JsonTableDataType, + **fields + ) + +def check_and_parse_decimal(x: int, y: int): + def check_float(v): + if v is None: + return None + try: + decimal_value = Decimal(v) + except InvalidOperation: + raise ValueError(f"Value {v} cannot be converted to Decimal.") + + decimal_str = str(decimal_value).strip() + + if '.' in decimal_str: + integer_part, decimal_part = decimal_str.split('.') + else: + integer_part, decimal_part = decimal_str, '' + + integer_count = len(integer_part.lstrip('-')) # Remove the negative sign length + decimal_count = len(decimal_part) + + if integer_count + min(decimal_count, y) > x: + raise ValueError(f"'{v}' Range out of Decimal({x}, {y})") + + if decimal_count > y: + quantize_str = '1.' + '0' * y + decimal_value = decimal_value.quantize(Decimal(quantize_str), rounding=ROUND_DOWN) + return decimal_value + return check_float + +class JsonTableDecimalFactory: + def __init__(self, ndigits: int, decimal_p: int): + self.ndigits = ndigits + self.decimal_p = decimal_p + + def get_json_table_decimal_type(self): + model_name = f"JsonTableDecimal_{self.ndigits}_{self.decimal_p}" + fields = { + 'type': (JType, JType.J_DECIMAL), + 'val': (Annotated[Optional[float], AfterValidator(check_and_parse_decimal(self.ndigits, self.decimal_p))], ...) + } + return create_model( + model_name, + __base__=JsonTableDataType, + **fields + ) + +class JsonTableInt(JsonTableDataType): + type: JType = Field(default=JType.J_INT) + val: Optional[int] + +def val2json(val): + if val is None: + return None + if isinstance(val, int) or isinstance(val, bool) or isinstance(val, str): + return val + if isinstance(val, datetime): + return val.isoformat() + if isinstance(val, Decimal): + return float(val) diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/__init__.py b/pyobsql-oceanbase-plugin/pyobsql/schema/__init__.py new file mode 100644 index 00000000..eafed45c --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/__init__.py @@ -0,0 +1,50 @@ +"""A extension for SQLAlchemy for OceanBase SQL schema definition. + +* ARRAY An extended data type in SQLAlchemy +* VECTOR An extended data type in SQLAlchemy for vector storage +* SPARSE_VECTOR An extended data type in SQLAlchemy for sparse vector storage +* POINT GIS Point data type +* ObTable Extension to Table for creating table +* ReplaceStmt Replace into statement +* FtsIndex Full Text Search Index +* CreateFtsIndex Full Text Search Index Creation statement clause +* MatchAgainst Full Text Search clause +* ST_GeomFromText GIS function: parse text to geometry object +* st_distance GIS function: calculate distance between Points +* st_dwithin GIS function: check if the distance between two points +* st_astext GIS function: return a Point in human-readable format +* OceanBaseDialect OceanBase SQLAlchemy dialect +* AsyncOceanBaseDialect OceanBase async SQLAlchemy dialect +""" +from .array import ARRAY +from .vector import VECTOR +from .sparse_vector import SPARSE_VECTOR +from .geo_srid_point import POINT +from .ob_table import ObTable +from .replace_stmt import ReplaceStmt +from .dialect import OceanBaseDialect, AsyncOceanBaseDialect +from .full_text_index import FtsIndex, CreateFtsIndex +from .match_against_func import MatchAgainst +from .gis_func import ST_GeomFromText, st_distance, st_dwithin, st_astext + +__all__ = [ + "ARRAY", + "VECTOR", + "SPARSE_VECTOR", + "POINT", + "ObTable", + "ReplaceStmt", + "FtsIndex", + "CreateFtsIndex", + "MatchAgainst", + "ST_GeomFromText", + "st_distance", + "st_dwithin", + "st_astext", + "OceanBaseDialect", + "AsyncOceanBaseDialect", +] + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/array.py b/pyobsql-oceanbase-plugin/pyobsql/schema/array.py new file mode 100644 index 00000000..52f8d1a9 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/array.py @@ -0,0 +1,122 @@ +"""ARRAY: An extended data type for SQLAlchemy""" +import json +from typing import Any, List, Optional, Sequence, Union + +from sqlalchemy.sql.type_api import TypeEngine +from sqlalchemy.types import UserDefinedType, String + + +class ARRAY(UserDefinedType): + """ARRAY data type definition with support for up to 6 levels of nesting.""" + cache_ok = True + _string = String() + + def __init__(self, item_type: Union[TypeEngine, type]): + """Construct an ARRAY. + + Args: + item_type: The data type of items in this array. For nested arrays, + pass another ARRAY type. + """ + super(UserDefinedType, self).__init__() + if isinstance(item_type, type): + item_type = item_type() + self.item_type = item_type + if isinstance(item_type, ARRAY): + self.dim = item_type.dim + 1 + else: + self.dim = 1 + if self.dim > 6: + raise ValueError("Maximum nesting level of 6 exceeded") + + def get_col_spec(self, **kw): # pylint: disable=unused-argument + """Parse to array data type definition in text SQL.""" + if hasattr(self.item_type, 'get_col_spec'): + base_type = self.item_type.get_col_spec(**kw) + else: + base_type = str(self.item_type) + return f"ARRAY({base_type})" + + def _get_list_depth(self, value: Any) -> int: + if not isinstance(value, list): + return 0 + max_depth = 0 + for element in value: + current_depth = self._get_list_depth(element) + if current_depth > max_depth: + max_depth = current_depth + return 1 + max_depth + + def _validate_dimension(self, value: list[Any]): + arr_depth = self._get_list_depth(value) + assert arr_depth == self.dim, "Array dimension mismatch, expected {}, got {}".format(self.dim, arr_depth) + + def bind_processor(self, dialect): + item_type = self.item_type + while isinstance(item_type, ARRAY): + item_type = item_type.item_type + + item_proc = item_type.dialect_impl(dialect).bind_processor(dialect) + + def process(value: Optional[Sequence[Any] | str]) -> Optional[str]: + if value is None: + return None + if isinstance(value, str): + self._validate_dimension(json.loads(value)) + return value + + def convert(val): + if isinstance(val, (list, tuple)): + return [convert(v) for v in val] + if item_proc: + return item_proc(val) + return val + + processed = convert(value) + self._validate_dimension(processed) + return json.dumps(processed) + + return process + + def result_processor(self, dialect, coltype): + item_type = self.item_type + while isinstance(item_type, ARRAY): + item_type = item_type.item_type + + item_proc = item_type.dialect_impl(dialect).result_processor(dialect, coltype) + + def process(value: Optional[str]) -> Optional[List[Any]]: + if value is None: + return None + + def convert(val): + if isinstance(val, (list, tuple)): + return [convert(v) for v in val] + if item_proc: + return item_proc(val) + return val + + value = json.loads(value) if isinstance(value, str) else value + return convert(value) + + return process + + def literal_processor(self, dialect): + item_type = self.item_type + while isinstance(item_type, ARRAY): + item_type = item_type.item_type + + item_proc = item_type.dialect_impl(dialect).literal_processor(dialect) + + def process(value: Sequence[Any]) -> str: + def convert(val): + if isinstance(val, (list, tuple)): + return [convert(v) for v in val] + if item_proc: + return item_proc(val) + return val + + processed = convert(value) + return json.dumps(processed) + + return process diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/dialect.py b/pyobsql-oceanbase-plugin/pyobsql/schema/dialect.py new file mode 100644 index 00000000..bac72d24 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/dialect.py @@ -0,0 +1,63 @@ +"""OceanBase dialect.""" +from sqlalchemy import util +from sqlalchemy.dialects.mysql import aiomysql, pymysql + +from .reflection import OceanBaseTableDefinitionParser +from .vector import VECTOR +from .sparse_vector import SPARSE_VECTOR +from .geo_srid_point import POINT + +class OceanBaseDialect(pymysql.MySQLDialect_pymysql): + # not change dialect name, since it is a subclass of pymysql.MySQLDialect_pymysql + # name = "oceanbase" + """OceanBase dialect.""" + supports_statement_cache = True + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.ischema_names["VECTOR"] = VECTOR + self.ischema_names["SPARSEVECTOR"] = SPARSE_VECTOR + self.ischema_names["point"] = POINT + + @util.memoized_property + def _tabledef_parser(self): + """return the MySQLTableDefinitionParser, generate if needed. + + The deferred creation ensures that the dialect has + retrieved server version information first. + + """ + preparer = self.identifier_preparer + default_schema = self.default_schema_name + return OceanBaseTableDefinitionParser( + self, preparer, default_schema=default_schema + ) + + +class AsyncOceanBaseDialect(aiomysql.MySQLDialect_aiomysql): + """OceanBase async dialect.""" + supports_statement_cache = True + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.ischema_names["VECTOR"] = VECTOR + self.ischema_names["SPARSEVECTOR"] = SPARSE_VECTOR + self.ischema_names["point"] = POINT + + @util.memoized_property + def _tabledef_parser(self): + """return the MySQLTableDefinitionParser, generate if needed. + + The deferred creation ensures that the dialect has + retrieved server version information first. + + """ + preparer = self.identifier_preparer + default_schema = self.default_schema_name + return OceanBaseTableDefinitionParser( + self, preparer, default_schema=default_schema + ) + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/full_text_index.py b/pyobsql-oceanbase-plugin/pyobsql/schema/full_text_index.py new file mode 100644 index 00000000..950d25f2 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/full_text_index.py @@ -0,0 +1,59 @@ +"""FullTextIndex: full text search index type""" +from sqlalchemy import Index +from sqlalchemy.schema import DDLElement +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.ddl import SchemaGenerator + + +class CreateFtsIndex(DDLElement): + """A new statement clause to create fts index. + + Attributes: + index : fts index schema + """ + def __init__(self, index): + self.index = index + + +class ObFtsSchemaGenerator(SchemaGenerator): + """A new schema generator to handle create fts index statement.""" + def visit_fts_index(self, index, create_ok=False): + """Handle create fts index statement compiling. + + Args: + index: fts index schema + create_ok: the schema is created or not + """ + if not create_ok and not self._can_create_index(index): + return + with self.with_ddl_events(index): + CreateFtsIndex(index)._invoke_with(self.connection) + +class FtsIndex(Index): + """Fts Index schema.""" + __visit_name__ = "fts_index" + + def __init__(self, name, fts_parser: str, *column_names, **kw): + self.fts_parser = fts_parser + super().__init__(name, *column_names, **kw) + + def create(self, bind, checkfirst: bool = False) -> None: + """Create fts index. + + Args: + bind: SQL engine or connection. + checkfirst: check the index exists or not. + """ + bind._run_ddl_visitor(ObFtsSchemaGenerator, self, checkfirst=checkfirst) + + +@compiles(CreateFtsIndex) +def compile_create_fts_index(element, compiler, **kw): # pylint: disable=unused-argument + """A decorator function to compile create fts index statement.""" + index = element.index + table_name = index.table.name + column_list = ", ".join([column.name for column in index.columns]) + fts_parser = index.fts_parser + if fts_parser is not None: + return f"CREATE FULLTEXT INDEX {index.name} ON {table_name} ({column_list}) WITH PARSER {fts_parser}" + return f"CREATE FULLTEXT INDEX {index.name} ON {table_name} ({column_list})" diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/geo_srid_point.py b/pyobsql-oceanbase-plugin/pyobsql/schema/geo_srid_point.py new file mode 100644 index 00000000..74642bbb --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/geo_srid_point.py @@ -0,0 +1,38 @@ +"""Point: OceanBase GIS data type for SQLAlchemy""" +from typing import Tuple, Optional +from sqlalchemy.types import UserDefinedType, String + +class POINT(UserDefinedType): + """Point data type definition.""" + cache_ok = True + _string = String() + + def __init__( + self, + # lat_long: Tuple[float, float], + srid: Optional[int] = None + ): + """Init Latitude and Longitude.""" + super(UserDefinedType, self).__init__() + # self.lat_long = lat_long + self.srid = srid + + def get_col_spec(self, **kw): # pylint: disable=unused-argument + """Parse to Point data type definition in text SQL.""" + if self.srid is None: + return "POINT" + return f"POINT SRID {self.srid}" + + @classmethod + def to_db(cls, value: Tuple[float, float]): + """Parse tuple to POINT literal""" + return f"POINT({value[0]} {value[1]})" + + def bind_processor(self, dialect): + raise ValueError("Never access Point directly.") + + def literal_processor(self, dialect): + raise ValueError("Never access Point directly.") + + def result_processor(self, dialect, coltype): + raise ValueError("Never access Point directly.") diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/gis_func.py b/pyobsql-oceanbase-plugin/pyobsql/schema/gis_func.py new file mode 100644 index 00000000..a403c4e8 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/gis_func.py @@ -0,0 +1,110 @@ +"""gis_func: An extended system function in GIS.""" + +import logging +from typing import Tuple + +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.functions import FunctionElement +from sqlalchemy import BINARY, Float, Boolean, Text + +from .geo_srid_point import POINT + +logger = logging.getLogger(__name__) + +class ST_GeomFromText(FunctionElement): + """ST_GeomFromText: parse text to geometry object. + + Attributes: + type : result type + """ + type = BINARY() + + def __init__(self, *args): + super().__init__() + self.args = args + +@compiles(ST_GeomFromText) +def compile_ST_GeomFromText(element, compiler, **kwargs): # pylint: disable=unused-argument + """Compile ST_GeomFromText function.""" + args = [] + for idx, arg in enumerate(element.args): + if idx == 0: + if ( + (not isinstance(arg, Tuple)) or + (len(arg) != 2) or + (not all(isinstance(x, float) for x in arg)) + ): + raise ValueError( + f"Tuple[float, float] is expected for Point literal," \ + f"while get {type(arg)}" + ) + args.append(f"'{POINT.to_db(arg)}'") + else: + args.append(str(arg)) + args_str = ", ".join(args) + # logger.info(f"{args_str}") + return f"ST_GeomFromText({args_str})" + +class st_distance(FunctionElement): + """st_distance: calculate distance between Points. + + Attributes: + type : result type + """ + type = Float() + inherit_cache = True + + def __init__(self, *args): + super().__init__() + self.args = args + +@compiles(st_distance) +def compile_st_distance(element, compiler, **kwargs): # pylint: disable=unused-argument + """Compile st_distance function.""" + args = ", ".join(compiler.process(arg) for arg in element.args) + return f"st_distance({args})" + +class st_dwithin(FunctionElement): + """st_dwithin: Checks if the distance between two points + is less than a specified distance. + + Attributes: + type : result type + """ + type = Boolean() + inherit_cache = True + + def __init__(self, *args): + super().__init__() + self.args = args + +@compiles(st_dwithin) +def compile_st_dwithin(element, compiler, **kwargs): # pylint: disable=unused-argument + """Compile st_dwithin function.""" + args = [] + for idx, arg in enumerate(element.args): + if idx == 2: + args.append(str(arg)) + else: + args.append(compiler.process(arg)) + args_str = ", ".join(args) + return f"_st_dwithin({args_str})" + +class st_astext(FunctionElement): + """st_astext: Returns a Point in human-readable format. + + Attributes: + type : result type + """ + type = Text() + inherit_cache = True + + def __init__(self, *args): + super().__init__() + self.args = args + +@compiles(st_astext) +def compile_st_astext(element, compiler, **kwargs): # pylint: disable=unused-argument + """Compile st_astext function.""" + args = ", ".join(compiler.process(arg) for arg in element.args) + return f"st_astext({args})" diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/match_against_func.py b/pyobsql-oceanbase-plugin/pyobsql/schema/match_against_func.py new file mode 100644 index 00000000..a522f81c --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/match_against_func.py @@ -0,0 +1,38 @@ +"""match_against_func: An extend system function in FTS.""" + +import logging + +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.functions import FunctionElement +from sqlalchemy import literal, column + +logger = logging.getLogger(__name__) + +class MatchAgainst(FunctionElement): + """MatchAgainst: match clause for full text search. + + Attributes: + type : result type + """ + inherit_cache = True + + def __init__(self, query, *columns): + columns = [column(col) if isinstance(col, str) else col for col in columns] + super().__init__(literal(query), *columns) + +@compiles(MatchAgainst) +def complie_MatchAgainst(element, compiler, **kwargs): # pylint: disable=unused-argument + """Compile MatchAgainst function.""" + clauses = list(element.clauses) + if len(clauses) < 2: + raise ValueError( + f"MatchAgainst should take a string expression and " \ + f"at least one column name string as parameters." + ) + + query_expr = clauses[0] + compiled_query = compiler.process(query_expr, **kwargs) + column_exprs = clauses[1:] + compiled_columns = [compiler.process(col, identifier_prepared=True) for col in column_exprs] + columns_str = ", ".join(compiled_columns) + return f"MATCH ({columns_str}) AGAINST ({compiled_query} IN NATURAL LANGUAGE MODE)" diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/ob_table.py b/pyobsql-oceanbase-plugin/pyobsql/schema/ob_table.py new file mode 100644 index 00000000..11a06d38 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/ob_table.py @@ -0,0 +1,24 @@ +"""ObTable: extension to Table for OceanBase-specific features.""" +from sqlalchemy import Table +from sqlalchemy.sql.ddl import SchemaGenerator + + +class ObSchemaGenerator(SchemaGenerator): + """Schema generator for ObTable (simplified version without vector index support).""" + pass + + +class ObTable(Table): + """A class extends SQLAlchemy Table for OceanBase-specific table creation.""" + def create(self, bind, checkfirst: bool = False) -> None: + """Create table with OceanBase-specific features. + + Args: + bind: SQL engine or connection + checkfirst: check if table exists before creating + """ + bind._run_ddl_visitor(ObSchemaGenerator, self, checkfirst=checkfirst) + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/reflection.py b/pyobsql-oceanbase-plugin/pyobsql/schema/reflection.py new file mode 100644 index 00000000..7b7426a9 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/reflection.py @@ -0,0 +1,157 @@ +"""OceanBase table definition reflection.""" +import re +import logging +from sqlalchemy.dialects.mysql.reflection import MySQLTableDefinitionParser, _re_compile, cleanup_text + +from .array import ARRAY + +logger = logging.getLogger(__name__) + +class OceanBaseTableDefinitionParser(MySQLTableDefinitionParser): + """OceanBase table definition parser.""" + def __init__(self, dialect, preparer, *, default_schema=None): + MySQLTableDefinitionParser.__init__(self, dialect, preparer) + self.default_schema = default_schema + + def _prep_regexes(self): + super()._prep_regexes() + + ### this block is copied from MySQLTableDefinitionParser._prep_regexes + _final = self.preparer.final_quote + quotes = dict( + zip( + ("iq", "fq", "esc_fq"), + [ + re.escape(s) + for s in ( + self.preparer.initial_quote, + _final, + self.preparer._escape_identifier(_final), + ) + ], + ) + ) + ### end of block + + self._re_array_column = _re_compile( + r"\s*" + r"%(iq)s(?P(?:%(esc_fq)s|[^%(fq)s])+)%(fq)s\s+" + r"(?P(?i:(?(?:NOT\s+)?NULL))?" + r"(?:\s+DEFAULT\s+(?P(?:NULL|'(?:''|[^'])*'|\(.+?\)|[\-\w\.\(\)]+)))?" + r"(?:\s+COMMENT\s+'(?P(?:''|[^'])*)')?" + r"\s*,?\s*$" % quotes + ) + + self._re_key = _re_compile( + r" " + r"(?:(FULLTEXT|SPATIAL|VECTOR|SPARSEVECTOR|(?P\S+)) )?KEY" + # r"(?:(?P\S+) )?KEY" + r"(?: +{iq}(?P(?:{esc_fq}|[^{fq}])+){fq})?" + r"(?: +USING +(?P\S+))?" + r" +\((?P[^)]+)\)" + r"(?: +USING +(?P\S+))?" + r"(?: +WITH +\((?P[^)]+)\))?" + r"(?: +WITH PARSER +(?P\S+))?" + r"(?: +PARSER_PROPERTIES=\((?P[^)]+)\))?" + r"(?: +(KEY_)?BLOCK_SIZE *[ =]? *(?P\S+) *(LOCAL)?)?" + r"(?: +COMMENT +(?P(\x27\x27|\x27([^\x27])*?\x27)+))?" + r"(?: +/\*(?P.+)\*/ *)?" + r",?$".format(iq=quotes["iq"], esc_fq=quotes["esc_fq"], fq=quotes["fq"]) + ) + + kw = quotes.copy() + kw["on"] = "RESTRICT|CASCADE|SET NULL|NO ACTION" + self._re_fk_constraint = _re_compile( + r" " + r"CONSTRAINT +" + r"{iq}(?P(?:{esc_fq}|[^{fq}])+){fq} +" + r"FOREIGN KEY +" + r"\((?P[^\)]+?)\) REFERENCES +" + r"(?P{iq}[^{fq}]+{fq}" + r"(?:\.{iq}[^{fq}]+{fq})?) *" + r"\((?P(?:{iq}[^{fq}]+{fq}(?: *, *)?)+)\)" + r"(?: +(?PMATCH \w+))?" + r"(?: +ON UPDATE (?P{on}))?" + r"(?: +ON DELETE (?P{on}))?".format( + iq=quotes["iq"], esc_fq=quotes["esc_fq"], fq=quotes["fq"], on=kw["on"] + ) + ) + + def _parse_column(self, line, state): + m = self._re_array_column.match(line) + if m: + spec = m.groupdict() + name, coltype_with_args = spec["name"].strip(), spec["coltype_with_args"].strip() + + item_pattern = re.compile( + r"^(?:array\s*\()*([\w]+)(?:\(([\d,]+)\))?\)*$", + re.IGNORECASE + ) + item_m = item_pattern.match(coltype_with_args) + if not item_m: + raise ValueError(f"Failed to find inner type from array column definition: {line}") + + item_type = self.dialect.ischema_names[item_m.group(1).lower()] + item_type_arg = item_m.group(2) + if item_type_arg is None or item_type_arg == "": + item_type_args = [] + elif item_type_arg[0] == "'" and item_type_arg[-1] == "'": + item_type_args = self._re_csv_str.findall(item_type_arg) + else: + item_type_args = [int(v) for v in self._re_csv_int.findall(item_type_arg)] + + nested_level = coltype_with_args.lower().count('array') + type_instance = item_type(*item_type_args) + for _ in range(nested_level): + type_instance = ARRAY(type_instance) + + col_kw = {} + + # NOT NULL + col_kw["nullable"] = True + if spec.get("notnull", False) == "NOT NULL": + col_kw["nullable"] = False + + # DEFAULT + default = spec.get("default", None) + + if default == "NULL": + # eliminates the need to deal with this later. + default = None + + comment = spec.get("comment", None) + + if comment is not None: + comment = cleanup_text(comment) + + col_d = dict( + name=name, type=type_instance, default=default, comment=comment + ) + col_d.update(col_kw) + state.columns.append(col_d) + else: + super()._parse_column(line, state) + + def _parse_constraints(self, line): + """Parse a CONSTRAINT line.""" + ret = super()._parse_constraints(line) + if ret: + tp, spec = ret + + if tp is None or tp == "partition" or not isinstance(spec, dict): + return ret + + if tp == "fk_constraint": + table = spec.get("table", []) + if isinstance(table, list) and len(table) == 2 and table[0] == self.default_schema: + spec["table"] = table[1:] + + for action in ["onupdate", "ondelete"]: + if (spec.get(action) or "").lower() == "restrict": + spec[action] = None + return ret + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/replace_stmt.py b/pyobsql-oceanbase-plugin/pyobsql/schema/replace_stmt.py new file mode 100644 index 00000000..6f591a1a --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/replace_stmt.py @@ -0,0 +1,19 @@ +"""ReplaceStmt: replace into statement compilation.""" +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.expression import Insert + +class ReplaceStmt(Insert): + """Replace into statement.""" + inherit_cache = True + +@compiles(ReplaceStmt) +def compile_replace_stmt(insert, compiler, **kw): + """Compile replace into statement. + + Args: + insert: replace clause + compiler: SQL compiler + """ + stmt_str = compiler.visit_insert(insert, **kw) + stmt_str = stmt_str.replace("INSERT INTO", "REPLACE INTO") + return stmt_str diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/sparse_vector.py b/pyobsql-oceanbase-plugin/pyobsql/schema/sparse_vector.py new file mode 100644 index 00000000..23574973 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/sparse_vector.py @@ -0,0 +1,36 @@ +"""SPARSE_VECTOR: An extended data type for SQLAlchemy""" +from sqlalchemy.types import UserDefinedType, String +from ..util import SparseVector + +class SPARSE_VECTOR(UserDefinedType): + """SPARSE_VECTOR data type definition.""" + cache_ok = True + _string = String() + + def __init__(self): + super(UserDefinedType, self).__init__() + + def get_col_spec(self, **kw): # pylint: disable=unused-argument + """Parse to sparse vector data type definition in text SQL.""" + return "SPARSEVECTOR" + + def bind_processor(self, dialect): + def process(value): + return SparseVector._to_db(value) + + return process + + def literal_processor(self, dialect): + string_literal_processor = self._string._cached_literal_processor(dialect) + + def process(value): + return string_literal_processor(SparseVector._to_db(value)) + + return process + + def result_processor(self, dialect, coltype): + def process(value): + return SparseVector._from_db(value) + + return process + diff --git a/pyobsql-oceanbase-plugin/pyobsql/schema/vector.py b/pyobsql-oceanbase-plugin/pyobsql/schema/vector.py new file mode 100644 index 00000000..ad1db0fa --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/schema/vector.py @@ -0,0 +1,40 @@ +"""VECTOR: An extended data type for SQLAlchemy""" +from sqlalchemy.types import UserDefinedType, String +from ..util import Vector + + +class VECTOR(UserDefinedType): + """VECTOR data type definition.""" + cache_ok = True + _string = String() + + def __init__(self, dim=None): + super(UserDefinedType, self).__init__() + self.dim = dim + + def get_col_spec(self, **kw): # pylint: disable=unused-argument + """Parse to vector data type definition in text SQL.""" + if self.dim is None: + return "VECTOR" + return f"VECTOR({self.dim})" + + def bind_processor(self, dialect): + def process(value): + return Vector._to_db(value, self.dim) + + return process + + def literal_processor(self, dialect): + string_literal_processor = self._string._cached_literal_processor(dialect) + + def process(value): + return string_literal_processor(Vector._to_db(value, self.dim)) + + return process + + def result_processor(self, dialect, coltype): + def process(value): + return Vector._from_db(value) + + return process + diff --git a/pyobsql-oceanbase-plugin/pyobsql/util/__init__.py b/pyobsql-oceanbase-plugin/pyobsql/util/__init__.py new file mode 100644 index 00000000..fd653075 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/util/__init__.py @@ -0,0 +1,15 @@ +"""A utility module for pyobsql. + +* ObVersion OceanBase cluster version class +* Vector Vector utility class for VECTOR data type +* SparseVector SparseVector utility class for SPARSE_VECTOR data type +""" +from .ob_version import ObVersion +from .vector import Vector +from .sparse_vector import SparseVector + +__all__ = ["ObVersion", "Vector", "SparseVector"] + + + + diff --git a/pyobsql-oceanbase-plugin/pyobsql/util/ob_version.py b/pyobsql-oceanbase-plugin/pyobsql/util/ob_version.py new file mode 100644 index 00000000..d0b2d51c --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/util/ob_version.py @@ -0,0 +1,48 @@ +"""OceanBase cluster version module.""" +import copy +from typing import List + + +class ObVersion: + """The class to describe OceanBase cluster version. + + Attributes: + version_nums (List[int]): version number of OceanBase cluster. For example, '4.3.3.0' + """ + def __init__(self, version_nums: List[int]): + self.version_nums = copy.deepcopy(version_nums) + + @classmethod + def from_db_version_string(cls, version: str): + """Construct ObVersion with a version string. + + Args: + version: a string of 4 numbers separated by '.' + """ + return cls([int(version_num) for version_num in version.split(".")]) + + @classmethod + def from_db_version_nums( + cls, main_ver, sub_ver1: int, sub_ver2: int, sub_ver3: int + ): + """Construct ObVersion with 4 version numbers. + + Args: + main_ver: main version + sub_ver1: first subversion + sub_ver2: second subversion + sub_ver3: third subversion + """ + return cls([main_ver, sub_ver1, sub_ver2, sub_ver3]) + + def __lt__(self, other): + if len(self.version_nums) != len(other.version_nums): + raise ValueError("version num list length is not equal") + idx, ilen = 0, len(self.version_nums) + while idx < ilen: + if self.version_nums[idx] < other.version_nums[idx]: + return True + if self.version_nums[idx] > other.version_nums[idx]: + return False + idx += 1 + return False diff --git a/pyobsql-oceanbase-plugin/pyobsql/util/sparse_vector.py b/pyobsql-oceanbase-plugin/pyobsql/util/sparse_vector.py new file mode 100644 index 00000000..66aaf4d5 --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/util/sparse_vector.py @@ -0,0 +1,49 @@ +"""A utility module for the extended data type class 'SPARSE_VECTOR'.""" +import ast + +class SparseVector: + """A transformer class between python dict and OceanBase SPARSE_VECTOR. + + Attributes: + _value (Dict) : a python dict + """ + def __init__(self, value): + if not isinstance(value, dict): + raise ValueError("Sparse Vector should be a dict in python") + + self._value = value + + def __repr__(self): + return f"{self._value}" + + def to_text(self): + return f"{self._value}" + + @classmethod + def from_text(cls, value: str): + """Construct Sparse Vector class with dict in string format. + + Args: + value: For example, '{1:1.1, 2:2.2}' + """ + return cls(ast.literal_eval(value)) + + @classmethod + def _to_db(cls, value): + if value is None: + return value + + if not isinstance(value, cls): + value = cls(value) + + return value.to_text() + + @classmethod + def _from_db(cls, value): + if value is None or isinstance(value, dict): + return value + + if isinstance(value, str): + return cls.from_text(value)._value + raise ValueError(f"unexpected sparse vector type: {type(value)}") + diff --git a/pyobsql-oceanbase-plugin/pyobsql/util/vector.py b/pyobsql-oceanbase-plugin/pyobsql/util/vector.py new file mode 100644 index 00000000..f753cf8a --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyobsql/util/vector.py @@ -0,0 +1,88 @@ +"""A utility module for the extended data type class 'VECTOR'.""" +import json +import numpy as np + + +class Vector: + """A transformer class between python numpy array and OceanBase VECTOR. + + Attributes: + _value (numpy.array): a numpy array + """ + def __init__(self, value): + # big-endian float32 + if not isinstance(value, np.ndarray) or value.dtype != ">f4": + value = np.asarray(value, dtype=">f4") + + if value.ndim != 1: + raise ValueError(f"expected ndim to be 1: {value} {value.ndim}") + + self._value = value + + def __repr__(self): + return f"{self._value.tolist()}" + + def dim(self): + """Get vector dimension.""" + return len(self._value) + + def to_list(self): + """Parse numpy array to python list.""" + return self._value.tolist() + + def to_numpy(self): + """Get numpy array.""" + return self._value + + def to_text(self): + """Parse numpy array to list string.""" + return "[" + ",".join([str(np.float32(v)) for v in self._value]) + "]" + + @classmethod + def from_text(cls, value: str): + """Construct Vector class with list string. + + Args: + value (str): For example, '[1,2,3]' + + Returns: + Vector: Vector instance + """ + return cls([float(v) for v in value[1:-1].split(",")]) + + @classmethod + def from_bytes(cls, value: bytes): + """Construct Vector class with raw bytes. + + Args: + value (bytes): the bytes of python list + + Returns: + Vector: Vector instance + """ + return cls(json.loads(value.decode())) + + @classmethod + def _to_db(cls, value, dim=None): + if value is None: + return value + + if not isinstance(value, cls): + value = cls(value) + + if dim is not None and value.dim() != dim: + raise ValueError(f"expected {dim} dimensions, not {value.dim()}") + + return value.to_text() + + @classmethod + def _from_db(cls, value): + if value is None or isinstance(value, np.ndarray): + return value + + if isinstance(value, str): + return cls.from_text(value).to_numpy().astype(np.float32) + if isinstance(value, bytes): + return cls.from_bytes(value).to_numpy().astype(np.float32) + raise ValueError("unexpected vector type") + diff --git a/pyobsql-oceanbase-plugin/pyproject.toml b/pyobsql-oceanbase-plugin/pyproject.toml new file mode 100644 index 00000000..e2f70c9b --- /dev/null +++ b/pyobsql-oceanbase-plugin/pyproject.toml @@ -0,0 +1,54 @@ +[project] +name = "pyobsql" +version = "0.1.0" +description = "A python SDK for OceanBase SQL, including JSON Table support and SQLAlchemy dialect extensions." +authors = [{name="shanhaikang.shk",email="shanhaikang.shk@oceanbase.com"}] +readme = "README.md" +license = {text = "Apache-2.0"} +keywords = ["oceanbase", "sql", "json-table", "sqlalchemy"] +classifiers = [ + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] +requires-python = ">=3.9" + +dependencies = [ + "numpy>=1.17.0,<2.0.0", + "sqlalchemy>=1.4,<=3", + "sqlglot>=26.0.1,<28.0.0", + "pydantic>=2.7.0,<3" +] + +[project.urls] +Homepage = "https://github.com/oceanbase/ecology-plugins/tree/main/pyobsql-oceanbase-plugin" +Repository = "https://github.com/oceanbase/ecology-plugins.git" + +[dependency-groups] +dev = [ + "pytest>=8.2.2", + "pylint>=3.2.7" +] + +[tool.pytest.ini_options] +log_level = "INFO" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatchling.metadata] +license = {text = "Apache-2.0"} + +[tool.hatchling.build] +packages = ["pyobsql"] + +[tool.hatch.build.targets.wheel] +metadata-version = "2.1" + + + + diff --git a/pyobsql-oceanbase-plugin/tests/__init__.py b/pyobsql-oceanbase-plugin/tests/__init__.py new file mode 100644 index 00000000..8d7521be --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/__init__.py @@ -0,0 +1,16 @@ +import unittest +from pyobsql.client import ObClient +from pyobsql import OceanBase +from sqlalchemy import Column, Integer, Table +from sqlalchemy.sql import func +from sqlalchemy.exc import NoSuchTableError + + +class ObClientTest(unittest.TestCase): + def setUp(self) -> None: + self.client = ObClient(echo=True) + + +if __name__ == "__main__": + unittest.main() + diff --git a/pyobsql-oceanbase-plugin/tests/test_comprehensive_oceanbase.py b/pyobsql-oceanbase-plugin/tests/test_comprehensive_oceanbase.py new file mode 100644 index 00000000..c04162cb --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/test_comprehensive_oceanbase.py @@ -0,0 +1,568 @@ +""" +Comprehensive test suite for OceanBase 4.4.1 using pyobsql whl package. +This test covers all major features of pyobsql. +""" +import unittest +import logging +from datetime import datetime + +from pyobsql.client import ObClient, ObRangePartition, RangeListPartInfo +from pyobsql.schema import ( + VECTOR, + SPARSE_VECTOR, + ARRAY, + POINT, + ST_GeomFromText, + st_distance +) +from pyobsql import ( + JsonTableBool, + JsonTableInt, + JsonTableTimestamp, + JsonTableVarcharFactory, + JsonTableDecimalFactory, + val2json, + json_value, +) +from sqlalchemy import Column, Integer, String, JSON, Table, Index, select, func, update + +logger = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) + + +class ComprehensiveOceanBaseTest(unittest.TestCase): + """Comprehensive test suite for OceanBase 4.4.1""" + + @classmethod + def setUpClass(cls): + """Set up test class with database connection""" + connection_args = { + "host": "11.124.9.21", + "port": "3881", + "user": "root@sun", + "password": "ShengTai@2024yyds", + "db_name": "langchain", + } + + # Construct URI from host and port + uri = f"{connection_args['host']}:{connection_args['port']}" + + logger.info(f"Connecting to OceanBase at {uri}") + cls.client = ObClient( + uri=uri, + user=connection_args['user'], + password=connection_args['password'], + db_name=connection_args['db_name'], + echo=False + ) + + logger.info(f"Connected successfully! OceanBase version: {cls.client.ob_version}") + + # Clean up any existing test tables + test_tables = [ + 'test_basic_table', + 'test_vector_table', + 'test_array_table', + 'test_partitioned_table', + 'test_json_table', + 'test_point_table', + 'test_sparse_vector_table', + ] + for table_name in test_tables: + try: + cls.client.drop_table_if_exist(table_name) + logger.info(f"Dropped existing table: {table_name}") + except Exception as e: + logger.warning(f"Error dropping table {table_name}: {e}") + + def test_01_connection(self): + """Test database connection""" + logger.info("=" * 60) + logger.info("Test 1: Database Connection") + logger.info("=" * 60) + + self.assertIsNotNone(self.client) + self.assertIsNotNone(self.client.engine) + logger.info(f"✓ Connection successful. OceanBase version: {self.client.ob_version}") + + def test_02_create_basic_table(self): + """Test creating a basic table""" + logger.info("=" * 60) + logger.info("Test 2: Create Basic Table") + logger.info("=" * 60) + + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('age', Integer), + Column('email', String(255)), + Column('metadata', JSON), + ] + + self.client.create_table('test_basic_table', columns=columns) + logger.info("✓ Basic table created successfully") + + # Verify table exists + self.assertTrue(self.client.check_table_exists('test_basic_table')) + logger.info("✓ Table existence verified") + + def test_03_insert_basic_data(self): + """Test inserting basic data""" + logger.info("=" * 60) + logger.info("Test 3: Insert Basic Data") + logger.info("=" * 60) + + # Insert single record + self.client.insert('test_basic_table', { + 'id': 1, + 'name': 'Alice', + 'age': 30, + 'email': 'alice@example.com', + 'metadata': {'department': 'Engineering', 'role': 'Developer'} + }) + logger.info("✓ Single record inserted") + + # Batch insert + data_list = [ + { + 'id': i, + 'name': f'User_{i}', + 'age': 20 + i, + 'email': f'user_{i}@example.com', + 'metadata': {'index': i, 'status': 'active'} + } + for i in range(2, 11) + ] + self.client.insert('test_basic_table', data_list) + logger.info(f"✓ Batch inserted {len(data_list)} records") + + def test_04_query_basic_data(self): + """Test querying basic data""" + logger.info("=" * 60) + logger.info("Test 4: Query Basic Data") + logger.info("=" * 60) + + # Query all records + result = self.client.get('test_basic_table') + rows = list(result) + logger.info(f"✓ Queried all records: {len(rows)} rows") + + # Query by primary key + result = self.client.get('test_basic_table', ids=1) + rows = list(result) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0].name, 'Alice') + logger.info("✓ Query by primary key successful") + + # Query with conditions + table = Table('test_basic_table', self.client.metadata_obj, autoload_with=self.client.engine) + result = self.client.get( + 'test_basic_table', + where_clause=[table.c.age > 25], + n_limits=5 + ) + rows = list(result) + logger.info(f"✓ Conditional query returned {len(rows)} rows") + + # Query with output columns + result = self.client.get( + 'test_basic_table', + output_column_name=['id', 'name', 'email'], + n_limits=3 + ) + rows = list(result) + logger.info(f"✓ Query with specific columns returned {len(rows)} rows") + + def test_05_update_data(self): + """Test updating data""" + logger.info("=" * 60) + logger.info("Test 5: Update Data") + logger.info("=" * 60) + + table = Table('test_basic_table', self.client.metadata_obj, autoload_with=self.client.engine) + + # Update single record - use dictionary for values_clause + # Note: SQLAlchemy update().values() expects dict or keyword args, not Column == value expressions + # We'll use a workaround by building the update statement manually + with self.client.engine.connect() as conn: + with conn.begin(): + update_stmt = update(table).where(table.c.id == 1).values( + age=31, + metadata={'department': 'Engineering', 'role': 'Senior Developer'} + ) + conn.execute(update_stmt) + logger.info("✓ Updated single record using manual update statement") + logger.info("✓ Updated single record") + + # Verify update + result = self.client.get('test_basic_table', ids=1) + rows = list(result) + self.assertEqual(rows[0].age, 31) + logger.info("✓ Update verified") + + def test_06_delete_data(self): + """Test deleting data""" + logger.info("=" * 60) + logger.info("Test 6: Delete Data") + logger.info("=" * 60) + + # Delete by primary key + self.client.delete('test_basic_table', ids=10) + logger.info("✓ Deleted record by primary key") + + # Verify deletion + result = self.client.get('test_basic_table', ids=10) + rows = list(result) + self.assertEqual(len(rows), 0) + logger.info("✓ Deletion verified") + + # Delete by condition + table = Table('test_basic_table', self.client.metadata_obj, autoload_with=self.client.engine) + self.client.delete( + 'test_basic_table', + where_clause=[table.c.age < 22] + ) + logger.info("✓ Deleted records by condition") + + def test_07_create_vector_table(self): + """Test creating table with VECTOR type""" + logger.info("=" * 60) + logger.info("Test 7: Create Vector Table") + logger.info("=" * 60) + + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('embedding', VECTOR(128)), + ] + + self.client.create_table('test_vector_table', columns=columns) + logger.info("✓ Vector table created") + + # Insert vector data + vector_data = [0.1 * i for i in range(128)] + self.client.insert('test_vector_table', { + 'id': 1, + 'name': 'vector_1', + 'embedding': vector_data + }) + logger.info("✓ Vector data inserted") + + def test_08_create_sparse_vector_table(self): + """Test creating table with SPARSE_VECTOR type""" + logger.info("=" * 60) + logger.info("Test 8: Create Sparse Vector Table") + logger.info("=" * 60) + + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('sparse_vec', SPARSE_VECTOR), + ] + + self.client.create_table('test_sparse_vector_table', columns=columns) + logger.info("✓ Sparse vector table created") + + # Insert sparse vector data + sparse_data = {1: 0.5, 5: 0.8, 10: 0.3, 20: 0.9} + self.client.insert('test_sparse_vector_table', { + 'id': 1, + 'name': 'sparse_1', + 'sparse_vec': sparse_data + }) + logger.info("✓ Sparse vector data inserted") + + def test_09_create_array_table(self): + """Test creating table with ARRAY type""" + logger.info("=" * 60) + logger.info("Test 9: Create Array Table") + logger.info("=" * 60) + + columns = [ + Column('id', Integer, primary_key=True), + Column('tags', ARRAY(String(50))), + Column('scores', ARRAY(Integer)), + ] + + self.client.create_table('test_array_table', columns=columns) + logger.info("✓ Array table created") + + # Insert array data + self.client.insert('test_array_table', { + 'id': 1, + 'tags': ['tag1', 'tag2', 'tag3'], + 'scores': [100, 200, 300] + }) + logger.info("✓ Array data inserted") + + def test_10_create_point_table(self): + """Test creating table with POINT type""" + logger.info("=" * 60) + logger.info("Test 10: Create Point Table") + logger.info("=" * 60) + + # Skip POINT test if not supported in this OceanBase version + try: + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('location', POINT(srid=4326)), + ] + + self.client.create_table('test_point_table', columns=columns) + logger.info("✓ Point table created") + + # Insert point data + self.client.insert('test_point_table', { + 'id': 1, + 'name': 'Beijing', + 'location': (116.3974, 39.9093) # (longitude, latitude) + }) + logger.info("✓ Point data inserted") + + # Query using GIS functions + table = Table('test_point_table', self.client.metadata_obj, autoload_with=self.client.engine) + stmt = select( + table.c.id, + table.c.name, + st_distance( + table.c.location, + ST_GeomFromText('POINT(116.3974 39.9093)', 4326) + ).label('distance') + ) + with self.client.engine.connect() as conn: + result = conn.execute(stmt) + rows = list(result) + logger.info(f"✓ GIS query returned {len(rows)} rows") + except Exception as e: + logger.warning(f"POINT type may not be supported in OceanBase 4.4.1: {e}") + logger.info("⚠ Skipping POINT table test") + self.skipTest(f"POINT type not supported: {e}") + + def test_11_create_partitioned_table(self): + """Test creating partitioned table""" + logger.info("=" * 60) + logger.info("Test 11: Create Partitioned Table") + logger.info("=" * 60) + + # Define Range partition strategy + range_partition = ObRangePartition( + is_range_columns=False, + range_part_infos=[ + RangeListPartInfo('p0', 100), + RangeListPartInfo('p1', 200), + RangeListPartInfo('p2', 'MAXVALUE') + ], + range_expr='id' + ) + + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('embedding', VECTOR(128)), + ] + + self.client.create_table('test_partitioned_table', columns=columns, partitions=range_partition) + logger.info("✓ Partitioned table created") + + # Insert data to specific partition + vector_data = [0.1] * 128 + self.client.insert( + 'test_partitioned_table', + {'id': 50, 'name': 'partitioned_1', 'embedding': vector_data}, + partition_name='p0' + ) + logger.info("✓ Data inserted to partition p0") + + def test_12_json_table_types(self): + """Test JSON Table virtual data types""" + logger.info("=" * 60) + logger.info("Test 12: JSON Table Types") + logger.info("=" * 60) + + # Test JsonTableBool - val2json expects the value, not the object + bool_type = JsonTableBool(val=True) + json_value = val2json(bool_type.val) + self.assertEqual(json_value, True) + logger.info("✓ JsonTableBool works") + + # Test JsonTableInt + int_type = JsonTableInt(val=42) + json_value = val2json(int_type.val) + self.assertEqual(json_value, 42) + logger.info("✓ JsonTableInt works") + + # Test JsonTableTimestamp + timestamp = datetime(2024, 12, 30, 3, 35, 30) + timestamp_type = JsonTableTimestamp(val=timestamp) + json_value = val2json(timestamp_type.val) + self.assertEqual(json_value, timestamp.isoformat()) + logger.info("✓ JsonTableTimestamp works") + + # Test JsonTableVarchar + varchar_factory = JsonTableVarcharFactory(length=255) + varchar_type = varchar_factory.get_json_table_varchar_type()(val="test") + json_value = val2json(varchar_type.val) + self.assertEqual(json_value, "test") + logger.info("✓ JsonTableVarchar works") + + # Test JsonTableDecimal - use ndigits and decimal_p instead of precision and scale + decimal_factory = JsonTableDecimalFactory(ndigits=10, decimal_p=2) + decimal_type = decimal_factory.get_json_table_decimal_type()(val=123.45) + json_value = val2json(decimal_type.val) + # val2json returns float for Decimal + self.assertAlmostEqual(json_value, 123.45, places=2) + logger.info(f"✓ JsonTableDecimal works (value: {json_value})") + + def test_13_json_value_function(self): + """Test json_value function""" + logger.info("=" * 60) + logger.info("Test 13: json_value Function") + logger.info("=" * 60) + + # Create a table with JSON column + columns = [ + Column('id', Integer, primary_key=True), + Column('metadata', JSON), + ] + self.client.create_table('test_json_table', columns=columns) + + # Insert JSON data + self.client.insert('test_json_table', { + 'id': 1, + 'metadata': {'key': 'value', 'number': 42, 'nested': {'deep': 'data'}} + }) + + # Query using json_value + table = Table('test_json_table', self.client.metadata_obj, autoload_with=self.client.engine) + stmt = select( + table.c.id, + json_value(table.c.metadata, '$.key', 'VARCHAR(100)').label('extracted_key') + ).where(table.c.id == 1) + + with self.client.engine.connect() as conn: + result = conn.execute(stmt) + rows = list(result) + self.assertEqual(len(rows), 1) + logger.info("✓ json_value function works") + + def test_14_upsert_operation(self): + """Test upsert (REPLACE INTO) operation""" + logger.info("=" * 60) + logger.info("Test 14: Upsert Operation") + logger.info("=" * 60) + + # Upsert with existing primary key (should replace) + self.client.upsert('test_basic_table', { + 'id': 1, + 'name': 'Alice_Updated', + 'age': 32, + 'email': 'alice_updated@example.com', + 'metadata': {'status': 'updated'} + }) + logger.info("✓ Upsert operation completed") + + # Verify upsert + result = self.client.get('test_basic_table', ids=1) + rows = list(result) + self.assertEqual(rows[0].name, 'Alice_Updated') + logger.info("✓ Upsert verified") + + def test_15_table_with_indexes(self): + """Test creating table with indexes""" + logger.info("=" * 60) + logger.info("Test 15: Table with Indexes") + logger.info("=" * 60) + + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('email', String(255)), + ] + + indexes = [ + Index('idx_name', 'name'), + Index('idx_email', 'email'), + ] + + self.client.create_table('test_indexed_table', columns=columns, indexes=indexes) + logger.info("✓ Table with indexes created") + + # Insert data + self.client.insert('test_indexed_table', { + 'id': 1, + 'name': 'Test User', + 'email': 'test@example.com' + }) + logger.info("✓ Data inserted into indexed table") + + def test_16_refresh_metadata(self): + """Test refreshing metadata""" + logger.info("=" * 60) + logger.info("Test 16: Refresh Metadata") + logger.info("=" * 60) + + # Refresh all metadata + self.client.refresh_metadata() + logger.info("✓ Refreshed all metadata") + + # Refresh specific table metadata + self.client.refresh_metadata(tables=['test_basic_table']) + logger.info("✓ Refreshed specific table metadata") + + def test_17_complex_query(self): + """Test complex queries using SQLAlchemy""" + logger.info("=" * 60) + logger.info("Test 17: Complex Queries") + logger.info("=" * 60) + + table = Table('test_basic_table', self.client.metadata_obj, autoload_with=self.client.engine) + + # Complex query with joins, aggregations, etc. + stmt = select( + table.c.id, + table.c.name, + func.json_extract(table.c.metadata, '$.department').label('department') + ).where( + table.c.age > 25 + ).order_by( + table.c.id.desc() + ).limit(5) + + with self.client.engine.connect() as conn: + result = conn.execute(stmt) + rows = list(result) + logger.info(f"✓ Complex query returned {len(rows)} rows") + + @classmethod + def tearDownClass(cls): + """Clean up test tables""" + logger.info("=" * 60) + logger.info("Cleaning up test tables") + logger.info("=" * 60) + + test_tables = [ + 'test_basic_table', + 'test_vector_table', + 'test_array_table', + 'test_partitioned_table', + 'test_json_table', + 'test_point_table', + 'test_sparse_vector_table', + 'test_indexed_table', + ] + + for table_name in test_tables: + try: + cls.client.drop_table_if_exist(table_name) + logger.info(f"✓ Dropped table: {table_name}") + except Exception as e: + logger.warning(f"Error dropping table {table_name}: {e}") + + +if __name__ == "__main__": + unittest.main(verbosity=2) + diff --git a/pyobsql-oceanbase-plugin/tests/test_json_table.py b/pyobsql-oceanbase-plugin/tests/test_json_table.py new file mode 100644 index 00000000..26ddc372 --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/test_json_table.py @@ -0,0 +1,56 @@ +import unittest +from pyobsql import ( + JsonTableBool, + JsonTableInt, + JsonTableTimestamp, + JsonTableVarcharFactory, + JsonTableDecimalFactory, + val2json) +import logging +from datetime import datetime +from decimal import Decimal + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class JsonTableTest(unittest.TestCase): + def setUp(self) -> None: + return super().setUp() + + def test_json_table_bool(self): + bool_type = JsonTableBool(val=True) + json_value = val2json(bool_type) + self.assertEqual(json_value, True) + + bool_type = JsonTableBool(val=False) + json_value = val2json(bool_type) + self.assertEqual(json_value, False) + + def test_json_table_int(self): + int_type = JsonTableInt(val=42) + json_value = val2json(int_type) + self.assertEqual(json_value, 42) + + def test_json_table_timestamp(self): + timestamp = datetime(2024, 12, 30, 3, 35, 30) + timestamp_type = JsonTableTimestamp(val=timestamp) + json_value = val2json(timestamp_type) + self.assertEqual(json_value, timestamp.isoformat()) + + def test_json_table_varchar(self): + varchar_factory = JsonTableVarcharFactory(length=255) + varchar_type = varchar_factory.get_json_table_varchar_type()(val="test") + json_value = val2json(varchar_type) + self.assertEqual(json_value, "test") + + def test_json_table_decimal(self): + decimal_factory = JsonTableDecimalFactory(precision=10, scale=2) + decimal_type = decimal_factory.get_json_table_decimal_type()(val=123.45) + json_value = val2json(decimal_type) + self.assertEqual(json_value, "123.45") + + +if __name__ == "__main__": + unittest.main() + diff --git a/pyobsql-oceanbase-plugin/tests/test_ob_client.py b/pyobsql-oceanbase-plugin/tests/test_ob_client.py new file mode 100644 index 00000000..53389929 --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/test_ob_client.py @@ -0,0 +1,37 @@ +import unittest +from pyobsql.client import ObClient +from sqlalchemy import Column, Integer, String, JSON +from pyobsql.schema import VECTOR, ARRAY +import logging + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class ObClientTest(unittest.TestCase): + def setUp(self) -> None: + # Note: These tests require actual database connections, may need to skip or use mock in CI + self.client = None # ObClient(echo=True) + + def test_client_initialization(self): + # Test client initialization (without actually connecting to database) + # If actual testing is needed, database connection information needs to be configured + pass + + def test_table_creation_structure(self): + # Test table structure definition + columns = [ + Column('id', Integer, primary_key=True), + Column('name', String(100)), + Column('embedding', VECTOR(128)), + Column('tags', ARRAY(String(50))), + Column('metadata', JSON) + ] + self.assertEqual(len(columns), 5) + self.assertEqual(columns[0].name, 'id') + self.assertEqual(columns[1].name, 'name') + + +if __name__ == "__main__": + unittest.main() + diff --git a/pyobsql-oceanbase-plugin/tests/test_oceanbase_dialect.py b/pyobsql-oceanbase-plugin/tests/test_oceanbase_dialect.py new file mode 100644 index 00000000..54366e39 --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/test_oceanbase_dialect.py @@ -0,0 +1,49 @@ +import unittest +from pyobsql import OceanBase +import logging + +from sqlglot import parse_one + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class OceanBaseDialectTest(unittest.TestCase): + def setUp(self) -> None: + return super().setUp() + + def test_drop_column(self): + sql = "ALTER TABLE users DROP COLUMN age" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + sql = "ALTER TABLE users DROP age" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + def test_modify_column(self): + sql = "ALTER TABLE users MODIFY COLUMN email VARCHAR(100) NOT NULL" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + sql = "ALTER TABLE users MODIFY email VARCHAR(100) NOT NULL" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + sql = "ALTER TABLE users MODIFY COLUMN email VARCHAR(100) NOT NULL DEFAULT 'ca'" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + def test_change_column(self): + sql = "ALTER TABLE users CHANGE COLUMN username user_name VARCHAR(50)" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + sql = "ALTER TABLE users CHANGE username user_name VARCHAR(50)" + ob_ast = parse_one(sql, dialect="oceanbase") + logger.info(f"\n{repr(ob_ast)}") + + +if __name__ == "__main__": + unittest.main() + diff --git a/pyobsql-oceanbase-plugin/tests/test_partition_compile.py b/pyobsql-oceanbase-plugin/tests/test_partition_compile.py new file mode 100644 index 00000000..4b3c17ae --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/test_partition_compile.py @@ -0,0 +1,120 @@ +import unittest +from pyobsql.client import ( + ObRangePartition, + ObSubRangePartition, + ObListPartition, + ObSubListPartition, + ObHashPartition, + ObSubHashPartition, + ObKeyPartition, + ObSubKeyPartition, + RangeListPartInfo, +) +import logging + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class ObPartitionTest(unittest.TestCase): + def setUp(self) -> None: + self.maxDiff = None + + def test_range_and_range_columns_partition(self): + range_part = ObRangePartition( + False, + range_part_infos=[ + RangeListPartInfo("p0", 100), + RangeListPartInfo("p1", "maxvalue"), + ], + range_expr="id", + ) + self.assertEqual( + range_part.do_compile(), + "PARTITION BY RANGE (id) (PARTITION p0 VALUES LESS THAN (100),PARTITION p1 VALUES LESS THAN (maxvalue))", + ) + + range_columns_part = ObRangePartition( + True, + range_part_infos=[ + RangeListPartInfo("M202001", "'2020/02/01'"), + RangeListPartInfo("M202002", "'2020/03/01'"), + RangeListPartInfo("M202003", "'2020/04/01'"), + RangeListPartInfo("MMAX", "MAXVALUE"), + ], + col_name_list=["log_date"], + ) + self.assertEqual( + range_columns_part.do_compile(), + "PARTITION BY RANGE COLUMNS (log_date) (PARTITION M202001 VALUES LESS THAN ('2020/02/01'),PARTITION M202002 VALUES LESS THAN ('2020/03/01'),PARTITION M202003 VALUES LESS THAN ('2020/04/01'),PARTITION MMAX VALUES LESS THAN (MAXVALUE))", + ) + + def test_list_and_list_columns_partition(self): + list_part = ObListPartition( + False, + list_part_infos=[ + RangeListPartInfo("p0", [1, 2, 3]), + RangeListPartInfo("p1", [5, 6]), + RangeListPartInfo("p2", "DEFAULT"), + ], + list_expr="col1", + ) + self.assertEqual( + list_part.do_compile(), + "PARTITION BY LIST (col1) (PARTITION p0 VALUES IN (1,2,3),PARTITION p1 VALUES IN (5,6),PARTITION p2 VALUES IN (DEFAULT))", + ) + + list_columns_part = ObListPartition( + True, + list_part_infos=[ + RangeListPartInfo("p0", ["'00'", "'01'"]), + RangeListPartInfo("p1", ["'02'", "'03'"]), + RangeListPartInfo("p2", "DEFAULT"), + ], + col_name_list=["partition_id"], + ) + self.assertEqual( + list_columns_part.do_compile(), + "PARTITION BY LIST COLUMNS (partition_id) (PARTITION p0 VALUES IN ('00','01'),PARTITION p1 VALUES IN ('02','03'),PARTITION p2 VALUES IN (DEFAULT))", + ) + + def test_hash_and_key_partition(self): + hash_part = ObHashPartition("col1", part_count=60) + self.assertEqual( + hash_part.do_compile(), "PARTITION BY HASH (col1) PARTITIONS 60" + ) + + key_part = ObKeyPartition(col_name_list=["id", "gmt_create"], part_count=10) + self.assertEqual( + key_part.do_compile(), "PARTITION BY KEY (id,gmt_create) PARTITIONS 10" + ) + + def test_range_columns_with_sub_partition(self): + range_columns_part = ObRangePartition( + True, + range_part_infos=[ + RangeListPartInfo("p0", 100), + RangeListPartInfo("p1", 200), + RangeListPartInfo("p2", 300), + ], + col_name_list=["col1"], + ) + range_sub_part = ObSubRangePartition( + False, + range_part_infos=[ + RangeListPartInfo("mp0", 1000), + RangeListPartInfo("mp1", 2000), + RangeListPartInfo("mp2", 3000), + ], + range_expr="col3", + ) + range_columns_part.add_subpartition(range_sub_part) + self.assertEqual( + range_columns_part.do_compile(), + "PARTITION BY RANGE COLUMNS (col1) SUBPARTITION BY RANGE (col3) SUBPARTITION TEMPLATE (SUBPARTITION mp0 VALUES LESS THAN (1000),SUBPARTITION mp1 VALUES LESS THAN (2000),SUBPARTITION mp2 VALUES LESS THAN (3000)) (PARTITION p0 VALUES LESS THAN (100),PARTITION p1 VALUES LESS THAN (200),PARTITION p2 VALUES LESS THAN (300))", + ) + + +if __name__ == "__main__": + unittest.main() + diff --git a/pyobsql-oceanbase-plugin/tests/test_reflection.py b/pyobsql-oceanbase-plugin/tests/test_reflection.py new file mode 100644 index 00000000..164d6a44 --- /dev/null +++ b/pyobsql-oceanbase-plugin/tests/test_reflection.py @@ -0,0 +1,40 @@ +import unittest +from pyobsql.schema import OceanBaseDialect +from sqlalchemy.dialects import registry +import logging + +logger = logging.getLogger(__name__) + + +class ObReflectionTest(unittest.TestCase): + def test_reflection(self): + dialect = OceanBaseDialect() + ddl = """CREATE TABLE `test_table` ( + `id` varchar(4096) NOT NULL, + `text` longtext DEFAULT NULL, + `embeddings` VECTOR(1024) DEFAULT NULL, + `metadata` json DEFAULT NULL, + PRIMARY KEY (`id`) +) DEFAULT CHARSET = utf8mb4 ROW_FORMAT = DYNAMIC COMPRESSION = 'zstd_1.3.8' REPLICA_NUM = 1 BLOCK_SIZE = 16384 USE_BLOOM_FILTER = FALSE TABLET_SIZE = 134217728 PCTFREE = 0 +""" + state = dialect._tabledef_parser.parse(ddl, "utf8") + assert len(state.columns) == 4 + assert len(state.keys) == 1 + + def test_dialect(self): + uri: str = "127.0.0.1:2881" + user: str = "root@test" + password: str = "" + db_name: str = "test" + registry.register("mysql.aoceanbase", "pyobsql.schema.dialect", "AsyncOceanBaseDialect") + connection_str = ( + f"mysql+aoceanbase://{user}:{password}@{uri}/{db_name}?charset=utf8mb4" + ) + # Note: This does not actually create an engine, just tests registration + # from sqlalchemy.ext.asyncio import create_async_engine + # self.engine = create_async_engine(connection_str) + + +if __name__ == "__main__": + unittest.main() +