pycharter 0.0.22__py3-none-any.whl → 0.0.24__py3-none-any.whl

pycharter/etl_generator/loaders/database.py

@@ -0,0 +1,274 @@
+"""
+Database loaders for ETL pipelines.
+"""
+
+import logging
+import time
+from typing import Any, Dict, List, Optional, Union
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+
+from pycharter.etl_generator.database import (
+    detect_database_type,
+    create_ssh_tunnel,
+    modify_url_for_tunnel,
+    load_data_postgresql,
+    load_data_mysql,
+    load_data_sqlite,
+    load_data_mssql,
+    DEFAULT_TUNNEL_LOCAL_PORT,
+    DB_POSTGRESQL,
+    DB_MYSQL,
+    DB_SQLITE,
+    DB_MSSQL,
+)
+from pycharter.etl_generator.loaders.base import BaseLoader
+from pycharter.etl_generator.result import LoadResult
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresLoader(BaseLoader):
+    """
+    Loader for PostgreSQL databases.
+    
+    Supports:
+    - Insert, upsert, replace, update, delete, truncate_and_load
+    - Bulk operations for efficiency
+    - SSH tunneling
+    
+    Example:
+        >>> loader = PostgresLoader(
+        ...     connection_string="postgresql://user:pass@localhost/db",
+        ...     table="users",
+        ...     write_method="upsert",
+        ...     primary_key="id",
+        ... )
+        >>> result = await loader.load(data)
+    """
+    
+    def __init__(
+        self,
+        connection_string: str,
+        table: str,
+        schema: str = "public",
+        write_method: str = "upsert",
+        primary_key: Optional[Union[str, List[str]]] = None,
+        batch_size: int = 1000,
+        ssh_tunnel: Optional[Dict[str, Any]] = None,
+    ):
+        self.connection_string = connection_string
+        self.table = table
+        self.schema = schema
+        self.write_method = write_method
+        self.primary_key = primary_key
+        self.batch_size = batch_size
+        self.ssh_tunnel = ssh_tunnel
+    
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "PostgresLoader":
+        """Create loader from configuration dict."""
+        db_config = config.get("database", {})
+        return cls(
+            connection_string=db_config.get("url") or config.get("connection_string"),
+            table=db_config.get("table") or config.get("table"),
+            schema=db_config.get("schema", config.get("schema", "public")),
+            write_method=db_config.get("write_method", config.get("write_method", "upsert")),
+            primary_key=db_config.get("primary_key") or config.get("primary_key"),
+            batch_size=config.get("batch_size", 1000),
+            ssh_tunnel=db_config.get("ssh_tunnel"),
+        )
+    
+    async def load(self, data: List[Dict[str, Any]], **params) -> LoadResult:
+        """Load data to PostgreSQL."""
+        start_time = time.time()
+        
+        if not data:
+            return LoadResult(success=True, rows_loaded=0)
+        
+        # Handle SSH tunnel if configured
+        tunnel = None
+        connection_string = self.connection_string
+        
+        if self.ssh_tunnel and self.ssh_tunnel.get("enabled"):
+            tunnel = create_ssh_tunnel(self.ssh_tunnel)
+            if tunnel:
+                local_port = int(self.ssh_tunnel.get("local_port", DEFAULT_TUNNEL_LOCAL_PORT))
+                connection_string = modify_url_for_tunnel(
+                    connection_string, local_port, DB_POSTGRESQL
+                )
+        
+        try:
+            # Use async engine for PostgreSQL
+            # Convert sync URL to async if needed
+            if "+asyncpg" not in connection_string:
+                async_url = connection_string.replace("postgresql://", "postgresql+asyncpg://")
+            else:
+                async_url = connection_string
+            
+            engine = create_async_engine(async_url, echo=False)
+            async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+            
+            async with async_session() as session:
+                result = await load_data_postgresql(
+                    data=data,
+                    session=session,
+                    schema_name=self.schema,
+                    table_name=self.table,
+                    write_method=self.write_method,
+                    primary_key=self.primary_key,
+                    batch_size=self.batch_size,
+                )
+            
+            await engine.dispose()
+            
+            duration = time.time() - start_time
+            logger.info(f"Loaded {result['total']} records to {self.schema}.{self.table} in {duration:.2f}s")
+            
+            return LoadResult(
+                success=True,
+                rows_loaded=result.get("inserted", 0) + result.get("updated", 0),
+                duration_seconds=duration,
+            )
+        
+        except Exception as e:
+            logger.error(f"PostgreSQL load failed: {e}", exc_info=True)
+            return LoadResult(
+                success=False,
+                error=str(e),
+                duration_seconds=time.time() - start_time,
+            )
+        finally:
+            if tunnel:
+                tunnel.stop()
+
+
+class DatabaseLoader(BaseLoader):
+    """
+    Generic database loader that auto-detects database type.
+    
+    Supports PostgreSQL, MySQL, SQLite, and MSSQL.
+    
+    Example:
+        >>> loader = DatabaseLoader(
+        ...     connection_string="mysql://user:pass@localhost/db",
+        ...     table="users",
+        ... )
+        >>> result = await loader.load(data)
+    """
+    
+    def __init__(
+        self,
+        connection_string: str,
+        table: str,
+        schema: Optional[str] = None,
+        write_method: str = "upsert",
+        primary_key: Optional[Union[str, List[str]]] = None,
+        batch_size: int = 1000,
+        ssh_tunnel: Optional[Dict[str, Any]] = None,
+    ):
+        self.connection_string = connection_string
+        self.table = table
+        self.schema = schema
+        self.write_method = write_method
+        self.primary_key = primary_key
+        self.batch_size = batch_size
+        self.ssh_tunnel = ssh_tunnel
+        self.db_type = detect_database_type(connection_string)
+    
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "DatabaseLoader":
+        """Create loader from configuration dict."""
+        db_config = config.get("database", {})
+        return cls(
+            connection_string=db_config.get("url") or config.get("connection_string"),
+            table=db_config.get("table") or config.get("table"),
+            schema=db_config.get("schema") or config.get("schema"),
+            write_method=db_config.get("write_method", config.get("write_method", "upsert")),
+            primary_key=db_config.get("primary_key") or config.get("primary_key"),
+            batch_size=config.get("batch_size", 1000),
+            ssh_tunnel=db_config.get("ssh_tunnel"),
+        )
+    
+    async def load(self, data: List[Dict[str, Any]], **params) -> LoadResult:
+        """Load data using appropriate database loader."""
+        if self.db_type == DB_POSTGRESQL:
+            loader = PostgresLoader(
+                connection_string=self.connection_string,
+                table=self.table,
+                schema=self.schema or "public",
+                write_method=self.write_method,
+                primary_key=self.primary_key,
+                batch_size=self.batch_size,
+                ssh_tunnel=self.ssh_tunnel,
+            )
+            return await loader.load(data, **params)
+        else:
+            # For non-PostgreSQL databases, use sync loading
+            return await self._load_sync(data, **params)
+    
+    async def _load_sync(self, data: List[Dict[str, Any]], **params) -> LoadResult:
+        """Load data using sync database operations."""
+        start_time = time.time()
+        
+        if not data:
+            return LoadResult(success=True, rows_loaded=0)
+        
+        # Handle SSH tunnel if configured
+        tunnel = None
+        connection_string = self.connection_string
+        
+        if self.ssh_tunnel and self.ssh_tunnel.get("enabled"):
+            tunnel = create_ssh_tunnel(self.ssh_tunnel)
+            if tunnel:
+                local_port = int(self.ssh_tunnel.get("local_port", DEFAULT_TUNNEL_LOCAL_PORT))
+                connection_string = modify_url_for_tunnel(
+                    connection_string, local_port, self.db_type
+                )
+        
+        try:
+            engine = create_engine(connection_string, echo=False)
+            Session = sessionmaker(bind=engine)
+            session = Session()
+            
+            # Select appropriate load function
+            if self.db_type == DB_MYSQL:
+                result = load_data_mysql(
+                    data, session, self.schema or "", self.table,
+                    self.write_method, self.primary_key, self.batch_size
+                )
+            elif self.db_type == DB_SQLITE:
+                result = load_data_sqlite(
+                    data, session, "", self.table,
+                    self.write_method, self.primary_key, self.batch_size
+                )
+            elif self.db_type == DB_MSSQL:
+                result = load_data_mssql(
+                    data, session, self.schema or "dbo", self.table,
+                    self.write_method, self.primary_key, self.batch_size
+                )
+            else:
+                raise ValueError(f"Unsupported database type: {self.db_type}")
+            
+            session.close()
+            engine.dispose()
+            
+            duration = time.time() - start_time
+            return LoadResult(
+                success=True,
+                rows_loaded=result.get("inserted", 0) + result.get("updated", 0),
+                duration_seconds=duration,
+            )
+        
+        except Exception as e:
+            logger.error(f"Database load failed: {e}", exc_info=True)
+            return LoadResult(
+                success=False,
+                error=str(e),
+                duration_seconds=time.time() - start_time,
+            )
+        finally:
+            if tunnel:
+                tunnel.stop()

pycharter/etl_generator/loaders/factory.py

@@ -0,0 +1,180 @@
+"""
+Loader factory for ETL pipelines.
+
+Provides a registry pattern to select and instantiate the appropriate loader
+based on the target type specified in load configuration.
+
+Usage:
+    from pycharter.etl_generator.loaders.factory import LoaderFactory
+    
+    # Create loader from config
+    loader = LoaderFactory.create(load_config)
+    
+    # Register custom loader
+    LoaderFactory.register("bigquery", BigQueryLoader)
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Type
+
+from pycharter.etl_generator.loaders.base import BaseLoader
+from pycharter.etl_generator.loaders.database import PostgresLoader, DatabaseLoader
+from pycharter.etl_generator.loaders.file import FileLoader
+from pycharter.etl_generator.loaders.cloud import CloudStorageLoader
+
+logger = logging.getLogger(__name__)
+
+
+class LoaderFactory:
+    """
+    Factory for creating loader instances based on target type.
+    
+    Supports:
+    - Explicit 'type' field (recommended)
+    - Legacy 'target_type' field
+    - Auto-detection from config keys (for backward compatibility)
+    
+    Example:
+        # With explicit type (recommended)
+        config = {"type": "postgres", "table": "users", "database": {"url": "..."}}
+        loader = LoaderFactory.create(config)
+        
+        # Auto-detected (legacy)
+        config = {"table": "users", "connection_string": "postgresql://..."}
+        loader = LoaderFactory.create(config)  # Detected as postgres
+    """
+    
+    # Registry of loaders by target type
+    _registry: Dict[str, Type[BaseLoader]] = {
+        "postgres": PostgresLoader,
+        "postgresql": PostgresLoader,
+        "database": DatabaseLoader,
+        "sqlite": DatabaseLoader,
+        "file": FileLoader,
+        "cloud_storage": CloudStorageLoader,
+    }
+    
+    @classmethod
+    def register(cls, type_name: str, loader_class: Type[BaseLoader]) -> None:
+        """
+        Register a custom loader class.
+        
+        Args:
+            type_name: Type identifier (e.g., 'bigquery', 'snowflake')
+            loader_class: Loader class that inherits from BaseLoader
+            
+        Example:
+            class BigQueryLoader(BaseLoader):
+                ...
+            
+            LoaderFactory.register("bigquery", BigQueryLoader)
+        """
+        if not issubclass(loader_class, BaseLoader):
+            raise TypeError(f"Loader class must inherit from BaseLoader: {loader_class}")
+        cls._registry[type_name.lower()] = loader_class
+        logger.info(f"Registered loader: {type_name} -> {loader_class.__name__}")
+    
+    @classmethod
+    def unregister(cls, type_name: str) -> None:
+        """Remove a loader from the registry."""
+        cls._registry.pop(type_name.lower(), None)
+    
+    @classmethod
+    def list_types(cls) -> List[str]:
+        """List all registered loader types."""
+        return list(cls._registry.keys())
+    
+    @classmethod
+    def create(cls, config: Dict[str, Any]) -> BaseLoader:
+        """
+        Create a loader instance from configuration.
+        
+        Args:
+            config: Load configuration dictionary
+        
+        Returns:
+            Configured loader instance
+        
+        Raises:
+            ValueError: If type cannot be determined or is not registered
+        """
+        # Get type from config (check 'type' first, then 'target_type' for legacy)
+        load_type = config.get("type") or config.get("target_type")
+        
+        # Auto-detect if not specified
+        if not load_type:
+            load_type = cls._detect_type(config)
+            if load_type:
+                logger.debug(f"Auto-detected loader type: {load_type}")
+            else:
+                raise ValueError(
+                    "Cannot determine loader type. "
+                    f"Add 'type' field with one of: {', '.join(set(cls._registry.keys()))}"
+                )
+        
+        load_type = load_type.lower()
+        
+        # Get loader class from registry
+        loader_class = cls._registry.get(load_type)
+        if not loader_class:
+            raise ValueError(
+                f"Unknown loader type: '{load_type}'. "
+                f"Available types: {', '.join(set(cls._registry.keys()))}. "
+                f"Register custom loaders with LoaderFactory.register()"
+            )
+        
+        # Create loader using from_config if available
+        if hasattr(loader_class, "from_config"):
+            loader = loader_class.from_config(config)
+        else:
+            loader = loader_class()
+        
+        logger.debug(f"Created {loader_class.__name__} for type: {load_type}")
+        return loader
+    
+    @classmethod
+    def _detect_type(cls, config: Dict[str, Any]) -> Optional[str]:
+        """
+        Auto-detect loader type from configuration keys.
+        
+        This is for backward compatibility. New configs should use explicit 'type'.
+        """
+        # Database indicators
+        if "table" in config:
+            if "connection_string" in config or "database" in config:
+                # Check if it's SQLite
+                conn_str = config.get("connection_string", "")
+                if not conn_str and "database" in config:
+                    conn_str = config["database"].get("url", "")
+                if "sqlite" in conn_str.lower():
+                    return "sqlite"
+                return "postgres"
+        
+        # File indicators
+        if any(key in config for key in ("path", "file_path")) and "storage" not in config:
+            return "file"
+        
+        # Cloud storage indicators
+        if any(key in config for key in ("storage", "bucket", "container")):
+            return "cloud_storage"
+        
+        return None
+    
+    # Legacy method name for consistency with ExtractorFactory
+    @classmethod
+    def get_loader(cls, load_config: Dict[str, Any]) -> BaseLoader:
+        """Legacy method. Use create() instead."""
+        return cls.create(load_config)
+
+
+def get_loader(load_config: Dict[str, Any]) -> BaseLoader:
+    """
+    Convenience function to get loader instance.
+    
+    Args:
+        load_config: Load configuration dictionary
+    
+    Returns:
+        Loader instance
+    """
+    return LoaderFactory.create(load_config)

pycharter/etl_generator/loaders/file.py

@@ -0,0 +1,72 @@
+"""
+File loader for ETL pipelines.
+"""
+
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from pycharter.etl_generator.loaders.base import BaseLoader
+from pycharter.etl_generator.loaders.file_loader import load_to_file
+from pycharter.etl_generator.result import LoadResult
+
+
+class FileLoader(BaseLoader):
+    """
+    Loader for local files.
+    
+    Supports JSON, CSV, Parquet, and JSONL formats.
+    
+    Example:
+        >>> loader = FileLoader(path="output/data.json", format="json")
+        >>> result = await loader.load(data)
+    """
+    
+    def __init__(
+        self,
+        path: str,
+        file_format: str = "json",
+        write_mode: str = "overwrite",
+    ):
+        self.path = path
+        self.file_format = file_format
+        self.write_mode = write_mode
+    
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "FileLoader":
+        """Create loader from configuration dict."""
+        return cls(
+            path=config.get("file_path") or config.get("path"),
+            file_format=config.get("format", "json"),
+            write_mode=config.get("write_mode", "overwrite"),
+        )
+    
+    async def load(self, data: List[Dict[str, Any]], **params) -> LoadResult:
+        """Load data to file."""
+        start_time = time.time()
+        
+        if not data:
+            return LoadResult(success=True, rows_loaded=0)
+        
+        try:
+            load_config = {
+                "file_path": self.path,
+                "format": self.file_format,
+                "write_mode": self.write_mode,
+            }
+            
+            result = load_to_file(data, load_config)
+            
+            duration = time.time() - start_time
+            return LoadResult(
+                success=True,
+                rows_loaded=result.get("written", 0),
+                duration_seconds=duration,
+            )
+        
+        except Exception as e:
+            return LoadResult(
+                success=False,
+                error=str(e),
+                duration_seconds=time.time() - start_time,
+            )

pycharter/etl_generator/loaders/file_loader.py

@@ -0,0 +1,130 @@
+"""
+File-based loader for ETL orchestrator.
+
+Writes transformed data to local files in JSON, CSV, Parquet, or JSONL format.
+"""
+
+import io
+import json
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from pycharter.utils.value_injector import resolve_values
+
+logger = logging.getLogger(__name__)
+
+SUPPORTED_FORMATS = ("json", "csv", "parquet", "jsonl")
+
+
+def load_to_file(
+    data: List[Dict[str, Any]],
+    load_config: Dict[str, Any],
+    contract_dir: Optional[Any] = None,
+    config_context: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Write transformed data to a local file.
+
+    Load config (destination_type: file):
+      file_path: Path to output file (required). Supports ${VAR} resolution.
+      format: json | csv | parquet | jsonl (default: json)
+      write_mode: overwrite | append (default: overwrite).
+        append: for jsonl/csv, appends lines; for json, read-merge-write (array concat).
+
+    Returns:
+        Dict with keys: written, total, path, format
+    """
+    source_file = str(contract_dir / "load.yaml") if contract_dir else None
+    file_path = load_config.get("file_path")
+    if not file_path:
+        raise ValueError(
+            "File loader requires 'file_path' in load configuration. "
+            "Example: file_path: ./output/data.json"
+        )
+    file_path = resolve_values(
+        file_path, context=config_context, source_file=source_file
+    )
+    path = Path(file_path)
+
+    fmt = (load_config.get("format") or "json").lower()
+    if fmt not in SUPPORTED_FORMATS:
+        raise ValueError(
+            f"File loader format must be one of {SUPPORTED_FORMATS}, got '{fmt}'"
+        )
+    write_mode = (load_config.get("write_mode") or "overwrite").lower()
+    if write_mode not in ("overwrite", "append"):
+        raise ValueError(
+            "File loader write_mode must be 'overwrite' or 'append', "
+            f"got '{write_mode}'"
+        )
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if fmt == "json":
+        _write_json(data, path, write_mode)
+    elif fmt == "jsonl":
+        _write_jsonl(data, path, write_mode)
+    elif fmt == "csv":
+        _write_csv(data, path, write_mode)
+    elif fmt == "parquet":
+        _write_parquet(data, path, write_mode)
+
+    logger.info(f"File loader wrote {len(data)} records to {path} ({fmt})")
+    return {"written": len(data), "total": len(data), "path": str(path), "format": fmt}
+
+
+def _write_json(
+    data: List[Dict[str, Any]], path: Path, write_mode: str
+) -> None:
+    if write_mode == "append" and path.exists():
+        with open(path, "r", encoding="utf-8") as f:
+            existing = json.load(f)
+        if isinstance(existing, list):
+            data = existing + data
+        else:
+            data = [existing] + data
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, default=str)
+
+
+def _write_jsonl(
+    data: List[Dict[str, Any]], path: Path, write_mode: str
+) -> None:
+    mode = "a" if write_mode == "append" and path.exists() else "w"
+    with open(path, mode, encoding="utf-8") as f:
+        for record in data:
+            f.write(json.dumps(record, default=str) + "\n")
+
+
+def _write_csv(
+    data: List[Dict[str, Any]], path: Path, write_mode: str
+) -> None:
+    if not data:
+        return
+    import csv
+
+    mode = "a" if write_mode == "append" and path.exists() else "w"
+    newfile = mode == "w"
+    with open(path, mode, encoding="utf-8", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=data[0].keys())
+        if newfile:
+            writer.writeheader()
+        writer.writerows(data)
+
+
+def _write_parquet(
+    data: List[Dict[str, Any]], path: Path, write_mode: str
+) -> None:
+    try:
+        import pandas as pd
+    except ImportError as e:
+        raise ImportError(
+            "pandas is required for Parquet file load. "
+            "Install with: pip install pandas pyarrow"
+        ) from e
+    df = pd.DataFrame(data)
+    if write_mode == "append" and path.exists():
+        existing = pd.read_parquet(path)
+        df = pd.concat([existing, df], ignore_index=True)
+    df.to_parquet(path, index=False)