datus-starrocks 0.1.0__tar.gz

datus_starrocks-0.1.0/.gitignore

@@ -0,0 +1,137 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+Pipfile.lock
+
+# uv
+uv.lock
+
+# PEP 582
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# IDEs
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db

datus_starrocks-0.1.0/PKG-INFO

@@ -0,0 +1,200 @@
+Metadata-Version: 2.4
+Name: datus-starrocks
+Version: 0.1.0
+Summary: StarRocks database adapter for Datus
+Project-URL: Homepage, https://github.com/Datus-ai/datus-db-adapters
+Project-URL: Repository, https://github.com/Datus-ai/datus-db-adapters
+Project-URL: Issues, https://github.com/Datus-ai/datus-db-adapters/issues
+Author-email: DatusAI <support@datus.ai>
+License: Apache-2.0
+Keywords: adapter,database,datus,starrocks
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.12
+Requires-Dist: datus-agent>0.2.1
+Requires-Dist: datus-mysql>=0.1.0
+Provides-Extra: test
+Requires-Dist: pytest-cov>=4.0.0; extra == 'test'
+Requires-Dist: pytest>=7.0.0; extra == 'test'
+Description-Content-Type: text/markdown
+
+# datus-starrocks
+
+StarRocks database adapter for Datus.
+
+## Overview
+
+StarRocks is a high-performance analytical database that uses the MySQL protocol. This adapter extends the MySQL connector with StarRocks-specific features:
+
+- Multi-catalog support
+- Materialized views
+- StarRocks-specific metadata queries
+
+## Installation
+
+```bash
+pip install datus-starrocks
+```
+
+This will automatically install the required dependencies:
+- `datus-agent`
+- `datus-mysql` (which includes `datus-sqlalchemy`)
+
+## Usage
+
+The adapter is automatically registered with Datus when installed. Configure your database connection:
+
+```yaml
+database:
+  type: starrocks
+  host: localhost
+  port: 9030
+  username: root
+  password: your_password
+  catalog: default_catalog
+  database: your_database
+```
+
+Or use programmatically:
+
+```python
+from datus_starrocks import StarRocksConnector
+
+# Create connector
+connector = StarRocksConnector(
+    host="localhost",
+    port=9030,
+    user="root",
+    password="your_password",
+    catalog="default_catalog",
+    database="mydb"
+)
+
+# Use context manager for automatic cleanup
+with connector:
+    # Test connection
+    connector.test_connection()
+
+    # Get catalogs
+    catalogs = connector.get_catalogs()
+    print(f"Catalogs: {catalogs}")
+
+    # Get databases in catalog
+    databases = connector.get_databases(catalog_name="default_catalog")
+    print(f"Databases: {databases}")
+
+    # Get tables
+    tables = connector.get_tables(catalog_name="default_catalog", database_name="mydb")
+    print(f"Tables: {tables}")
+
+    # Get materialized views
+    mvs = connector.get_materialized_views(database_name="mydb")
+    print(f"Materialized Views: {mvs}")
+
+    # Get materialized views with DDL
+    mvs_with_ddl = connector.get_materialized_views_with_ddl(database_name="mydb")
+    for mv in mvs_with_ddl:
+        print(f"\n{mv['table_name']}:")
+        print(mv['definition'])
+
+    # Execute query
+    result = connector.execute_query("SELECT * FROM users LIMIT 10")
+    print(result.sql_return)
+```
+
+## Features
+
+### StarRocks-Specific Features
+- **Multi-catalog support**: Query across multiple catalogs
+- **Materialized views**: Full support for StarRocks materialized views
+- **Catalog management**: Switch between catalogs seamlessly
+
+### Inherited from MySQL
+- Full CRUD operations (SELECT, INSERT, UPDATE, DELETE)
+- DDL execution (CREATE, ALTER, DROP)
+- Metadata retrieval (tables, views, schemas)
+- Sample data extraction
+- Multiple result formats (pandas, arrow, csv, list)
+- Connection pooling and management
+
+## StarRocks-Specific Examples
+
+### Working with Catalogs
+
+```python
+# List all catalogs
+catalogs = connector.get_catalogs()
+
+# Switch catalog
+connector.switch_context(catalog_name="hive_catalog")
+
+# Query with explicit catalog
+tables = connector.get_tables(
+    catalog_name="hive_catalog",
+    database_name="my_hive_db"
+)
+```
+
+### Materialized Views
+
+```python
+# Get materialized views
+mvs = connector.get_materialized_views(database_name="mydb")
+
+# Get materialized views with full DDL
+mvs_with_ddl = connector.get_materialized_views_with_ddl(database_name="mydb")
+
+for mv in mvs_with_ddl:
+    print(f"Name: {mv['table_name']}")
+    print(f"Database: {mv['database_name']}")
+    print(f"Catalog: {mv['catalog_name']}")
+    print(f"Definition: {mv['definition']}")
+```
+
+### Fully-Qualified Names
+
+StarRocks supports three-part names: `catalog.database.table`
+
+```python
+# Build full name
+full_name = connector.full_name(
+    catalog_name="default_catalog",
+    database_name="mydb",
+    table_name="users"
+)
+# Result: `default_catalog`.`mydb`.`users`
+
+# Query with full name
+result = connector.execute_query(f"SELECT * FROM {full_name} LIMIT 10")
+```
+
+## Requirements
+
+- Python >= 3.10
+- StarRocks >= 2.0
+- datus-agent >= 0.3.0
+- datus-mysql >= 0.1.0
+
+## Connection Cleanup
+
+The connector includes special handling for PyMySQL cleanup errors that can occur with StarRocks connections. Use the context manager pattern for automatic cleanup:
+
+```python
+with StarRocksConnector(...) as connector:
+    # Your code here
+    pass
+# Connection automatically cleaned up
+```
+
+## License
+
+Apache License 2.0
+
+## Related Packages
+
+- `datus-mysql` - MySQL adapter (base for StarRocks)
+- `datus-sqlalchemy` - SQLAlchemy base connector
+- `datus-snowflake` - Snowflake adapter

