datus-spark 0.1.0__tar.gz

datus_spark-0.1.0/.gitignore

@@ -0,0 +1,140 @@
+# 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
+
+
+.omc

datus_spark-0.1.0/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: datus-spark
+Version: 0.1.0
+Summary: Spark SQL 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,hive,spark
+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-sqlalchemy>=0.1.0
+Requires-Dist: pure-sasl>=0.6.2
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: pyhive>=0.7.0
+Requires-Dist: thrift-sasl>=0.4.3
+Requires-Dist: thrift>=0.16.0
+Provides-Extra: test
+Requires-Dist: pytest-cov>=4.0.0; extra == 'test'
+Requires-Dist: pytest>=7.0.0; extra == 'test'

datus_spark-0.1.0/README.md

@@ -0,0 +1,203 @@
+# datus-spark
+
+Spark SQL database adapter for Datus, connecting via HiveServer2/Thrift protocol.
+
+## Installation
+
+```bash
+pip install datus-spark
+```
+
+This will automatically install the required dependencies:
+- `datus-agent`
+- `datus-sqlalchemy`
+- `pyhive`
+- `thrift`
+- `thrift-sasl`
+- `pure-sasl`
+
+## Usage
+
+The adapter is automatically registered with Datus when installed. Configure your database connection in your Datus configuration:
+
+```yaml
+database:
+  type: spark
+  host: localhost
+  port: 10000
+  username: spark
+  database: default
+  auth_mechanism: NONE
+```
+
+Or use programmatically:
+
+```python
+from datus_spark import SparkConnector, SparkConfig
+
+# Using config object
+config = SparkConfig(
+    host="localhost",
+    port=10000,
+    username="spark",
+    password="",
+    database="default",
+    auth_mechanism="NONE",
+)
+connector = SparkConnector(config)
+
+# Or using dict
+connector = SparkConnector({
+    "host": "localhost",
+    "port": 10000,
+    "username": "spark",
+    "database": "default",
+})
+
+# Test connection
+connector.test_connection()
+
+# Execute query
+result = connector.execute({"sql_query": "SELECT * FROM `default`.`my_table` LIMIT 10"})
+print(result.sql_return)
+
+# Get table list
+tables = connector.get_tables(database_name="default")
+print(f"Tables: {tables}")
+
+# Get table schema
+schema = connector.get_schema(database_name="default", table_name="my_table")
+for column in schema:
+    print(f"{column['name']}: {column['type']}")
+```
+
+## Configuration Options
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| host | str | "127.0.0.1" | Spark Thrift Server host |
+| port | int | 10000 | Spark Thrift Server port |
+| username | str | (required) | Username |
+| password | str | "" | Password |
+| database | str | None | Default database (falls back to `default`) |
+| auth_mechanism | str | "NONE" | Authentication mechanism (NONE, PLAIN, KERBEROS) |
+| timeout_seconds | int | 30 | Connection timeout |
+
+## Features
+
+- Query execution via Spark SQL (SELECT)
+- DDL execution (CREATE, ALTER, DROP)
+- Metadata retrieval (databases, tables, views, columns)
+- Sample data extraction
+- Multiple result formats (pandas, arrow, csv, list)
+- Connection pooling and management
+- Context manager support
+
+## Testing
+
+### Quick Start
+
+```bash
+cd datus-spark
+
+# Unit tests (no database required)
+uv run pytest tests/ -m "not integration" -v
+
+# All tests with coverage
+uv run pytest tests/ -v --cov=datus_spark --cov-report=term-missing
+```
+
+### Integration Tests (Requires Spark Thrift Server)
+
+```bash
+# Start Spark Thrift Server container
+docker compose up -d
+
+# Wait for container to become healthy (~60s)
+docker compose ps
+
+# Run integration tests
+uv run pytest tests/integration/ -v
+
+# Run only TPC-H tests
+uv run pytest tests/integration/test_tpch.py -v
+
+# Run acceptance tests (core functionality)
+uv run pytest tests/ -m acceptance -v
+
+# Stop Spark
+docker compose down
+```
+
+### TPC-H Test Data
+
+Integration tests include TPC-H benchmark data for realistic query testing. The `tpch_setup` fixture (session-scoped) automatically creates 5 tables with sample data:
+
+| Table | Rows | Description |
+|-------|------|-------------|
+| `tpch_region` | 5 | Standard TPC-H regions |
+| `tpch_nation` | 25 | Standard TPC-H nations |
+| `tpch_customer` | 10 | Simplified customer data |
+| `tpch_orders` | 15 | Simplified order data |
+| `tpch_supplier` | 5 | Simplified supplier data |
+
+Tables are created at the start of the test session and dropped after all tests complete.
+
+#### Initialize TPC-H Data Manually
+
+To create TPC-H data for use with Datus (outside of tests):
+
+```bash
+# Basic usage
+uv run python scripts/init_tpch_data.py
+
+# Drop existing tables and re-create
+uv run python scripts/init_tpch_data.py --drop
+
+# Custom connection
+uv run python scripts/init_tpch_data.py --host 192.168.1.100 --port 10000
+```
+
+### Test Statistics
+
+- **Unit Tests**: 46 tests (config validation, connector logic, identifiers)
+- **Integration Tests**: 24 tests (connection, metadata, SQL execution, TPC-H)
+- **Total**: 70 tests
+
+### Test Markers
+
+| Marker | Description |
+|--------|-------------|
+| `integration` | Requires running Spark Thrift Server |
+| `acceptance` | Core functionality validation for CI/CD |
+
+## Development
+
+### Setup
+
+```bash
+# From workspace root
+uv sync --all-packages
+
+# Or install in editable mode
+uv pip install -e .
+```
+
+### Code Quality
+
+```bash
+black datus_spark tests
+isort datus_spark tests
+ruff check datus_spark tests
+```
+
+## Requirements
+
+- Python >= 3.12
+- Apache Spark >= 3.0 with Thrift Server enabled
+- datus-agent > 0.2.1
+- datus-sqlalchemy >= 0.1.0
+
+## License
+
+Apache License 2.0

