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

pycharter/etl_generator/loaders/mongodb.py

@@ -0,0 +1,416 @@
+"""
+MongoDB loader for ETL orchestrator.
+
+Supports loading data to MongoDB collections using:
+- insert: Insert new documents (fails on duplicates)
+- upsert: Insert or update based on primary key
+- replace: Replace entire documents
+- update: Update specific fields in existing documents
+- delete: Delete documents by primary key
+
+Configuration example:
+    ```yaml
+    type: mongodb
+    mongodb:
+      url: mongodb://user:pass@localhost:27017
+      database: mydb
+      collection: users
+      ssh_tunnel:
+        enabled: false
+        host: bastion.example.com
+        username: tunnel_user
+        remote_host: mongo-internal.example.com
+        remote_port: 27017
+    
+    write_method: upsert  # insert, upsert, replace, update, delete
+    primary_key: _id      # field to use for matching (default: _id)
+    
+    # Optional: ordered bulk writes (default: true)
+    # Set to false for better performance when order doesn't matter
+    ordered: false
+    
+    # Optional: bypass document validation
+    bypass_validation: false
+    ```
+"""
+
+import logging
+import time
+from typing import Any, Dict, List, Optional, Union
+
+from bson import ObjectId
+from pymongo import MongoClient
+from pymongo.operations import (
+    DeleteOne,
+    InsertOne,
+    ReplaceOne,
+    UpdateOne,
+)
+
+from pycharter.etl_generator.loaders.base import BaseLoader
+from pycharter.etl_generator.result import LoadResult
+from pycharter.etl_generator.database import (
+    create_ssh_tunnel,
+    DEFAULT_TUNNEL_LOCAL_PORT,
+)
+from pycharter.utils.value_injector import resolve_values
+
+logger = logging.getLogger(__name__)
+
+# Default MongoDB port
+DEFAULT_MONGODB_PORT = 27017
+
+
+def _modify_mongodb_url_for_tunnel(url: str, local_port: int) -> str:
+    """
+    Modify MongoDB URL to use local tunnel port.
+    
+    Args:
+        url: Original MongoDB URL
+        local_port: Local tunnel port
+        
+    Returns:
+        Modified URL pointing to localhost tunnel
+    """
+    import re
+    pattern = r'mongodb(\+srv)?://([^@]+@)?([^/:]+)(:\d+)?'
+    
+    def replace_host(match):
+        protocol = match.group(1) or ''
+        auth = match.group(2) or ''
+        return f'mongodb://{auth}127.0.0.1:{local_port}'
+    
+    return re.sub(pattern, replace_host, url)
+
+
+def _convert_id_field(doc: Dict[str, Any], primary_key: str) -> Dict[str, Any]:
+    """
+    Convert primary key to ObjectId if it's a valid ObjectId string.
+    
+    Args:
+        doc: Document dictionary
+        primary_key: Primary key field name
+        
+    Returns:
+        Document with converted _id if applicable
+    """
+    result = doc.copy()
+    
+    # If primary key is '_id' and value is a valid ObjectId string, convert it
+    if primary_key == '_id' and '_id' in result:
+        id_value = result['_id']
+        if isinstance(id_value, str) and ObjectId.is_valid(id_value):
+            result['_id'] = ObjectId(id_value)
+    
+    return result
+
+
+class MongoDBLoader(BaseLoader):
+    """
+    Loader for MongoDB data targets.
+    
+    Supports:
+    - Insert (insert_many for bulk inserts)
+    - Upsert (bulk_write with ReplaceOne operations)
+    - Replace (bulk_write with ReplaceOne)
+    - Update (bulk_write with UpdateOne)
+    - Delete (bulk_write with DeleteOne)
+    
+    Example (programmatic API):
+        >>> loader = MongoDBLoader(
+        ...     connection_string="mongodb://localhost:27017",
+        ...     database="mydb",
+        ...     collection="users",
+        ...     write_method="upsert",
+        ...     primary_key="_id",
+        ... )
+        >>> result = await loader.load(data)
+    
+    Example (config-driven):
+        >>> loader = MongoDBLoader.from_config(config)
+        >>> result = await loader.load(data)
+    """
+    
+    def __init__(
+        self,
+        connection_string: str,
+        database: str,
+        collection: str,
+        write_method: str = "upsert",
+        primary_key: Union[str, List[str]] = "_id",
+        batch_size: int = 1000,
+        ordered: bool = True,
+        bypass_validation: bool = False,
+        ssh_tunnel: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize MongoDB loader.
+        
+        Args:
+            connection_string: MongoDB connection URL
+            database: Database name
+            collection: Collection name
+            write_method: Write method (insert, upsert, replace, update, delete)
+            primary_key: Field(s) to use for matching documents
+            batch_size: Number of operations per bulk write
+            ordered: Whether bulk writes should be ordered
+            bypass_validation: Whether to bypass document validation
+            ssh_tunnel: SSH tunnel configuration
+        """
+        self.connection_string = connection_string
+        self.database = database
+        self.collection = collection
+        self.write_method = write_method
+        self.primary_key = primary_key
+        self.batch_size = batch_size
+        self.ordered = ordered
+        self.bypass_validation = bypass_validation
+        self.ssh_tunnel = ssh_tunnel
+    
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "MongoDBLoader":
+        """Create loader from configuration dict."""
+        mongo_config = config.get("mongodb", {})
+        return cls(
+            connection_string=mongo_config.get("url") or config.get("connection_string"),
+            database=mongo_config.get("database") or config.get("database"),
+            collection=mongo_config.get("collection") or config.get("collection"),
+            write_method=config.get("write_method", "upsert"),
+            primary_key=config.get("primary_key", "_id"),
+            batch_size=config.get("batch_size", 1000),
+            ordered=config.get("ordered", True),
+            bypass_validation=config.get("bypass_validation", False),
+            ssh_tunnel=mongo_config.get("ssh_tunnel"),
+        )
+    
+    def _build_filter(self, doc: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Build a filter document for matching.
+        
+        Args:
+            doc: Source document
+            
+        Returns:
+            Filter document using primary key field(s)
+        """
+        if isinstance(self.primary_key, list):
+            # Composite key
+            return {key: doc.get(key) for key in self.primary_key if key in doc}
+        else:
+            # Single key
+            key_value = doc.get(self.primary_key)
+            if key_value is None:
+                raise ValueError(
+                    f"Document missing primary key field '{self.primary_key}': {doc}"
+                )
+            return {self.primary_key: key_value}
+    
+    async def load(self, data: List[Dict[str, Any]], **params) -> LoadResult:
+        """
+        Load data to MongoDB collection.
+        
+        Args:
+            data: List of documents to load
+            **params: Additional parameters (config_context, contract_dir)
+            
+        Returns:
+            LoadResult with statistics
+        """
+        start_time = time.time()
+        
+        if not data:
+            return LoadResult(success=True, rows_loaded=0)
+        
+        if not self.connection_string:
+            raise ValueError("MongoDB connection string is required")
+        if not self.database:
+            raise ValueError("MongoDB database name is required")
+        if not self.collection:
+            raise ValueError("MongoDB collection name is required")
+        
+        # Handle SSH tunnel if configured
+        tunnel = None
+        connection_string = self.connection_string
+        
+        config_context = params.get('config_context')
+        contract_dir = params.get('contract_dir')
+        source_file = str(contract_dir / "load.yaml") if contract_dir else None
+        
+        if self.ssh_tunnel:
+            ssh_config = resolve_values(
+                self.ssh_tunnel, 
+                context=config_context, 
+                source_file=source_file
+            )
+            enabled_value = ssh_config.get('enabled', False)
+            if isinstance(enabled_value, str):
+                enabled_lower = enabled_value.lower()
+                ssh_config['enabled'] = enabled_lower in ('true', '1', 'yes', 'on')
+            elif not isinstance(enabled_value, bool):
+                ssh_config['enabled'] = bool(enabled_value)
+            
+            if ssh_config.get('enabled', False):
+                if 'remote_port' not in ssh_config:
+                    ssh_config['remote_port'] = DEFAULT_MONGODB_PORT
+                
+                tunnel = create_ssh_tunnel(ssh_config)
+                if tunnel:
+                    local_port = int(ssh_config.get('local_port', DEFAULT_TUNNEL_LOCAL_PORT))
+                    connection_string = _modify_mongodb_url_for_tunnel(
+                        connection_string, local_port
+                    )
+        
+        # Create MongoDB client
+        client = MongoClient(connection_string)
+        db = client[self.database]
+        collection = db[self.collection]
+        
+        inserted = 0
+        updated = 0
+        deleted = 0
+        errors = []
+        
+        try:
+            logger.info(
+                f"Loading to MongoDB: {self.database}.{self.collection} "
+                f"(method: {self.write_method}, records: {len(data)})"
+            )
+            
+            # Process in batches
+            for i in range(0, len(data), self.batch_size):
+                batch = data[i:i + self.batch_size]
+                
+                if self.write_method == 'insert':
+                    # Simple insert_many
+                    result = collection.insert_many(
+                        batch,
+                        ordered=self.ordered,
+                        bypass_document_validation=self.bypass_validation,
+                    )
+                    inserted += len(result.inserted_ids)
+                
+                elif self.write_method == 'upsert':
+                    # Bulk upsert using ReplaceOne with upsert=True
+                    operations = []
+                    for doc in batch:
+                        doc = _convert_id_field(doc, self.primary_key if isinstance(self.primary_key, str) else '_id')
+                        filter_doc = self._build_filter(doc)
+                        operations.append(
+                            ReplaceOne(filter_doc, doc, upsert=True)
+                        )
+                    
+                    if operations:
+                        result = collection.bulk_write(
+                            operations,
+                            ordered=self.ordered,
+                            bypass_document_validation=self.bypass_validation,
+                        )
+                        inserted += result.upserted_count
+                        updated += result.modified_count
+                
+                elif self.write_method == 'replace':
+                    # Bulk replace (no upsert - only updates existing)
+                    operations = []
+                    for doc in batch:
+                        doc = _convert_id_field(doc, self.primary_key if isinstance(self.primary_key, str) else '_id')
+                        filter_doc = self._build_filter(doc)
+                        operations.append(
+                            ReplaceOne(filter_doc, doc, upsert=False)
+                        )
+                    
+                    if operations:
+                        result = collection.bulk_write(
+                            operations,
+                            ordered=self.ordered,
+                            bypass_document_validation=self.bypass_validation,
+                        )
+                        updated += result.modified_count
+                
+                elif self.write_method == 'update':
+                    # Bulk update using $set
+                    operations = []
+                    for doc in batch:
+                        doc = _convert_id_field(doc, self.primary_key if isinstance(self.primary_key, str) else '_id')
+                        filter_doc = self._build_filter(doc)
+                        # Remove primary key from update fields
+                        update_fields = {
+                            k: v for k, v in doc.items() 
+                            if k != self.primary_key and k not in (
+                                self.primary_key if isinstance(self.primary_key, list) else [self.primary_key]
+                            )
+                        }
+                        operations.append(
+                            UpdateOne(filter_doc, {"$set": update_fields})
+                        )
+                    
+                    if operations:
+                        result = collection.bulk_write(
+                            operations,
+                            ordered=self.ordered,
+                            bypass_document_validation=self.bypass_validation,
+                        )
+                        updated += result.modified_count
+                
+                elif self.write_method == 'delete':
+                    # Bulk delete by primary key
+                    operations = []
+                    for doc in batch:
+                        doc = _convert_id_field(doc, self.primary_key if isinstance(self.primary_key, str) else '_id')
+                        filter_doc = self._build_filter(doc)
+                        operations.append(DeleteOne(filter_doc))
+                    
+                    if operations:
+                        result = collection.bulk_write(
+                            operations,
+                            ordered=self.ordered,
+                        )
+                        deleted += result.deleted_count
+                
+                elif self.write_method == 'truncate_and_load':
+                    # Delete all documents first (only on first batch)
+                    if i == 0:
+                        collection.delete_many({})
+                        logger.info(f"Truncated collection {self.database}.{self.collection}")
+                    
+                    # Then insert
+                    result = collection.insert_many(
+                        batch,
+                        ordered=self.ordered,
+                        bypass_document_validation=self.bypass_validation,
+                    )
+                    inserted += len(result.inserted_ids)
+                
+                else:
+                    raise ValueError(f"Unsupported write method: {self.write_method}")
+            
+            duration = time.time() - start_time
+            total_affected = inserted + updated + deleted
+            
+            logger.info(
+                f"MongoDB load completed: {total_affected} records affected "
+                f"(inserted: {inserted}, updated: {updated}, deleted: {deleted}) "
+                f"in {duration:.2f}s"
+            )
+            
+            return LoadResult(
+                success=True,
+                rows_loaded=inserted + updated,
+                duration_seconds=duration,
+                metadata={
+                    "inserted": inserted,
+                    "updated": updated,
+                    "deleted": deleted,
+                },
+            )
+        
+        except Exception as e:
+            logger.error(f"MongoDB load failed: {e}", exc_info=True)
+            return LoadResult(
+                success=False,
+                error=str(e),
+                duration_seconds=time.time() - start_time,
+            )
+        finally:
+            client.close()
+            if tunnel:
+                tunnel.stop()