datus_starrocks-0.1.0/README.md

@@ -0,0 +1,177 @@
+# datus-starrocks
+
+StarRocks database adapter for Datus.
+
+## Overview
+
+StarRocks is a high-performance analytical database that uses the MySQL protocol. This adapter extends the MySQL connector with StarRocks-specific features:
+
+- Multi-catalog support
+- Materialized views
+- StarRocks-specific metadata queries
+
+## Installation
+
+```bash
+pip install datus-starrocks
+```
+
+This will automatically install the required dependencies:
+- `datus-agent`
+- `datus-mysql` (which includes `datus-sqlalchemy`)
+
+## Usage
+
+The adapter is automatically registered with Datus when installed. Configure your database connection:
+
+```yaml
+database:
+  type: starrocks
+  host: localhost
+  port: 9030
+  username: root
+  password: your_password
+  catalog: default_catalog
+  database: your_database
+```
+
+Or use programmatically:
+
+```python
+from datus_starrocks import StarRocksConnector
+
+# Create connector
+connector = StarRocksConnector(
+    host="localhost",
+    port=9030,
+    user="root",
+    password="your_password",
+    catalog="default_catalog",
+    database="mydb"
+)
+
+# Use context manager for automatic cleanup
+with connector:
+    # Test connection
+    connector.test_connection()
+
+    # Get catalogs
+    catalogs = connector.get_catalogs()
+    print(f"Catalogs: {catalogs}")
+
+    # Get databases in catalog
+    databases = connector.get_databases(catalog_name="default_catalog")
+    print(f"Databases: {databases}")
+
+    # Get tables
+    tables = connector.get_tables(catalog_name="default_catalog", database_name="mydb")
+    print(f"Tables: {tables}")
+
+    # Get materialized views
+    mvs = connector.get_materialized_views(database_name="mydb")
+    print(f"Materialized Views: {mvs}")
+
+    # Get materialized views with DDL
+    mvs_with_ddl = connector.get_materialized_views_with_ddl(database_name="mydb")
+    for mv in mvs_with_ddl:
+        print(f"\n{mv['table_name']}:")
+        print(mv['definition'])
+
+    # Execute query
+    result = connector.execute_query("SELECT * FROM users LIMIT 10")
+    print(result.sql_return)
+```
+
+## Features
+
+### StarRocks-Specific Features
+- **Multi-catalog support**: Query across multiple catalogs
+- **Materialized views**: Full support for StarRocks materialized views
+- **Catalog management**: Switch between catalogs seamlessly
+
+### Inherited from MySQL
+- Full CRUD operations (SELECT, INSERT, UPDATE, DELETE)
+- DDL execution (CREATE, ALTER, DROP)
+- Metadata retrieval (tables, views, schemas)
+- Sample data extraction
+- Multiple result formats (pandas, arrow, csv, list)
+- Connection pooling and management
+
+## StarRocks-Specific Examples
+
+### Working with Catalogs
+
+```python
+# List all catalogs
+catalogs = connector.get_catalogs()
+
+# Switch catalog
+connector.switch_context(catalog_name="hive_catalog")
+
+# Query with explicit catalog
+tables = connector.get_tables(
+    catalog_name="hive_catalog",
+    database_name="my_hive_db"
+)
+```
+
+### Materialized Views
+
+```python
+# Get materialized views
+mvs = connector.get_materialized_views(database_name="mydb")
+
+# Get materialized views with full DDL
+mvs_with_ddl = connector.get_materialized_views_with_ddl(database_name="mydb")
+
+for mv in mvs_with_ddl:
+    print(f"Name: {mv['table_name']}")
+    print(f"Database: {mv['database_name']}")
+    print(f"Catalog: {mv['catalog_name']}")
+    print(f"Definition: {mv['definition']}")
+```
+
+### Fully-Qualified Names
+
+StarRocks supports three-part names: `catalog.database.table`
+
+```python
+# Build full name
+full_name = connector.full_name(
+    catalog_name="default_catalog",
+    database_name="mydb",
+    table_name="users"
+)
+# Result: `default_catalog`.`mydb`.`users`
+
+# Query with full name
+result = connector.execute_query(f"SELECT * FROM {full_name} LIMIT 10")
+```
+
+## Requirements
+
+- Python >= 3.10
+- StarRocks >= 2.0
+- datus-agent >= 0.3.0
+- datus-mysql >= 0.1.0
+
+## Connection Cleanup
+
+The connector includes special handling for PyMySQL cleanup errors that can occur with StarRocks connections. Use the context manager pattern for automatic cleanup:
+
+```python
+with StarRocksConnector(...) as connector:
+    # Your code here
+    pass
+# Connection automatically cleaned up
+```
+
+## License
+
+Apache License 2.0
+
+## Related Packages
+
+- `datus-mysql` - MySQL adapter (base for StarRocks)
+- `datus-sqlalchemy` - SQLAlchemy base connector
+- `datus-snowflake` - Snowflake adapter

datus_starrocks-0.1.0/datus_starrocks/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2025-present DatusAI, Inc.
+# Licensed under the Apache License, Version 2.0.
+# See http://www.apache.org/licenses/LICENSE-2.0 for details.
+
+from .config import StarRocksConfig
+from .connector import StarRocksConnector
+
+__version__ = "0.1.0"
+__all__ = ["StarRocksConnector", "StarRocksConfig", "register"]
+
+
+def register():
+    """Register StarRocks connector with Datus registry."""
+    from datus.tools.db_tools import connector_registry
+
+    connector_registry.register("starrocks", StarRocksConnector)

datus_starrocks-0.1.0/datus_starrocks/config.py

@@ -0,0 +1,23 @@
+# Copyright 2025-present DatusAI, Inc.
+# Licensed under the Apache License, Version 2.0.
+# See http://www.apache.org/licenses/LICENSE-2.0 for details.
+
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class StarRocksConfig(BaseModel):
+    """StarRocks-specific configuration."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    host: str = Field(..., description="StarRocks server host")
+    port: int = Field(default=9030, description="StarRocks server port")
+    username: str = Field(..., description="StarRocks username")
+    password: str = Field(default="", description="StarRocks password")
+    catalog: str = Field(default="default_catalog", description="Default catalog name")
+    database: Optional[str] = Field(default=None, description="Default database name")
+    charset: str = Field(default="utf8mb4", description="Character set to use")
+    autocommit: bool = Field(default=True, description="Enable autocommit mode")
+    timeout_seconds: int = Field(default=30, description="Connection timeout in seconds")

datus_starrocks-0.1.0/datus_starrocks/connector.py

@@ -0,0 +1,340 @@
+# Copyright 2025-present DatusAI, Inc.
+# Licensed under the Apache License, Version 2.0.
+# See http://www.apache.org/licenses/LICENSE-2.0 for details.
+
+from typing import Any, Dict, List, Union, override
+
+from datus.tools.db_tools.base import list_to_in_str
+from datus.tools.db_tools.mixins import CatalogSupportMixin, MaterializedViewSupportMixin
+from datus.utils.constants import DBType
+from datus.utils.loggings import get_logger
+from datus_mysql import MySQLConnector
+
+from .config import StarRocksConfig
+
+logger = get_logger(__name__)
+
+
+class StarRocksConnector(MySQLConnector, CatalogSupportMixin, MaterializedViewSupportMixin):
+    """
+    StarRocks database connector.
+
+    StarRocks uses MySQL protocol but adds multi-catalog support and materialized views.
+    This connector implements CatalogSupportMixin and MaterializedViewSupportMixin.
+    """
+
+    def __init__(self, config: Union[StarRocksConfig, dict]):
+        """
+        Initialize StarRocks connector.
+
+        Args:
+            config: StarRocksConfig object or dict with configuration
+        """
+        # Handle config object or dict
+        if isinstance(config, dict):
+            config = StarRocksConfig(**config)
+        elif not isinstance(config, StarRocksConfig):
+            raise TypeError(f"config must be StarRocksConfig or dict, got {type(config)}")
+
+        self.starrocks_config = config
+
+        # Pass MySQL config to parent connector
+        from datus_mysql import MySQLConfig
+
+        mysql_config = MySQLConfig(
+            host=config.host,
+            port=config.port,
+            username=config.username,
+            password=config.password,
+            database=config.database or "",
+            charset=config.charset,
+            autocommit=config.autocommit,
+            timeout_seconds=config.timeout_seconds,
+        )
+        super().__init__(mysql_config)
+
+        self.catalog_name = config.catalog
+
+        # Override dialect to StarRocks
+        self.dialect = DBType.STARROCKS
+
+    # ==================== Context Manager Support ====================
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit with cleanup."""
+        self.close()
+        return False  # Don't suppress exceptions
+
+    # ==================== Catalog Management (CatalogSupportMixin) ====================
+
+    @override
+    def default_catalog(self) -> str:
+        """StarRocks default catalog."""
+        return "default_catalog"
+
+    @override
+    def get_catalogs(self) -> List[str]:
+        """Get list of catalogs."""
+        result = self._execute_pandas("SHOW CATALOGS")
+        if result.empty:
+            return []
+        return result["Catalog"].tolist()
+
+    @override
+    def switch_catalog(self, catalog_name: str) -> None:
+        """Switch to a different catalog.
+
+        Args:
+            catalog_name: Name of the catalog to switch to
+        """
+        self.switch_context(catalog_name=catalog_name)
+        self.catalog_name = catalog_name
+
+    def reset_catalog_to_default(self, catalog: str) -> str:
+        """Reset the catalog to the default catalog if it is not set or is 'def'."""
+        if not catalog or catalog == "def":
+            return self.default_catalog()
+        return catalog
+
+    def _before_metadata_query(self, catalog_name: str = "", database_name: str = "") -> None:
+        """Switch catalog before metadata queries if needed."""
+        target_catalog = catalog_name or self.catalog_name or self.default_catalog()
+        if target_catalog and target_catalog != self.catalog_name:
+            self.switch_context(catalog_name=target_catalog)
+
+    # ==================== Metadata Retrieval ====================
+
+    def _get_metadata(
+        self,
+        table_type: str = "table",
+        catalog_name: str = "",
+        database_name: str = "",
+    ) -> List[Dict[str, str]]:
+        """
+        Get metadata for tables/views with catalog support.
+
+        Args:
+            table_type: Type of object (table, view, mv)
+            catalog_name: Catalog name
+            database_name: Database name to query
+
+        Returns:
+            List of metadata dictionaries with catalog_name properly set
+        """
+        # Determine the target catalog
+        current_catalog = self.reset_catalog_to_default(catalog_name or self.catalog_name)
+
+        # Switch to the correct catalog before querying
+        self._before_metadata_query(catalog_name=current_catalog, database_name=database_name)
+
+        # Get base metadata from parent
+        result = super()._get_metadata(table_type, catalog_name, database_name)
+
+        # Set the correct catalog_name and filter results by catalog as safety check
+        filtered_result = []
+        for item in result:
+            # Filter by catalog if the item has catalog_name set
+            if "catalog_name" in item and item["catalog_name"] and item["catalog_name"] != current_catalog:
+                continue
+
+            item["catalog_name"] = current_catalog
+            # Update identifier to include catalog
+            item["identifier"] = self.identifier(
+                catalog_name=current_catalog, database_name=item["database_name"], table_name=item["table_name"]
+            )
+            filtered_result.append(item)
+
+        return filtered_result
+
+    @override
+    def get_tables(self, catalog_name: str = "", database_name: str = "", schema_name: str = "") -> List[str]:
+        """Get list of table names."""
+        result = self._get_metadata(table_type="table", catalog_name=catalog_name, database_name=database_name)
+        return [table["table_name"] for table in result]
+
+    @override
+    def get_views(self, catalog_name: str = "", database_name: str = "", schema_name: str = "") -> List[str]:
+        """Get list of view names."""
+        try:
+            result = self._get_metadata(table_type="view", catalog_name=catalog_name, database_name=database_name)
+            return [view["table_name"] for view in result]
+        except Exception as e:
+            logger.warning(f"Failed to get views: {e}")
+            return []
+
+    def get_materialized_views(
+        self, catalog_name: str = "", database_name: str = "", schema_name: str = ""
+    ) -> List[str]:
+        """Get list of materialized view names."""
+        try:
+            result = self._get_metadata(table_type="mv", catalog_name=catalog_name, database_name=database_name)
+            return [mv["table_name"] for mv in result]
+        except Exception as e:
+            logger.warning(f"Failed to get materialized views: {e}")
+            return []
+
+    def get_materialized_views_with_ddl(
+        self, catalog_name: str = "", database_name: str = "", schema_name: str = ""
+    ) -> List[Dict[str, str]]:
+        """
+        Get materialized views with DDL definitions.
+
+        Args:
+            catalog_name: Catalog name
+            database_name: Database name
+            schema_name: Schema name (unused in StarRocks)
+
+        Returns:
+            List of materialized view metadata with DDL
+        """
+        current_catalog = self.reset_catalog_to_default(catalog_name or self.catalog_name)
+
+        self._before_metadata_query(catalog_name=current_catalog, database_name=database_name)
+
+        # Query materialized views from information_schema
+        query_sql = (
+            "SELECT TABLE_SCHEMA, TABLE_NAME, MATERIALIZED_VIEW_DEFINITION "
+            "FROM information_schema.materialized_views"
+        )
+
+        if database_name:
+            query_sql = f"{query_sql} WHERE TABLE_SCHEMA = '{database_name}'"
+        else:
+            ignore_dbs = list(self._sys_databases())
+            query_sql = f"{query_sql} {list_to_in_str('WHERE TABLE_SCHEMA NOT IN', ignore_dbs)}"
+
+        result = self._execute_pandas(query_sql)
+
+        mv_list = []
+        for i in range(len(result)):
+            mv_list.append(
+                {
+                    "identifier": self.identifier(
+                        catalog_name=current_catalog,
+                        database_name=str(result["TABLE_SCHEMA"][i]),
+                        table_name=str(result["TABLE_NAME"][i]),
+                    ),
+                    "catalog_name": current_catalog,
+                    "database_name": result["TABLE_SCHEMA"][i],
+                    "schema_name": "",
+                    "table_name": result["TABLE_NAME"][i],
+                    "definition": result["MATERIALIZED_VIEW_DEFINITION"][i],
+                    "table_type": "mv",
+                }
+            )
+
+        return mv_list
+
+    # ==================== Database Management ====================
+
+    @override
+    def get_databases(self, catalog_name: str = "", include_sys: bool = False) -> List[str]:
+        """Get list of databases in the catalog."""
+        return super().get_databases(catalog_name, include_sys=include_sys)
+
+    # ==================== Full Name Construction ====================
+
+    @override
+    def full_name(
+        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", table_name: str = ""
+    ) -> str:
+        """
+        Build fully-qualified table name with catalog support.
+
+        StarRocks format: `catalog`.`database`.`table`
+        """
+        catalog_name = self.reset_catalog_to_default(catalog_name)
+
+        if catalog_name:
+            if database_name:
+                return f"`{catalog_name}`.`{database_name}`.`{table_name}`"
+            else:
+                return f"`{table_name}`"
+        else:
+            if database_name:
+                return f"`{database_name}`.`{table_name}`"
+            return f"`{table_name}`"
+
+    @override
+    def _sqlalchemy_schema(self, catalog_name: str = "", database_name: str = "", schema_name: str = "") -> str:
+        """Get schema name for SQLAlchemy Inspector with catalog support."""
+        database_name = database_name or self.database_name
+
+        if self.support_catalog():
+            catalog_name = catalog_name or self.catalog_name or self.default_catalog()
+            if database_name:
+                return f"{catalog_name}.{database_name}"
+            return None
+        else:
+            return database_name if database_name else None
+
+    # ==================== Connection Cleanup ====================
+
+    @override
+    def close(self):
+        """
+        Close connection with special handling for PyMySQL cleanup errors.
+
+        StarRocks may trigger PyMySQL struct.pack errors during cleanup,
+        which we safely ignore.
+        """
+        try:
+            super().close()
+        except Exception as e:
+            error_str = str(e)
+
+            # Check for known PyMySQL cleanup errors
+            pymysql_errors = ["struct.error", "struct.pack", "COMMAND.COM_QUIT", "required argument is not an integer"]
+
+            if any(err in error_str for err in pymysql_errors):
+                logger.debug(f"Ignoring PyMySQL cleanup error: {e}")
+
+                # Force cleanup of connection variables
+                if hasattr(self, "connection"):
+                    self.connection = None
+                if hasattr(self, "engine"):
+                    try:
+                        if self.engine:
+                            self.engine.dispose()
+                    except Exception:
+                        pass
+                    finally:
+                        self.engine = None
+            else:
+                # Re-raise unexpected errors
+                logger.error(f"Unexpected close error: {e}")
+                raise
+
+    # ==================== Utility Methods ====================
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert connector to serializable dictionary."""
+        return {
+            "db_type": DBType.STARROCKS,
+            "host": self.host,
+            "port": self.port,
+            "user": self.user,
+            "catalog": self.catalog_name,
+            "database": self.database_name,
+        }
+
+    def get_type(self) -> str:
+        """Return the database type."""
+        return DBType.STARROCKS
+
+    @override
+    def test_connection(self) -> bool:
+        """Test the database connection with proper cleanup."""
+        try:
+            return super().test_connection()
+        finally:
+            # Ensure connection is closed after test
+            try:
+                self.close()
+            except Exception as e:
+                logger.debug(f"Ignoring cleanup error during test: {e}")

datus_starrocks-0.1.0/pyproject.toml

@@ -0,0 +1,47 @@
+[project]
+name = "datus-starrocks"
+version = "0.1.0"
+description = "StarRocks database adapter for Datus"
+readme = "README.md"
+requires-python = ">=3.12"
+license = {text = "Apache-2.0"}
+authors = [
+    {name = "DatusAI", email = "support@datus.ai"}
+]
+keywords = ["datus", "database", "starrocks", "adapter"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+]
+
+dependencies = [
+    "datus-agent>0.2.1",
+    "datus-mysql>=0.1.0",
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/Datus-ai/datus-db-adapters"
+Repository = "https://github.com/Datus-ai/datus-db-adapters"
+Issues = "https://github.com/Datus-ai/datus-db-adapters/issues"
+
+[project.entry-points."datus.adapters"]
+starrocks = "datus_starrocks:register"
+
+[tool.uv.sources]
+datus-mysql = { workspace = true }
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["datus_starrocks"]

datus_starrocks-0.1.0/tests/test_connector.py

@@ -0,0 +1,429 @@
+# Copyright 2025-present DatusAI, Inc.
+# Licensed under the Apache License, Version 2.0.
+# See http://www.apache.org/licenses/LICENSE-2.0 for details.
+
+import os
+import uuid
+from typing import Generator
+
+import pytest
+from datus.tools.db_tools.mixins import CatalogSupportMixin, MaterializedViewSupportMixin
+from datus.utils.exceptions import DatusException, ErrorCode
+from datus_starrocks import StarRocksConfig, StarRocksConnector
+
+
+@pytest.fixture
+def config() -> StarRocksConfig:
+    """Create StarRocks configuration from environment or defaults."""
+    return StarRocksConfig(
+        host=os.getenv("STARROCKS_HOST", "localhost"),
+        port=int(os.getenv("STARROCKS_PORT", "9030")),
+        username=os.getenv("STARROCKS_USER", "root"),
+        password=os.getenv("STARROCKS_PASSWORD", ""),
+        catalog=os.getenv("STARROCKS_CATALOG", "default_catalog"),
+        database=os.getenv("STARROCKS_DATABASE", "quickstart"),
+    )
+
+
+@pytest.fixture
+def connector(config: StarRocksConfig) -> Generator[StarRocksConnector, None, None]:
+    """Create and cleanup StarRocks connector."""
+    conn = StarRocksConnector(config)
+    yield conn
+    conn.close()
+
+
+# ==================== Mixin Tests ====================
+
+
+def test_connector_implements_catalog_mixin(connector: StarRocksConnector):
+    """Verify StarRocks connector implements CatalogSupportMixin."""
+    assert isinstance(connector, CatalogSupportMixin)
+
+
+def test_connector_implements_materialized_view_mixin(connector: StarRocksConnector):
+    """Verify StarRocks connector implements MaterializedViewSupportMixin."""
+    assert isinstance(connector, MaterializedViewSupportMixin)
+
+
+# ==================== Connection Tests ====================
+
+
+def test_connection_with_config_object(config: StarRocksConfig):
+    """Test connection using config object."""
+    conn = StarRocksConnector(config)
+    assert conn.test_connection()
+    conn.close()
+
+
+def test_connection_with_dict():
+    """Test connection using dict config."""
+    conn = StarRocksConnector(
+        {
+            "host": os.getenv("STARROCKS_HOST", "localhost"),
+            "port": int(os.getenv("STARROCKS_PORT", "9030")),
+            "username": os.getenv("STARROCKS_USER", "root"),
+            "password": os.getenv("STARROCKS_PASSWORD", ""),
+        }
+    )
+    assert conn.test_connection()
+    conn.close()
+
+
+def test_context_manager(config: StarRocksConfig):
+    """Test connector as context manager."""
+    with StarRocksConnector(config) as conn:
+        assert conn.test_connection()
+
+
+# ==================== Catalog Tests (CatalogSupportMixin) ====================
+
+
+def test_get_catalogs(connector: StarRocksConnector):
+    """Test getting list of catalogs."""
+    catalogs = connector.get_catalogs()
+    assert len(catalogs) > 0
+    assert connector.default_catalog() in catalogs
+
+
+def test_default_catalog(connector: StarRocksConnector):
+    """Test default catalog."""
+    assert connector.default_catalog() == "default_catalog"
+
+
+def test_switch_catalog(connector: StarRocksConnector):
+    """Test switching catalogs."""
+    original_catalog = connector.catalog_name
+    catalogs = connector.get_catalogs()
+
+    if len(catalogs) > 1:
+        target_catalog = [c for c in catalogs if c != original_catalog][0]
+        connector.switch_catalog(target_catalog)
+        assert connector.catalog_name == target_catalog
+
+        # Switch back
+        connector.switch_catalog(original_catalog)
+        assert connector.catalog_name == original_catalog
+
+
+# ==================== Database Tests ====================
+
+
+def test_get_databases(connector: StarRocksConnector):
+    """Test getting list of databases."""
+    databases = connector.get_databases()
+    assert len(databases) > 0
+
+
+# ==================== Table Metadata Tests ====================
+
+
+def test_get_tables(connector: StarRocksConnector):
+    """Test getting table list."""
+    tables = connector.get_tables()
+    assert isinstance(tables, list)
+
+
+def test_get_tables_with_ddl(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test getting tables with DDL."""
+    tables = connector.get_tables_with_ddl(catalog_name=config.catalog)
+
+    if len(tables) > 0:
+        table = tables[0]
+        assert "table_name" in table
+        assert "definition" in table
+        assert table["table_type"] == "table"
+        assert "database_name" in table
+        assert table["schema_name"] == ""
+        assert table["catalog_name"] == config.catalog
+        assert "identifier" in table
+        assert len(table["identifier"].split(".")) == 3
+
+
+# ==================== View Tests ====================
+
+
+def test_get_views(connector: StarRocksConnector):
+    """Test getting view list."""
+    views = connector.get_views()
+    assert isinstance(views, list)
+
+
+def test_get_views_with_ddl(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test getting views with DDL."""
+    views = connector.get_views_with_ddl(catalog_name=config.catalog)
+
+    if len(views) > 0:
+        view = views[0]
+        assert "table_name" in view
+        assert "definition" in view
+        assert view["table_type"] == "view"
+        assert "database_name" in view
+        assert view["schema_name"] == ""
+        assert "catalog_name" in view
+
+        identifier_parts = view["identifier"].split(".")
+        assert len(identifier_parts) == 3
+        assert identifier_parts[0] == view["catalog_name"]
+        assert identifier_parts[1] == view["database_name"]
+        assert identifier_parts[2] == view["table_name"]
+
+
+# ==================== Materialized View Tests (MaterializedViewSupportMixin) ====================
+
+
+def test_get_materialized_views(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test getting materialized view list."""
+    mvs = connector.get_materialized_views(catalog_name=config.catalog)
+    assert isinstance(mvs, list)
+
+
+def test_get_materialized_views_with_ddl(connector: StarRocksConnector):
+    """Test getting materialized views with DDL."""
+    mvs = connector.get_materialized_views_with_ddl()
+
+    if len(mvs) > 0:
+        mv = mvs[0]
+        assert "table_name" in mv
+        assert "definition" in mv
+        assert mv["table_type"] == "mv"
+        assert "database_name" in mv
+        assert mv["schema_name"] == ""
+        assert "catalog_name" in mv
+
+        identifier_parts = mv["identifier"].split(".")
+        assert len(identifier_parts) == 3
+
+
+# ==================== Sample Data Tests ====================
+
+
+def test_get_sample_rows_default(connector: StarRocksConnector):
+    """Test getting sample rows with defaults."""
+    sample_rows = connector.get_sample_rows()
+    assert isinstance(sample_rows, list)
+
+
+def test_get_sample_rows_with_database(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test getting sample rows for specific database."""
+    sample_rows = connector.get_sample_rows(catalog_name=config.catalog, database_name=config.database)
+
+    if len(sample_rows) > 0:
+        item = sample_rows[0]
+        assert "database_name" in item
+        assert "table_name" in item
+        assert "catalog_name" in item
+        assert item["schema_name"] == ""
+        assert "identifier" in item
+        assert len(item["identifier"].split(".")) == 3
+        assert "sample_rows" in item
+
+
+def test_get_sample_rows_specific_tables(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test getting sample rows for specific tables."""
+    # First get available tables
+    tables = connector.get_tables(catalog_name=config.catalog, database_name=config.database)
+
+    if len(tables) > 0:
+        table_name = tables[0]
+        sample_rows = connector.get_sample_rows(
+            catalog_name=config.catalog, database_name=config.database, tables=[table_name], top_n=3
+        )
+
+        assert len(sample_rows) == 1
+        assert sample_rows[0]["table_name"] == table_name
+
+
+# ==================== SQL Execution Tests ====================
+
+
+def test_execute_query(connector: StarRocksConnector):
+    """Test executing simple query."""
+    result = connector.execute({"sql_query": "SELECT 1 as num"}, result_format="list")
+    assert result.success
+    assert not result.error
+    assert result.sql_return == [{"num": 1}]
+
+
+def test_execute_explain(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test executing EXPLAIN query."""
+    tables = connector.get_tables(catalog_name=config.catalog, database_name=config.database)
+
+    if len(tables) > 0:
+        table_name = tables[0]
+        full_name = connector.full_name(
+            catalog_name=config.catalog, database_name=config.database, table_name=table_name
+        )
+
+        result = connector.execute({"sql_query": f"EXPLAIN SELECT * FROM {full_name} LIMIT 1"})
+        assert result.success
+        assert not result.error
+        assert result.sql_return
+
+
+def test_execute_ddl_create_drop(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test DDL operations (CREATE/DROP)."""
+    suffix = uuid.uuid4().hex[:8]
+    table_name = f"datus_test_{suffix}"
+
+    connector.switch_context(database_name=config.database)
+
+    create_sql = f"""
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        `id` BIGINT NOT NULL,
+        `name` VARCHAR(64)
+    ) ENGINE=OLAP
+    PRIMARY KEY (`id`)
+    DISTRIBUTED BY HASH(`id`) BUCKETS 1
+    PROPERTIES (
+        "replication_num" = "1"
+    );
+    """
+
+    try:
+        create_result = connector.execute_ddl(create_sql)
+        assert create_result.success, f"Failed to create table: {create_result.error}"
+    finally:
+        connector.execute_ddl(f"DROP TABLE IF EXISTS {table_name}")
+
+
+def test_execute_insert(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test INSERT operation."""
+    suffix = uuid.uuid4().hex[:8]
+    table_name = f"datus_insert_test_{suffix}"
+
+    connector.switch_context(database_name=config.database)
+
+    create_sql = f"""
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        `id` BIGINT NOT NULL,
+        `name` VARCHAR(64)
+    ) ENGINE=OLAP
+    PRIMARY KEY (`id`)
+    DISTRIBUTED BY HASH(`id`) BUCKETS 1
+    PROPERTIES (
+        "replication_num" = "1"
+    );
+    """
+
+    try:
+        create_result = connector.execute_ddl(create_sql)
+        if not create_result.success:
+            pytest.skip(f"Unable to create test table: {create_result.error}")
+
+        # Insert data
+        insert_result = connector.execute_insert(f"INSERT INTO {table_name} (id, name) VALUES (1, 'Alice'), (2, 'Bob')")
+        assert insert_result.success
+
+        # Verify
+        query_result = connector.execute(
+            {"sql_query": f"SELECT id, name FROM {table_name} ORDER BY id"}, result_format="list"
+        )
+        assert query_result.success
+        assert query_result.sql_return == [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]
+    finally:
+        connector.execute_ddl(f"DROP TABLE IF EXISTS {table_name}")
+
+
+def test_execute_update(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test UPDATE operation."""
+    suffix = uuid.uuid4().hex[:8]
+    table_name = f"datus_update_test_{suffix}"
+
+    connector.switch_context(database_name=config.database)
+
+    create_sql = f"""
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        `id` BIGINT NOT NULL,
+        `name` VARCHAR(64)
+    ) ENGINE=OLAP
+    PRIMARY KEY (`id`)
+    DISTRIBUTED BY HASH(`id`) BUCKETS 1
+    PROPERTIES (
+        "replication_num" = "1"
+    );
+    """
+
+    try:
+        create_result = connector.execute_ddl(create_sql)
+        if not create_result.success:
+            pytest.skip(f"Unable to create test table: {create_result.error}")
+
+        # Insert initial data
+        connector.execute_insert(f"INSERT INTO {table_name} (id, name) VALUES (1, 'Alice'), (2, 'Bob')")
+
+        # Update
+        update_result = connector.execute(
+            {"sql_query": f"UPDATE {table_name} SET name = 'Alice Updated' WHERE id = 1"}, result_format="list"
+        )
+        assert update_result.success
+
+        # Verify
+        query_result = connector.execute(
+            {"sql_query": f"SELECT id, name FROM {table_name} ORDER BY id"}, result_format="list"
+        )
+        assert query_result.sql_return == [{"id": 1, "name": "Alice Updated"}, {"id": 2, "name": "Bob"}]
+    finally:
+        connector.execute_ddl(f"DROP TABLE IF EXISTS {table_name}")
+
+
+def test_execute_delete(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test DELETE operation."""
+    suffix = uuid.uuid4().hex[:8]
+    table_name = f"datus_delete_test_{suffix}"
+
+    connector.switch_context(database_name=config.database)
+
+    create_sql = f"""
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        `id` BIGINT NOT NULL,
+        `name` VARCHAR(64)
+    ) ENGINE=OLAP
+    PRIMARY KEY (`id`)
+    DISTRIBUTED BY HASH(`id`) BUCKETS 1
+    PROPERTIES (
+        "replication_num" = "1"
+    );
+    """
+
+    try:
+        create_result = connector.execute_ddl(create_sql)
+        if not create_result.success:
+            pytest.skip(f"Unable to create test table: {create_result.error}")
+
+        # Insert initial data
+        connector.execute_insert(f"INSERT INTO {table_name} (id, name) VALUES (1, 'Alice'), (2, 'Bob')")
+
+        # Delete
+        delete_result = connector.execute({"sql_query": f"DELETE FROM {table_name} WHERE id = 2"}, result_format="list")
+        assert delete_result.success
+
+        # Verify
+        query_result = connector.execute(
+            {"sql_query": f"SELECT id, name FROM {table_name} ORDER BY id"}, result_format="list"
+        )
+        assert query_result.sql_return == [{"id": 1, "name": "Alice"}]
+    finally:
+        connector.execute_ddl(f"DROP TABLE IF EXISTS {table_name}")
+
+
+# ==================== Error Handling Tests ====================
+
+
+def test_exception_on_nonexistent_table(connector: StarRocksConnector, config: StarRocksConfig):
+    """Test exception handling for non-existent table."""
+    with pytest.raises(DatusException, match=ErrorCode.DB_EXECUTION_ERROR.code):
+        connector.get_sample_rows(catalog_name=config.catalog, tables=["nonexistent_table_" + uuid.uuid4().hex])
+
+
+def test_execute_merge_returns_error(connector: StarRocksConnector):
+    """Test MERGE statement error handling."""
+    merge_sql = (
+        "MERGE INTO nonexistent_target AS t USING nonexistent_source AS s ON t.id = s.id "
+        "WHEN MATCHED THEN UPDATE SET t.value = s.value "
+        "WHEN NOT MATCHED THEN INSERT (id, value) VALUES (s.id, s.value)"
+    )
+
+    result = connector.execute({"sql_query": merge_sql})
+    assert result.sql_query == merge_sql
+    assert not result.success or result.error  # Either fails or returns error