datus_spark-0.1.0/datus_spark/__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 SparkConfig
+from .connector import SparkConnector
+
+__version__ = "0.1.0"
+__all__ = ["SparkConnector", "SparkConfig", "register"]
+
+
+def register():
+    """Register Spark connector with Datus registry."""
+    from datus.tools.db_tools import connector_registry
+
+    connector_registry.register("spark", SparkConnector, config_class=SparkConfig)

datus_spark-0.1.0/datus_spark/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 Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class SparkConfig(BaseModel):
+    """Spark SQL (via HiveServer2/Thrift) specific configuration."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    host: str = Field(default="127.0.0.1", description="Spark Thrift Server host")
+    port: int = Field(default=10000, description="Spark Thrift Server port")
+    username: str = Field(..., description="Spark username")
+    password: str = Field(default="", description="Spark password", json_schema_extra={"input_type": "password"})
+    database: Optional[str] = Field(default=None, description="Default database name")
+    auth_mechanism: Literal["NONE", "PLAIN", "KERBEROS"] = Field(
+        default="NONE", description="Authentication mechanism (NONE, PLAIN, KERBEROS)"
+    )
+    timeout_seconds: int = Field(default=30, description="Connection timeout in seconds")

datus_spark-0.1.0/datus_spark/connector.py

@@ -0,0 +1,236 @@
+# 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, Optional, Set, Union, override
+from urllib.parse import quote_plus
+
+from datus.utils.loggings import get_logger
+from datus_sqlalchemy import SQLAlchemyConnector
+
+from .config import SparkConfig
+
+logger = get_logger(__name__)
+
+SPARK_DIALECT = "spark"
+
+
+class SparkConnector(SQLAlchemyConnector):
+    """
+    Spark SQL database connector via HiveServer2/Thrift protocol.
+
+    Spark uses a two-level hierarchy: database -> table.
+    Connects via the Hive SQLAlchemy dialect (pyhive).
+    """
+
+    def __init__(self, config: Union[SparkConfig, dict]):
+        """
+        Initialize Spark connector.
+
+        Args:
+            config: SparkConfig object or dict with configuration
+        """
+        if isinstance(config, dict):
+            config = SparkConfig(**config)
+        elif not isinstance(config, SparkConfig):
+            raise TypeError(f"config must be SparkConfig or dict, got {type(config)}")
+
+        self.spark_config = config
+        self.host = config.host
+        self.port = config.port
+        self.user = config.username
+
+        database = config.database or "default"
+
+        # Build connection string: hive://user:pass@host:port/database
+        encoded_username = quote_plus(config.username)
+        encoded_password = quote_plus(config.password) if config.password else ""
+        if config.password:
+            auth_part = f"{encoded_username}:{encoded_password}@"
+        else:
+            auth_part = f"{encoded_username}@"
+
+        # Build connection string with auth mechanism
+        connection_string = f"hive://{auth_part}{config.host}:{config.port}/{database}"
+
+        if config.auth_mechanism and config.auth_mechanism != "NONE":
+            connection_string += f"?auth={config.auth_mechanism}"
+
+        super().__init__(connection_string, dialect=SPARK_DIALECT, timeout_seconds=config.timeout_seconds)
+
+        self.dialect = SPARK_DIALECT
+        self.database_name = database
+
+    # ==================== 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
+
+    # ==================== System Resources ====================
+
+    @override
+    def _sys_databases(self) -> Set[str]:
+        """System databases to filter out."""
+        return {"information_schema"}
+
+    @override
+    def _sys_schemas(self) -> Set[str]:
+        """System schemas to filter out (same as databases for Spark)."""
+        return self._sys_databases()
+
+    # ==================== Metadata Retrieval ====================
+
+    @override
+    def get_databases(self, catalog_name: str = "", include_sys: bool = False) -> List[str]:
+        """Get list of databases."""
+        result = self._execute_pandas("SHOW DATABASES")
+        if result.empty:
+            return []
+        databases = result.iloc[:, 0].tolist()
+        if not include_sys:
+            sys_dbs = self._sys_databases()
+            databases = [d for d in databases if d.lower() not in sys_dbs]
+        return databases
+
+    @override
+    def get_schemas(self, catalog_name: str = "", database_name: str = "", include_sys: bool = False) -> List[str]:
+        """Spark doesn't have separate schemas, return empty list."""
+        return []
+
+    @override
+    def get_tables(self, catalog_name: str = "", database_name: str = "", schema_name: str = "") -> List[str]:
+        """Get list of table names."""
+        db = database_name or self.database_name
+        result = self._execute_pandas(f"SHOW TABLES IN {self._quote_identifier(db)}")
+        if result.empty:
+            return []
+        # SHOW TABLES returns (namespace, tableName, isTemporary) in Spark 3.x
+        # Use the second column (tableName) when available, otherwise first
+        if len(result.columns) >= 2:
+            name_col = result.columns[1]
+        else:
+            name_col = result.columns[0]
+        return result[name_col].tolist()
+
+    @override
+    def get_views(self, catalog_name: str = "", database_name: str = "", schema_name: str = "") -> List[str]:
+        """Get list of view names."""
+        db = database_name or self.database_name
+        try:
+            result = self._execute_pandas(f"SHOW VIEWS IN {self._quote_identifier(db)}")
+            if result.empty:
+                return []
+            if len(result.columns) >= 2:
+                name_col = result.columns[1]
+            else:
+                name_col = result.columns[0]
+            return result[name_col].tolist()
+        except Exception as e:
+            logger.warning(f"Failed to get views: {e}")
+            return []
+
+    @override
+    def get_schema(
+        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", table_name: str = ""
+    ) -> List[Dict[str, Any]]:
+        """Get table schema information using DESCRIBE."""
+        if not table_name:
+            return []
+
+        db = database_name or self.database_name
+        full_name = self.full_name(database_name=db, table_name=table_name)
+
+        query_result = self._execute_pandas(f"DESCRIBE {full_name}")
+
+        result = []
+        for i in range(len(query_result)):
+            col_name = query_result.iloc[i, 0]
+            # Skip partition/metadata separator lines
+            if col_name is None or str(col_name).startswith("#") or str(col_name).strip() == "":
+                continue
+            result.append(
+                {
+                    "cid": len(result),
+                    "name": col_name,
+                    "type": str(query_result.iloc[i, 1]) if len(query_result.columns) > 1 else "",
+                    "nullable": True,  # Spark doesn't expose nullable in DESCRIBE
+                    "default_value": None,
+                    "pk": False,
+                    "comment": str(query_result.iloc[i, 2]) if len(query_result.columns) > 2 else None,
+                }
+            )
+        return result
+
+    # ==================== Database Management ====================
+
+    @override
+    def _sqlalchemy_schema(
+        self, catalog_name: str = "", database_name: str = "", schema_name: str = ""
+    ) -> Optional[str]:
+        """Get schema name for SQLAlchemy Inspector (database name in Spark)."""
+        return database_name or self.database_name
+
+    @override
+    def do_switch_context(self, catalog_name: str = "", database_name: str = "", schema_name: str = ""):
+        """Switch database context using USE statement."""
+        if database_name:
+            from sqlalchemy import text
+
+            with self.engine.connect() as conn:
+                conn.execute(text(f"USE {self._quote_identifier(database_name)}"))
+                conn.commit()
+            self.database_name = database_name
+
+    # ==================== Utility Methods ====================
+
+    @staticmethod
+    def _quote_identifier(identifier: str) -> str:
+        """Safely wrap identifiers with backticks for Spark."""
+        escaped = identifier.replace("`", "``")
+        return f"`{escaped}`"
+
+    @override
+    def full_name(
+        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", table_name: str = ""
+    ) -> str:
+        """
+        Build fully-qualified table name.
+
+        Spark format: `database`.`table`
+        """
+        db = database_name or self.database_name
+        if db:
+            return f"{self._quote_identifier(db)}.{self._quote_identifier(table_name)}"
+        return self._quote_identifier(table_name)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert connector to serializable dictionary."""
+        return {
+            "db_type": SPARK_DIALECT,
+            "host": self.host,
+            "port": self.port,
+            "user": self.user,
+            "database": self.database_name,
+        }
+
+    def get_type(self) -> str:
+        """Return the database type."""
+        return SPARK_DIALECT
+
+    @override
+    def test_connection(self) -> bool:
+        """Test the database connection."""
+        try:
+            return super().test_connection()
+        finally:
+            try:
+                self.close()
+            except Exception as e:
+                logger.debug(f"Ignoring cleanup error during test: {e}")

datus_spark-0.1.0/docker-compose.yml

@@ -0,0 +1,25 @@
+services:
+  spark-thrift:
+    image: apache/spark:3.5.0
+    container_name: datus-spark-test
+    command: >
+      /opt/spark/sbin/start-thriftserver.sh
+      --master local[*]
+      --hiveconf hive.server2.thrift.port=10000
+      --hiveconf hive.server2.thrift.bind.host=0.0.0.0
+    ports:
+      - "10000:10000"  # Thrift port
+      - "4040:4040"    # Spark UI
+    environment:
+      - SPARK_NO_DAEMONIZE=true
+    healthcheck:
+      test: ["CMD-SHELL", "bash -c '(echo > /dev/tcp/localhost/10000) 2>/dev/null || exit 1'"]
+      interval: 10s
+      timeout: 5s
+      retries: 15
+      start_period: 60s
+    volumes:
+      - spark_data:/opt/spark/work-dir
+
+volumes:
+  spark_data: