diff --git a/semantica/ingest/api_ingestor.py b/semantica/ingest/api_ingestor.py new file mode 100644 index 00000000..226325cb --- /dev/null +++ b/semantica/ingest/api_ingestor.py @@ -0,0 +1,372 @@ +""" +REST API Ingestion Module + +This module provides comprehensive REST API ingestion capabilities for the +Semantica framework, enabling data extraction from any REST API endpoint. + +Key Features: + - Generic REST API client + - Pagination support + - Authentication (API key, OAuth, Bearer token) + - Batch request handling + - Error handling and retry logic + +Main Classes: + - RESTIngestor: Main REST API ingestion class + - APIData: Data representation for API ingestion + +Example Usage: + >>> from semantica.ingest import RESTIngestor + >>> ingestor = RESTIngestor() + >>> data = ingestor.ingest_endpoint("https://api.example.com/data", headers={"Authorization": "Bearer token"}) + >>> paginated_data = ingestor.paginated_fetch("https://api.example.com/data", page_size=100) +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Optional, Union +from urllib.parse import urljoin, urlparse + +import requests +from requests.adapters import HTTPAdapter +try: + from urllib3.util.retry import Retry +except ImportError: + from requests.packages.urllib3.util.retry import Retry + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + + +@dataclass +class APIData: + """REST API data representation.""" + + data: Union[List[Dict[str, Any]], Dict[str, Any]] + response_status: int + endpoint: str + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class RESTIngestor: + """ + REST API ingestion handler. + + This class provides comprehensive REST API ingestion capabilities, + supporting generic API endpoints with authentication, pagination, and + error handling. + + Features: + - Generic REST API client + - Pagination support + - Authentication (API key, OAuth, Bearer token) + - Batch request handling + - Error handling and retry logic + + Example Usage: + >>> ingestor = RESTIngestor() + >>> data = ingestor.ingest_endpoint("https://api.example.com/data") + >>> paginated_data = ingestor.paginated_fetch("https://api.example.com/data", page_size=100) + """ + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize REST API ingestor. + + Args: + config: Optional REST API ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + self.logger = get_logger("api_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize session with retry strategy + self.session = requests.Session() + retry_strategy = Retry( + total=self.config.get("max_retries", 3), + backoff_factor=self.config.get("backoff_factor", 1), + status_forcelist=[429, 500, 502, 503, 504], + ) + adapter = HTTPAdapter(max_retries=retry_strategy) + self.session.mount("http://", adapter) + self.session.mount("https://", adapter) + + # Set default headers + default_headers = self.config.get("headers", {}) + if default_headers: + self.session.headers.update(default_headers) + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("REST API ingestor initialized") + + def ingest_endpoint( + self, + endpoint: str, + method: str = "GET", + headers: Optional[Dict[str, str]] = None, + params: Optional[Dict[str, Any]] = None, + data: Optional[Union[Dict, str]] = None, + json_data: Optional[Dict[str, Any]] = None, + **options, + ) -> APIData: + """ + Ingest data from REST API endpoint. + + This method makes a request to a REST API endpoint and returns the response data. + + Args: + endpoint: API endpoint URL + method: HTTP method (GET, POST, PUT, DELETE, etc.) + headers: Optional request headers + params: Optional query parameters + data: Optional request body (for form data) + json_data: Optional JSON request body + **options: Additional request options + + Returns: + APIData: Ingested data object containing: + - data: Response data (parsed JSON or raw text) + - response_status: HTTP status code + - endpoint: Endpoint URL + - metadata: Additional metadata + + Raises: + ValidationError: If endpoint is invalid + ProcessingError: If request fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=endpoint, + module="ingest", + submodule="RESTIngestor", + message=f"Requesting: {method} {endpoint}", + ) + + try: + # Prepare request + request_headers = self.session.headers.copy() + if headers: + request_headers.update(headers) + + # Make request + response = self.session.request( + method=method, + url=endpoint, + headers=request_headers, + params=params, + data=data, + json=json_data, + timeout=self.config.get("timeout", 30), + **options, + ) + + # Check for errors + response.raise_for_status() + + # Parse response + try: + response_data = response.json() + except ValueError: + # Not JSON, return as text + response_data = response.text + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Request successful: {response.status_code}", + ) + + self.logger.info( + f"API request completed: {method} {endpoint} - {response.status_code}" + ) + + return APIData( + data=response_data, + response_status=response.status_code, + endpoint=endpoint, + metadata={ + "method": method, + "headers": dict(response.headers), + "content_type": response.headers.get("Content-Type"), + }, + ) + + except requests.exceptions.RequestException as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest endpoint {endpoint}: {e}") + raise ProcessingError(f"Failed to ingest endpoint: {e}") from e + + def paginated_fetch( + self, + endpoint: str, + page_size: int = 100, + page_param: str = "page", + size_param: str = "size", + limit: Optional[int] = None, + **options, + ) -> List[APIData]: + """ + Fetch paginated data from REST API endpoint. + + This method handles pagination automatically, fetching all pages or up to + a specified limit. + + Args: + endpoint: API endpoint URL + page_size: Number of items per page + page_param: Query parameter name for page number + size_param: Query parameter name for page size + limit: Maximum number of items to fetch (optional) + **options: Additional request options + + Returns: + List of APIData objects, one per page + + Raises: + ProcessingError: If pagination fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=endpoint, + module="ingest", + submodule="RESTIngestor", + message=f"Fetching paginated data from: {endpoint}", + ) + + try: + all_pages = [] + page = 1 + total_fetched = 0 + + while True: + # Build query parameters + params = options.get("params", {}).copy() + params[page_param] = page + params[size_param] = page_size + + # Fetch page + page_options = options.copy() + page_options["params"] = params + + page_data = self.ingest_endpoint(endpoint, **page_options) + + # Extract items from response + if isinstance(page_data.data, list): + items = page_data.data + elif isinstance(page_data.data, dict): + # Try common pagination response formats + items = ( + page_data.data.get("items", []) + or page_data.data.get("data", []) + or page_data.data.get("results", []) + or [page_data.data] + ) + else: + items = [] + + if not items: + # No more items, stop pagination + break + + all_pages.append(page_data) + total_fetched += len(items) + + # Check limit + if limit and total_fetched >= limit: + break + + # Check if there are more pages + if isinstance(page_data.data, dict): + has_more = ( + page_data.data.get("has_more", False) + or page_data.data.get("next", None) is not None + ) + if not has_more: + break + + page += 1 + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Fetched {len(all_pages)} pages, {total_fetched} items", + ) + + self.logger.info( + f"Paginated fetch completed: {len(all_pages)} page(s), {total_fetched} item(s)" + ) + + return all_pages + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to fetch paginated data: {e}") + raise ProcessingError(f"Failed to fetch paginated data: {e}") from e + + def batch_request( + self, + endpoints: List[str], + method: str = "GET", + **options, + ) -> List[APIData]: + """ + Make batch requests to multiple endpoints. + + This method makes requests to multiple endpoints and returns all results. + + Args: + endpoints: List of endpoint URLs + method: HTTP method + **options: Additional request options + + Returns: + List of APIData objects, one per endpoint + + Raises: + ProcessingError: If batch request fails + """ + tracking_id = self.progress_tracker.start_tracking( + file="batch", + module="ingest", + submodule="RESTIngestor", + message=f"Batch request: {len(endpoints)} endpoints", + ) + + try: + results = [] + for endpoint in endpoints: + try: + data = self.ingest_endpoint(endpoint, method=method, **options) + results.append(data) + except Exception as e: + self.logger.warning(f"Failed to fetch {endpoint}: {e}") + if self.config.get("fail_fast", False): + raise + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Batch request completed: {len(results)}/{len(endpoints)} successful", + ) + + self.logger.info( + f"Batch request completed: {len(results)}/{len(endpoints)} successful" + ) + + return results + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to execute batch request: {e}") + raise ProcessingError(f"Failed to execute batch request: {e}") from e + diff --git a/semantica/ingest/duckdb_ingestor.py b/semantica/ingest/duckdb_ingestor.py new file mode 100644 index 00000000..fc4036ed --- /dev/null +++ b/semantica/ingest/duckdb_ingestor.py @@ -0,0 +1,432 @@ +""" +DuckDB Ingestion Module + +This module provides comprehensive DuckDB ingestion capabilities for the +Semantica framework, enabling data extraction from CSV, Parquet, and Excel files +using DuckDB's SQL interface. + +Key Features: + - CSV file ingestion with SQL queries + - Parquet file ingestion + - Excel file ingestion + - SQL query execution on files + - Schema extraction + - Large dataset handling + +Main Classes: + - DuckDBIngestor: Main DuckDB ingestion class + - DuckDBData: Data representation for DuckDB ingestion + +Example Usage: + >>> from semantica.ingest import DuckDBIngestor + >>> ingestor = DuckDBIngestor() + >>> data = ingestor.ingest_csv("data.csv") + >>> parquet_data = ingestor.ingest_parquet("data.parquet") + >>> excel_data = ingestor.ingest_excel("data.xlsx") +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +try: + import duckdb +except ImportError: + duckdb = None + + +@dataclass +class DuckDBData: + """DuckDB data representation.""" + + data: List[Dict[str, Any]] + row_count: int + columns: List[str] + query: Optional[str] = None + source_file: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class DuckDBIngestor: + """ + DuckDB ingestion handler. + + This class provides comprehensive DuckDB ingestion capabilities, + supporting ingestion from CSV, Parquet, and Excel files using DuckDB's + SQL interface for efficient querying. + + Features: + - CSV file ingestion + - Parquet file ingestion + - Excel file ingestion + - SQL query execution on files + - Schema extraction + - Large dataset handling + + Example Usage: + >>> ingestor = DuckDBIngestor() + >>> data = ingestor.ingest_csv("data.csv") + >>> result = ingestor.execute_query("SELECT * FROM 'data.csv' LIMIT 100") + """ + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize DuckDB ingestor. + + Args: + config: Optional DuckDB ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + if duckdb is None: + raise ImportError( + "duckdb is required for DuckDBIngestor. Install it with: pip install duckdb" + ) + + self.logger = get_logger("duckdb_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize DuckDB connection + self.conn = duckdb.connect() + if self.config.get("memory_limit"): + self.conn.execute(f"SET memory_limit='{self.config['memory_limit']}'") + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("DuckDB ingestor initialized") + + def __del__(self): + """Close DuckDB connection on cleanup.""" + if hasattr(self, "conn") and self.conn: + try: + self.conn.close() + except Exception: + pass + + def ingest_csv( + self, + file_path: Union[str, Path], + limit: Optional[int] = None, + where: Optional[str] = None, + **options, + ) -> DuckDBData: + """ + Ingest data from CSV file. + + This method reads a CSV file using DuckDB and returns the data. + + Args: + file_path: Path to CSV file + limit: Maximum number of rows to return (optional) + where: WHERE clause for filtering (optional) + **options: Additional query options + + Returns: + DuckDBData: Ingested data object + + Raises: + ValidationError: If CSV file not found + ProcessingError: If CSV reading fails + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise ValidationError(f"CSV file not found: {file_path}") + + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="ingest", + submodule="DuckDBIngestor", + message=f"Ingesting CSV: {file_path.name}", + ) + + try: + # Build SQL query + query = f"SELECT * FROM read_csv_auto('{file_path}')" + + if where: + query += f" WHERE {where}" + + if limit: + query += f" LIMIT {limit}" + + # Execute query + result = self.conn.execute(query).fetchall() + columns = [desc[0] for desc in self.conn.description] + + # Convert to list of dictionaries + data = [dict(zip(columns, row)) for row in result] + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested CSV: {len(data)} rows", + ) + + self.logger.info(f"CSV ingestion completed: {len(data)} row(s)") + + return DuckDBData( + data=data, + row_count=len(data), + columns=columns, + query=query, + source_file=str(file_path), + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest CSV {file_path}: {e}") + raise ProcessingError(f"Failed to ingest CSV: {e}") from e + + def ingest_parquet( + self, + file_path: Union[str, Path], + limit: Optional[int] = None, + where: Optional[str] = None, + **options, + ) -> DuckDBData: + """ + Ingest data from Parquet file. + + This method reads a Parquet file using DuckDB and returns the data. + + Args: + file_path: Path to Parquet file + limit: Maximum number of rows to return (optional) + where: WHERE clause for filtering (optional) + **options: Additional query options + + Returns: + DuckDBData: Ingested data object + + Raises: + ValidationError: If Parquet file not found + ProcessingError: If Parquet reading fails + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise ValidationError(f"Parquet file not found: {file_path}") + + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="ingest", + submodule="DuckDBIngestor", + message=f"Ingesting Parquet: {file_path.name}", + ) + + try: + # Build SQL query + query = f"SELECT * FROM read_parquet('{file_path}')" + + if where: + query += f" WHERE {where}" + + if limit: + query += f" LIMIT {limit}" + + # Execute query + result = self.conn.execute(query).fetchall() + columns = [desc[0] for desc in self.conn.description] + + # Convert to list of dictionaries + data = [dict(zip(columns, row)) for row in result] + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested Parquet: {len(data)} rows", + ) + + self.logger.info(f"Parquet ingestion completed: {len(data)} row(s)") + + return DuckDBData( + data=data, + row_count=len(data), + columns=columns, + query=query, + source_file=str(file_path), + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest Parquet {file_path}: {e}") + raise ProcessingError(f"Failed to ingest Parquet: {e}") from e + + def ingest_excel( + self, + file_path: Union[str, Path], + sheet_name: Optional[str] = None, + limit: Optional[int] = None, + where: Optional[str] = None, + **options, + ) -> DuckDBData: + """ + Ingest data from Excel file. + + This method reads an Excel file using DuckDB (via Parquet conversion) + and returns the data. + + Args: + file_path: Path to Excel file + sheet_name: Name of sheet to read (optional, reads first sheet if not specified) + limit: Maximum number of rows to return (optional) + where: WHERE clause for filtering (optional) + **options: Additional query options + + Returns: + DuckDBData: Ingested data object + + Raises: + ValidationError: If Excel file not found + ProcessingError: If Excel reading fails + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise ValidationError(f"Excel file not found: {file_path}") + + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="ingest", + submodule="DuckDBIngestor", + message=f"Ingesting Excel: {file_path.name}", + ) + + try: + # DuckDB doesn't directly support Excel, so we need to use pandas + # to read Excel and then query it + try: + import pandas as pd + except ImportError: + raise ImportError( + "pandas and openpyxl are required for Excel ingestion. " + "Install with: pip install pandas openpyxl" + ) + + # Read Excel with pandas + if sheet_name: + df = pd.read_excel(file_path, sheet_name=sheet_name, **options) + else: + df = pd.read_excel(file_path, **options) + + # Register DataFrame as a DuckDB table + table_name = f"excel_data_{id(df)}" + self.conn.register(table_name, df) + + # Build SQL query + query = f"SELECT * FROM {table_name}" + + if where: + query += f" WHERE {where}" + + if limit: + query += f" LIMIT {limit}" + + # Execute query + result = self.conn.execute(query).fetchall() + columns = [desc[0] for desc in self.conn.description] + + # Convert to list of dictionaries + data = [dict(zip(columns, row)) for row in result] + + # Unregister table + self.conn.unregister(table_name) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested Excel: {len(data)} rows", + ) + + self.logger.info(f"Excel ingestion completed: {len(data)} row(s)") + + return DuckDBData( + data=data, + row_count=len(data), + columns=columns, + query=query, + source_file=str(file_path), + metadata={"sheet_name": sheet_name}, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest Excel {file_path}: {e}") + raise ProcessingError(f"Failed to ingest Excel: {e}") from e + + def execute_query( + self, + query: str, + file_path: Optional[Union[str, Path]] = None, + **options, + ) -> DuckDBData: + """ + Execute SQL query on file or in-memory data. + + This method executes a SQL query using DuckDB. If file_path is provided, + the query can reference the file directly. + + Args: + query: SQL query to execute + file_path: Optional file path to reference in query + **options: Additional query options + + Returns: + DuckDBData: Query result data object + + Raises: + ProcessingError: If query execution fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=file_path or "query", + module="ingest", + submodule="DuckDBIngestor", + message="Executing SQL query", + ) + + try: + # Execute query + result = self.conn.execute(query).fetchall() + columns = [desc[0] for desc in self.conn.description] + + # Convert to list of dictionaries + data = [dict(zip(columns, row)) for row in result] + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Query executed: {len(data)} rows", + ) + + self.logger.info(f"Query executed: {len(data)} row(s) returned") + + return DuckDBData( + data=data, + row_count=len(data), + columns=columns, + query=query, + source_file=str(file_path) if file_path else None, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to execute query: {e}") + raise ProcessingError(f"Failed to execute query: {e}") from e + diff --git a/semantica/ingest/elastic_ingestor.py b/semantica/ingest/elastic_ingestor.py new file mode 100644 index 00000000..fc9a7499 --- /dev/null +++ b/semantica/ingest/elastic_ingestor.py @@ -0,0 +1,356 @@ +""" +Elasticsearch Ingestion Module + +This module provides comprehensive Elasticsearch ingestion capabilities for the +Semantica framework, enabling data extraction from Elasticsearch indices. + +Key Features: + - Index ingestion + - Search query execution + - Index export + - Schema extraction + - Large dataset handling with scroll API + +Main Classes: + - ElasticIngestor: Main Elasticsearch ingestion class + - ElasticData: Data representation for Elasticsearch ingestion + +Example Usage: + >>> from semantica.ingest import ElasticIngestor + >>> ingestor = ElasticIngestor() + >>> data = ingestor.ingest_index("http://localhost:9200", "my_index") + >>> results = ingestor.search_documents("http://localhost:9200", "my_index", {"query": {"match_all": {}}}) +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +try: + from elasticsearch import Elasticsearch + from elasticsearch.helpers import scan +except ImportError: + Elasticsearch = None + scan = None + + +@dataclass +class ElasticData: + """Elasticsearch data representation.""" + + documents: List[Dict[str, Any]] + document_count: int + index_name: str + schema: Dict[str, Any] + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class ElasticIngestor: + """ + Elasticsearch ingestion handler. + + This class provides comprehensive Elasticsearch ingestion capabilities, + connecting to Elasticsearch, querying indices, and exporting data. + + Features: + - Index ingestion + - Search query execution + - Index export + - Schema extraction + - Large dataset handling with scroll API + + Example Usage: + >>> ingestor = ElasticIngestor() + >>> data = ingestor.ingest_index("http://localhost:9200", "my_index") + >>> results = ingestor.search_documents("http://localhost:9200", "my_index", {"query": {"match_all": {}}}) + """ + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize Elasticsearch ingestor. + + Args: + config: Optional Elasticsearch ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + if Elasticsearch is None: + raise ImportError( + "elasticsearch is required for ElasticIngestor. Install it with: pip install elasticsearch" + ) + + self.logger = get_logger("elastic_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("Elasticsearch ingestor initialized") + + def _get_client(self, connection_string: str) -> Elasticsearch: + """ + Get Elasticsearch client. + + Args: + connection_string: Elasticsearch connection string or host + + Returns: + Elasticsearch: Elasticsearch client object + """ + # Parse connection string + if connection_string.startswith("http://") or connection_string.startswith( + "https://" + ): + hosts = [connection_string] + else: + # Assume it's a host:port format + hosts = [connection_string] + + # Create client + client_config = self.config.get("client_config", {}) + client = Elasticsearch(hosts=hosts, **client_config) + + # Test connection + if not client.ping(): + raise ProcessingError("Failed to connect to Elasticsearch") + + return client + + def ingest_index( + self, + connection_string: str, + index_name: str, + limit: Optional[int] = None, + **options, + ) -> ElasticData: + """ + Ingest data from Elasticsearch index. + + This method connects to Elasticsearch, retrieves documents from an index, + and extracts schema information. + + Args: + connection_string: Elasticsearch connection string (e.g., "http://localhost:9200") + index_name: Name of the index + limit: Maximum number of documents to retrieve (optional) + **options: Additional processing options + + Returns: + ElasticData: Ingested data object containing: + - documents: List of document dictionaries + - document_count: Number of documents + - index_name: Index name + - schema: Schema information dictionary + + Raises: + ProcessingError: If ingestion fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=index_name, + module="ingest", + submodule="ElasticIngestor", + message=f"Ingesting index: {index_name}", + ) + + try: + # Get client + client = self._get_client(connection_string) + + # Check if index exists + if not client.indices.exists(index=index_name): + raise ValidationError(f"Index not found: {index_name}") + + # Get total document count + count_response = client.count(index=index_name) + total_count = count_response["count"] + + # Retrieve documents using scroll API for large datasets + self.progress_tracker.update_tracking( + tracking_id, message=f"Retrieving documents (total: {total_count})..." + ) + + query = options.get("query", {"match_all": {}}) + scroll_size = options.get("scroll_size", 1000) + + documents = [] + if limit: + # Use regular search with size limit + response = client.search( + index=index_name, body={"query": query}, size=limit + ) + documents = [hit["_source"] for hit in response["hits"]["hits"]] + else: + # Use scroll API for all documents + for doc in scan( + client, query={"query": query}, index=index_name, size=scroll_size + ): + documents.append(doc["_source"]) + if limit and len(documents) >= limit: + break + + # Extract schema information + schema = self._extract_schema(documents) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested {len(documents)} documents", + ) + + self.logger.info( + f"Index ingestion completed: {len(documents)} document(s)" + ) + + return ElasticData( + documents=documents, + document_count=len(documents), + index_name=index_name, + schema=schema, + metadata={"total_count": total_count, "query": query}, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest index: {e}") + raise ProcessingError(f"Failed to ingest index: {e}") from e + + def search_documents( + self, + connection_string: str, + index_name: str, + search_query: Dict[str, Any], + limit: Optional[int] = None, + **options, + ) -> ElasticData: + """ + Search documents in Elasticsearch index. + + This method executes a search query on an Elasticsearch index and returns + matching documents. + + Args: + connection_string: Elasticsearch connection string + index_name: Name of the index + search_query: Elasticsearch search query dictionary + limit: Maximum number of documents to return (optional) + **options: Additional search options + + Returns: + ElasticData: Search result data object + + Raises: + ProcessingError: If search fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=index_name, + module="ingest", + submodule="ElasticIngestor", + message=f"Searching index: {index_name}", + ) + + try: + # Get client + client = self._get_client(connection_string) + + # Execute search + size = limit or options.get("size", 100) + response = client.search( + index=index_name, body=search_query, size=size, **options + ) + + # Extract documents + documents = [hit["_source"] for hit in response["hits"]["hits"]] + + # Extract schema information + schema = self._extract_schema(documents) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Search returned {len(documents)} documents", + ) + + self.logger.info(f"Search completed: {len(documents)} document(s)") + + return ElasticData( + documents=documents, + document_count=len(documents), + index_name=index_name, + schema=schema, + metadata={ + "total_hits": response["hits"]["total"]["value"], + "query": search_query, + }, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to search documents: {e}") + raise ProcessingError(f"Failed to search documents: {e}") from e + + def export_index( + self, + connection_string: str, + index_name: str, + **options, + ) -> ElasticData: + """ + Export entire index. + + This method exports all documents from an Elasticsearch index. + + Args: + connection_string: Elasticsearch connection string + index_name: Name of the index + **options: Additional export options + + Returns: + ElasticData: Exported data object + + Raises: + ProcessingError: If export fails + """ + return self.ingest_index(connection_string, index_name, limit=None, **options) + + def _extract_schema(self, documents: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Extract schema information from documents. + + Args: + documents: List of Elasticsearch documents + + Returns: + Dictionary containing schema information + """ + if not documents: + return {"fields": [], "field_types": {}} + + # Collect all field names and types + field_types = {} + for doc in documents: + for key, value in doc.items(): + if key not in field_types: + field_types[key] = set() + field_types[key].add(type(value).__name__) + + # Convert sets to lists + schema = { + "fields": list(field_types.keys()), + "field_types": { + k: list(v) if len(v) > 1 else list(v)[0] + for k, v in field_types.items() + }, + } + + return schema + diff --git a/semantica/ingest/gdrive_ingestor.py b/semantica/ingest/gdrive_ingestor.py new file mode 100644 index 00000000..7d750957 --- /dev/null +++ b/semantica/ingest/gdrive_ingestor.py @@ -0,0 +1,422 @@ +""" +Google Drive Ingestion Module + +This module provides comprehensive Google Drive ingestion capabilities for the +Semantica framework, enabling data extraction from Google Drive files and folders. + +Key Features: + - Folder ingestion + - File ingestion + - Drive export + - File type detection + - OAuth authentication + +Main Classes: + - GDriveIngestor: Main Google Drive ingestion class + - GDriveData: Data representation for Google Drive ingestion + +Example Usage: + >>> from semantica.ingest import GDriveIngestor + >>> ingestor = GDriveIngestor(credentials_path="credentials.json") + >>> data = ingestor.ingest_folder("folder_id") + >>> file_data = ingestor.ingest_file("file_id") +""" + +import json +import os +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +try: + from google.oauth2.credentials import Credentials + from google_auth_oauthlib.flow import InstalledAppFlow + from google.auth.transport.requests import Request + from googleapiclient.discovery import build + from googleapiclient.errors import HttpError + from googleapiclient.http import MediaIoBaseDownload + import io +except ImportError: + Credentials = None + InstalledAppFlow = None + Request = None + build = None + HttpError = None + MediaIoBaseDownload = None + io = None + + +@dataclass +class GDriveData: + """Google Drive data representation.""" + + files: List[Dict[str, Any]] + file_count: int + folder_id: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class GDriveIngestor: + """ + Google Drive ingestion handler. + + This class provides comprehensive Google Drive ingestion capabilities, + connecting to Google Drive, listing files, and downloading content. + + Features: + - Folder ingestion + - File ingestion + - Drive export + - File type detection + - OAuth authentication + + Example Usage: + >>> ingestor = GDriveIngestor(credentials_path="credentials.json") + >>> data = ingestor.ingest_folder("folder_id") + """ + + # Google Drive API scopes + SCOPES = ["https://www.googleapis.com/auth/drive.readonly"] + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize Google Drive ingestor. + + Args: + config: Optional Google Drive ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + if build is None: + raise ImportError( + "google-api-python-client and google-auth-oauthlib are required for GDriveIngestor. " + "Install with: pip install google-api-python-client google-auth-oauthlib" + ) + + self.logger = get_logger("gdrive_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize Google Drive service + self.service = None + self.credentials_path = self.config.get("credentials_path") + self.token_path = self.config.get("token_path", "token.json") + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("Google Drive ingestor initialized") + + def _authenticate(self): + """ + Authenticate with Google Drive API. + + Raises: + ProcessingError: If authentication fails + """ + if self.service: + return + + creds = None + + # Load existing token + if os.path.exists(self.token_path): + try: + creds = Credentials.from_authorized_user_file( + self.token_path, self.SCOPES + ) + except Exception as e: + self.logger.warning(f"Failed to load token: {e}") + + # If there are no (valid) credentials available, let the user log in + if not creds or not creds.valid: + if creds and creds.expired and creds.refresh_token: + creds.refresh(Request()) + else: + if not self.credentials_path: + raise ValidationError( + "credentials_path is required for Google Drive authentication. " + "Provide path to OAuth2 credentials JSON file." + ) + + if not os.path.exists(self.credentials_path): + raise ValidationError( + f"Credentials file not found: {self.credentials_path}" + ) + + flow = InstalledAppFlow.from_client_secrets_file( + self.credentials_path, self.SCOPES + ) + creds = flow.run_local_server(port=0) + + # Save the credentials for the next run + with open(self.token_path, "w") as token: + token.write(creds.to_json()) + + # Build service + self.service = build("drive", "v3", credentials=creds) + self.logger.info("Authenticated with Google Drive") + + def ingest_folder( + self, + folder_id: str, + include_subfolders: bool = False, + file_types: Optional[List[str]] = None, + **options, + ) -> GDriveData: + """ + Ingest data from Google Drive folder. + + This method lists all files in a Google Drive folder and retrieves + their metadata. + + Args: + folder_id: Google Drive folder ID + include_subfolders: Whether to include files from subfolders + file_types: Optional list of file MIME types to filter + **options: Additional processing options + + Returns: + GDriveData: Ingested data object containing: + - files: List of file metadata dictionaries + - file_count: Number of files + - folder_id: Folder ID + - metadata: Additional metadata + + Raises: + ProcessingError: If ingestion fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=folder_id, + module="ingest", + submodule="GDriveIngestor", + message=f"Ingesting folder: {folder_id}", + ) + + try: + # Authenticate + self._authenticate() + + # Query files in folder + query = f"'{folder_id}' in parents and trashed=false" + if file_types: + mime_types = " or ".join([f"mimeType='{ft}'" for ft in file_types]) + query += f" and ({mime_types})" + + self.progress_tracker.update_tracking( + tracking_id, message="Listing files in folder..." + ) + + files = [] + page_token = None + + while True: + results = ( + self.service.files() + .list( + q=query, + pageSize=1000, + fields="nextPageToken, files(id, name, mimeType, size, modifiedTime, createdTime)", + pageToken=page_token, + ) + .execute() + ) + + items = results.get("files", []) + files.extend(items) + + page_token = results.get("nextPageToken") + if not page_token: + break + + # Handle subfolders if requested + if include_subfolders: + # Get all subfolders + subfolder_query = f"'{folder_id}' in parents and mimeType='application/vnd.google-apps.folder' and trashed=false" + subfolders = [] + page_token = None + + while True: + results = ( + self.service.files() + .list( + q=subfolder_query, + pageSize=1000, + fields="nextPageToken, files(id, name)", + pageToken=page_token, + ) + .execute() + ) + + items = results.get("files", []) + subfolders.extend(items) + + page_token = results.get("nextPageToken") + if not page_token: + break + + # Recursively ingest subfolders + for subfolder in subfolders: + try: + subfolder_data = self.ingest_folder( + subfolder["id"], + include_subfolders=True, + file_types=file_types, + **options, + ) + files.extend(subfolder_data.files) + except Exception as e: + self.logger.warning( + f"Failed to ingest subfolder {subfolder['name']}: {e}" + ) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested {len(files)} files", + ) + + self.logger.info(f"Folder ingestion completed: {len(files)} file(s)") + + return GDriveData( + files=files, + file_count=len(files), + folder_id=folder_id, + metadata={"include_subfolders": include_subfolders}, + ) + + except HttpError as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest folder: {e}") + raise ProcessingError(f"Failed to ingest folder: {e}") from e + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest folder: {e}") + raise ProcessingError(f"Failed to ingest folder: {e}") from e + + def ingest_file( + self, + file_id: str, + download: bool = False, + **options, + ) -> Dict[str, Any]: + """ + Ingest data from Google Drive file. + + This method retrieves metadata and optionally downloads content + from a Google Drive file. + + Args: + file_id: Google Drive file ID + download: Whether to download file content + **options: Additional processing options + + Returns: + Dictionary containing file metadata and optionally content + + Raises: + ProcessingError: If ingestion fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=file_id, + module="ingest", + submodule="GDriveIngestor", + message=f"Ingesting file: {file_id}", + ) + + try: + # Authenticate + self._authenticate() + + # Get file metadata + file_metadata = ( + self.service.files().get(fileId=file_id, fields="*").execute() + ) + + result = {"metadata": file_metadata} + + # Download content if requested + if download: + self.progress_tracker.update_tracking( + tracking_id, message="Downloading file content..." + ) + + # Check if file is Google Workspace file (needs export) + mime_type = file_metadata.get("mimeType", "") + if mime_type.startswith("application/vnd.google-apps"): + # Export Google Workspace file + export_mime_type = options.get( + "export_mime_type", "application/pdf" + ) + request = self.service.files().export_media( + fileId=file_id, mimeType=export_mime_type + ) + else: + # Download regular file + request = self.service.files().get_media(fileId=file_id) + + # Download to bytes + file_content = io.BytesIO() + downloader = MediaIoBaseDownload(file_content, request) + done = False + while not done: + status, done = downloader.next_chunk() + + result["content"] = file_content.getvalue() + + self.progress_tracker.stop_tracking( + tracking_id, status="completed", message="File ingested" + ) + + self.logger.info(f"File ingestion completed: {file_id}") + + return result + + except HttpError as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest file: {e}") + raise ProcessingError(f"Failed to ingest file: {e}") from e + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest file: {e}") + raise ProcessingError(f"Failed to ingest file: {e}") from e + + def export_drive( + self, + folder_id: Optional[str] = None, + **options, + ) -> GDriveData: + """ + Export entire Google Drive or a folder. + + This method exports all files from a Google Drive folder or the entire drive. + + Args: + folder_id: Optional folder ID (exports entire drive if not provided) + **options: Additional export options + + Returns: + GDriveData: Exported data object + + Raises: + ProcessingError: If export fails + """ + if folder_id: + return self.ingest_folder(folder_id, include_subfolders=True, **options) + else: + # Export entire drive (root folder) + return self.ingest_folder("root", include_subfolders=True, **options) + diff --git a/semantica/ingest/huggingface_ingestor.py b/semantica/ingest/huggingface_ingestor.py new file mode 100644 index 00000000..7725dc13 --- /dev/null +++ b/semantica/ingest/huggingface_ingestor.py @@ -0,0 +1,319 @@ +""" +HuggingFace Ingestion Module + +This module provides comprehensive HuggingFace datasets ingestion capabilities for the +Semantica framework, enabling data extraction from HuggingFace Hub datasets. + +Key Features: + - Dataset ingestion + - Streaming dataset support + - Split export + - Large dataset handling + +Main Classes: + - HuggingFaceIngestor: Main HuggingFace ingestion class + - HFData: Data representation for HuggingFace ingestion + +Example Usage: + >>> from semantica.ingest import HuggingFaceIngestor + >>> ingestor = HuggingFaceIngestor() + >>> data = ingestor.ingest_dataset("squad", split="train") + >>> stream_data = ingestor.stream_dataset("squad", split="train") +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +try: + from datasets import load_dataset, Dataset, IterableDataset +except ImportError: + load_dataset = None + Dataset = None + IterableDataset = None + + +@dataclass +class HFData: + """HuggingFace data representation.""" + + data: List[Dict[str, Any]] + row_count: int + columns: List[str] + dataset_name: str + split: Optional[str] = None + schema: Dict[str, Any] = field(default_factory=dict) + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class HuggingFaceIngestor: + """ + HuggingFace ingestion handler. + + This class provides comprehensive HuggingFace datasets ingestion capabilities, + loading datasets from HuggingFace Hub and converting them to standard formats. + + Features: + - Dataset ingestion + - Streaming dataset support + - Split export + - Large dataset handling + + Example Usage: + >>> ingestor = HuggingFaceIngestor() + >>> data = ingestor.ingest_dataset("squad", split="train") + >>> stream_data = ingestor.stream_dataset("squad", split="train") + """ + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize HuggingFace ingestor. + + Args: + config: Optional HuggingFace ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + if load_dataset is None: + raise ImportError( + "datasets is required for HuggingFaceIngestor. Install it with: pip install datasets" + ) + + self.logger = get_logger("huggingface_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("HuggingFace ingestor initialized") + + def ingest_dataset( + self, + dataset_name: str, + split: Optional[str] = None, + limit: Optional[int] = None, + **options, + ) -> HFData: + """ + Ingest data from HuggingFace dataset. + + This method loads a dataset from HuggingFace Hub and converts it to + a standard format. + + Args: + dataset_name: Name of the dataset (e.g., "squad", "glue") + split: Dataset split to load (e.g., "train", "test", "validation") + limit: Maximum number of rows to return (optional) + **options: Additional options passed to load_dataset() + + Returns: + HFData: Ingested data object containing: + - data: List of row dictionaries + - row_count: Number of rows + - columns: List of column names + - dataset_name: Dataset name + - split: Split name + - schema: Schema information + + Raises: + ValidationError: If dataset not found + ProcessingError: If ingestion fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=dataset_name, + module="ingest", + submodule="HuggingFaceIngestor", + message=f"Loading dataset: {dataset_name}", + ) + + try: + # Load dataset + self.progress_tracker.update_tracking( + tracking_id, message="Loading from HuggingFace Hub..." + ) + + if split: + dataset = load_dataset(dataset_name, split=split, **options) + else: + # Load all splits + dataset = load_dataset(dataset_name, **options) + # Use first split if multiple splits + if isinstance(dataset, dict): + split = list(dataset.keys())[0] + dataset = dataset[split] + + # Convert to list of dictionaries + self.progress_tracker.update_tracking( + tracking_id, message="Converting to list format..." + ) + + if limit: + data = [dataset[i] for i in range(min(limit, len(dataset)))] + else: + data = list(dataset) + + # Extract columns and schema + columns = list(data[0].keys()) if data else [] + schema = { + "columns": columns, + "features": dataset.features.to_dict() if hasattr(dataset, "features") else {}, + } + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested {len(data)} rows", + ) + + self.logger.info( + f"Dataset ingestion completed: {len(data)} row(s) from {dataset_name}" + ) + + return HFData( + data=data, + row_count=len(data), + columns=columns, + dataset_name=dataset_name, + split=split, + schema=schema, + metadata={"features": schema.get("features", {})}, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest dataset {dataset_name}: {e}") + raise ProcessingError(f"Failed to ingest dataset: {e}") from e + + def stream_dataset( + self, + dataset_name: str, + split: Optional[str] = None, + limit: Optional[int] = None, + **options, + ) -> HFData: + """ + Stream dataset from HuggingFace Hub. + + This method loads a dataset in streaming mode, which is more memory-efficient + for large datasets. + + Args: + dataset_name: Name of the dataset + split: Dataset split to load (optional) + limit: Maximum number of rows to return (optional) + **options: Additional options passed to load_dataset() + + Returns: + HFData: Streamed data object + + Raises: + ProcessingError: If streaming fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=dataset_name, + module="ingest", + submodule="HuggingFaceIngestor", + message=f"Streaming dataset: {dataset_name}", + ) + + try: + # Load dataset in streaming mode + self.progress_tracker.update_tracking( + tracking_id, message="Loading in streaming mode..." + ) + + streaming_options = options.copy() + streaming_options["streaming"] = True + + if split: + dataset = load_dataset( + dataset_name, split=split, **streaming_options + ) + else: + dataset = load_dataset(dataset_name, **streaming_options) + if isinstance(dataset, dict): + split = list(dataset.keys())[0] + dataset = dataset[split] + + # Convert streaming dataset to list + self.progress_tracker.update_tracking( + tracking_id, message="Streaming and converting data..." + ) + + data = [] + columns = None + + for i, item in enumerate(dataset): + if columns is None: + columns = list(item.keys()) + + data.append(item) + + if limit and len(data) >= limit: + break + + # Extract schema + schema = { + "columns": columns or [], + "features": dataset.features.to_dict() if hasattr(dataset, "features") else {}, + } + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Streamed {len(data)} rows", + ) + + self.logger.info( + f"Dataset streaming completed: {len(data)} row(s) from {dataset_name}" + ) + + return HFData( + data=data, + row_count=len(data), + columns=columns or [], + dataset_name=dataset_name, + split=split, + schema=schema, + metadata={"streaming": True}, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to stream dataset {dataset_name}: {e}") + raise ProcessingError(f"Failed to stream dataset: {e}") from e + + def export_split( + self, + dataset_name: str, + split: str, + **options, + ) -> HFData: + """ + Export a specific split from a dataset. + + This method exports a specific split from a HuggingFace dataset. + + Args: + dataset_name: Name of the dataset + split: Split name to export + **options: Additional options + + Returns: + HFData: Exported split data object + + Raises: + ProcessingError: If export fails + """ + return self.ingest_dataset(dataset_name, split=split, **options) + diff --git a/semantica/ingest/mongo_ingestor.py b/semantica/ingest/mongo_ingestor.py new file mode 100644 index 00000000..188d08f2 --- /dev/null +++ b/semantica/ingest/mongo_ingestor.py @@ -0,0 +1,470 @@ +""" +MongoDB Ingestion Module + +This module provides comprehensive MongoDB ingestion capabilities for the +Semantica framework, enabling data extraction from MongoDB collections and databases. + +Key Features: + - Collection ingestion + - Query-based document retrieval + - Database export + - Schema extraction + - Large dataset handling with pagination + +Main Classes: + - MongoIngestor: Main MongoDB ingestion class + - MongoConnector: MongoDB connection handler + - MongoData: Data representation for MongoDB ingestion + +Example Usage: + >>> from semantica.ingest import MongoIngestor + >>> ingestor = MongoIngestor() + >>> data = ingestor.ingest_collection("mongodb://localhost:27017", "mydb", "mycollection") + >>> docs = ingestor.query_documents("mongodb://...", "mydb", "mycollection", {"status": "active"}) +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +try: + from pymongo import MongoClient + from pymongo.errors import ConnectionFailure, OperationFailure +except ImportError: + MongoClient = None + ConnectionFailure = None + OperationFailure = None + + +@dataclass +class MongoData: + """MongoDB data representation.""" + + documents: List[Dict[str, Any]] + document_count: int + collection_name: str + database_name: str + schema: Dict[str, Any] + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class MongoConnector: + """ + MongoDB connection management. + + This class manages connections to MongoDB, handles connection pooling, + and provides a unified interface for database operations. + + Example Usage: + >>> connector = MongoConnector() + >>> client = connector.connect("mongodb://localhost:27017") + >>> connector.disconnect() + """ + + def __init__(self, **config): + """ + Initialize MongoDB connector. + + Args: + **config: Connection configuration options + """ + if MongoClient is None: + raise ImportError( + "pymongo is required for MongoConnector. Install it with: pip install pymongo" + ) + + self.logger = get_logger("mongo_connector") + self.config = config + self.client: Optional[MongoClient] = None + + self.logger.debug("MongoDB connector initialized") + + def connect(self, connection_string: str) -> MongoClient: + """ + Establish MongoDB connection. + + Args: + connection_string: MongoDB connection string + Example: "mongodb://localhost:27017" + + Returns: + MongoClient: MongoDB client object + + Raises: + ProcessingError: If connection fails + """ + try: + self.client = MongoClient(connection_string, **self.config) + # Test connection + self.client.admin.command("ping") + self.logger.info("Connected to MongoDB") + return self.client + except Exception as e: + self.logger.error(f"Failed to connect to MongoDB: {e}") + raise ProcessingError(f"Failed to connect to MongoDB: {e}") from e + + def disconnect(self): + """Close MongoDB connection.""" + if self.client: + self.client.close() + self.client = None + self.logger.info("Disconnected from MongoDB") + + def test_connection(self, connection_string: str) -> bool: + """ + Test MongoDB connection without creating a persistent connection. + + Args: + connection_string: MongoDB connection string to test + + Returns: + bool: True if connection successful, False otherwise + """ + try: + client = MongoClient(connection_string, serverSelectionTimeoutMS=5000) + client.admin.command("ping") + client.close() + return True + except Exception as e: + self.logger.debug(f"Connection test failed: {e}") + return False + + +class MongoIngestor: + """ + MongoDB ingestion handler. + + This class provides comprehensive MongoDB ingestion capabilities, + connecting to MongoDB, querying collections, and exporting data. + + Features: + - Collection ingestion + - Query-based document retrieval + - Database export + - Schema extraction + - Large dataset handling with pagination + + Example Usage: + >>> ingestor = MongoIngestor() + >>> data = ingestor.ingest_collection("mongodb://...", "mydb", "mycollection") + >>> docs = ingestor.query_documents("mongodb://...", "mydb", "mycollection", {"status": "active"}) + """ + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize MongoDB ingestor. + + Args: + config: Optional MongoDB ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + self.logger = get_logger("mongo_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize connector + self.connector = MongoConnector(**self.config) + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("MongoDB ingestor initialized") + + def ingest_collection( + self, + connection_string: str, + database_name: str, + collection_name: str, + limit: Optional[int] = None, + **options, + ) -> MongoData: + """ + Ingest data from MongoDB collection. + + This method connects to MongoDB, retrieves documents from a collection, + and extracts schema information. + + Args: + connection_string: MongoDB connection string + database_name: Name of the database + collection_name: Name of the collection + limit: Maximum number of documents to retrieve (optional) + **options: Additional processing options + + Returns: + MongoData: Ingested data object containing: + - documents: List of document dictionaries + - document_count: Number of documents + - collection_name: Collection name + - database_name: Database name + - schema: Schema information dictionary + + Raises: + ProcessingError: If ingestion fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=f"{database_name}.{collection_name}", + module="ingest", + submodule="MongoIngestor", + message=f"Collection: {database_name}.{collection_name}", + ) + + try: + # Connect to MongoDB + client = self.connector.connect(connection_string) + db = client[database_name] + collection = db[collection_name] + + # Get document count + total_count = collection.count_documents({}) + + # Retrieve documents + self.progress_tracker.update_tracking( + tracking_id, message=f"Retrieving documents (total: {total_count})..." + ) + + query = options.get("query", {}) + cursor = collection.find(query) + + if limit: + cursor = cursor.limit(limit) + + documents = list(cursor) + + # Extract schema information + schema = self._extract_schema(documents) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested {len(documents)} documents", + ) + + self.logger.info( + f"Collection ingestion completed: {len(documents)} document(s)" + ) + + return MongoData( + documents=documents, + document_count=len(documents), + collection_name=collection_name, + database_name=database_name, + schema=schema, + metadata={"total_count": total_count, "query": query}, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest collection: {e}") + raise ProcessingError(f"Failed to ingest collection: {e}") from e + finally: + self.connector.disconnect() + + def query_documents( + self, + connection_string: str, + database_name: str, + collection_name: str, + query: Dict[str, Any], + limit: Optional[int] = None, + **options, + ) -> MongoData: + """ + Query documents from MongoDB collection. + + This method executes a query on a MongoDB collection and returns + matching documents. + + Args: + connection_string: MongoDB connection string + database_name: Name of the database + collection_name: Name of the collection + query: MongoDB query dictionary + limit: Maximum number of documents to return (optional) + **options: Additional query options + + Returns: + MongoData: Query result data object + + Raises: + ProcessingError: If query fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=f"{database_name}.{collection_name}", + module="ingest", + submodule="MongoIngestor", + message=f"Querying collection: {database_name}.{collection_name}", + ) + + try: + # Connect to MongoDB + client = self.connector.connect(connection_string) + db = client[database_name] + collection = db[collection_name] + + # Execute query + cursor = collection.find(query) + + if limit: + cursor = cursor.limit(limit) + + if "sort" in options: + cursor = cursor.sort(options["sort"]) + + documents = list(cursor) + + # Extract schema information + schema = self._extract_schema(documents) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Query returned {len(documents)} documents", + ) + + self.logger.info(f"Query completed: {len(documents)} document(s)") + + return MongoData( + documents=documents, + document_count=len(documents), + collection_name=collection_name, + database_name=database_name, + schema=schema, + metadata={"query": query}, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to query documents: {e}") + raise ProcessingError(f"Failed to query documents: {e}") from e + finally: + self.connector.disconnect() + + def export_database( + self, + connection_string: str, + database_name: str, + include_collections: Optional[List[str]] = None, + exclude_collections: Optional[List[str]] = None, + **options, + ) -> Dict[str, MongoData]: + """ + Export entire database. + + This method exports all collections from a MongoDB database. + + Args: + connection_string: MongoDB connection string + database_name: Name of the database + include_collections: List of collection names to include (optional) + exclude_collections: List of collection names to exclude (optional) + **options: Additional export options + + Returns: + Dictionary mapping collection names to MongoData objects + + Raises: + ProcessingError: If export fails + """ + tracking_id = self.progress_tracker.start_tracking( + file=database_name, + module="ingest", + submodule="MongoIngestor", + message=f"Exporting database: {database_name}", + ) + + try: + # Connect to MongoDB + client = self.connector.connect(connection_string) + db = client[database_name] + + # Get all collection names + all_collections = db.list_collection_names() + + # Apply filters + if include_collections: + collections = [c for c in all_collections if c in include_collections] + else: + exclude_collections = exclude_collections or [] + collections = [c for c in all_collections if c not in exclude_collections] + + self.progress_tracker.update_tracking( + tracking_id, message=f"Exporting {len(collections)} collections..." + ) + + # Export each collection + result = {} + for collection_name in collections: + try: + data = self.ingest_collection( + connection_string, database_name, collection_name, **options + ) + result[collection_name] = data + except Exception as e: + self.logger.error( + f"Failed to export collection {collection_name}: {e}" + ) + if self.config.get("fail_fast", False): + raise + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Exported {len(result)} collections", + ) + + self.logger.info(f"Database export completed: {len(result)} collection(s)") + + return result + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to export database: {e}") + raise ProcessingError(f"Failed to export database: {e}") from e + finally: + self.connector.disconnect() + + def _extract_schema(self, documents: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Extract schema information from documents. + + Args: + documents: List of MongoDB documents + + Returns: + Dictionary containing schema information + """ + if not documents: + return {"fields": [], "field_types": {}} + + # Collect all field names and types + field_types = {} + for doc in documents: + for key, value in doc.items(): + if key not in field_types: + field_types[key] = set() + field_types[key].add(type(value).__name__) + + # Convert sets to lists + schema = { + "fields": list(field_types.keys()), + "field_types": { + k: list(v) if len(v) > 1 else list(v)[0] + for k, v in field_types.items() + }, + } + + return schema + diff --git a/semantica/ingest/pandas_ingestor.py b/semantica/ingest/pandas_ingestor.py new file mode 100644 index 00000000..a58188ae --- /dev/null +++ b/semantica/ingest/pandas_ingestor.py @@ -0,0 +1,341 @@ +""" +Pandas Ingestion Module + +This module provides comprehensive pandas DataFrame ingestion capabilities for the +Semantica framework, enabling data extraction from DataFrames, CSV, JSON, and dictionaries. + +Key Features: + - DataFrame ingestion from various sources + - CSV, JSON, dictionary conversion + - Data transformation and validation + - Schema extraction + - Large dataset handling with chunking + +Main Classes: + - PandasIngestor: Main pandas ingestion class + - PandasData: Data representation for pandas ingestion + +Example Usage: + >>> from semantica.ingest import PandasIngestor + >>> import pandas as pd + >>> ingestor = PandasIngestor() + >>> df = pd.DataFrame({"name": ["Alice", "Bob"], "age": [25, 30]}) + >>> data = ingestor.ingest_dataframe(df) + >>> csv_data = ingestor.from_csv("data.csv") +""" + +import json +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +try: + import pandas as pd +except ImportError: + pd = None + + +@dataclass +class PandasData: + """Pandas data representation.""" + + dataframe: Any # pd.DataFrame + row_count: int + column_count: int + columns: List[str] + dtypes: Dict[str, str] + schema: Dict[str, Any] + metadata: Dict[str, Any] = field(default_factory=dict) + ingested_at: datetime = field(default_factory=datetime.now) + + +class PandasIngestor: + """ + Pandas ingestion handler. + + This class provides comprehensive pandas DataFrame ingestion capabilities, + supporting ingestion from DataFrames, CSV files, JSON files, and dictionaries. + + Features: + - DataFrame ingestion + - CSV file reading + - JSON file reading + - Dictionary conversion + - Schema extraction + - Data validation + - Large dataset chunking + + Example Usage: + >>> ingestor = PandasIngestor() + >>> data = ingestor.ingest_dataframe(df) + >>> csv_data = ingestor.from_csv("data.csv") + """ + + def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs): + """ + Initialize pandas ingestor. + + Args: + config: Optional pandas ingestion configuration dictionary + **kwargs: Additional configuration parameters (merged into config) + """ + if pd is None: + raise ImportError( + "pandas is required for PandasIngestor. Install it with: pip install pandas" + ) + + self.logger = get_logger("pandas_ingestor") + self.config = config or {} + self.config.update(kwargs) + + # Initialize progress tracker + self.progress_tracker = get_progress_tracker() + + self.logger.debug("Pandas ingestor initialized") + + def ingest_dataframe( + self, dataframe: Any, **options + ) -> PandasData: + """ + Ingest data from pandas DataFrame. + + This method processes a pandas DataFrame and extracts metadata, + schema information, and data statistics. + + Args: + dataframe: pandas DataFrame to ingest + **options: Additional processing options + + Returns: + PandasData: Ingested data object containing: + - dataframe: Original DataFrame + - row_count: Number of rows + - column_count: Number of columns + - columns: List of column names + - dtypes: Dictionary mapping column names to data types + - schema: Schema information dictionary + - metadata: Additional metadata + + Raises: + ValidationError: If dataframe is not a valid pandas DataFrame + ProcessingError: If ingestion fails + """ + if not isinstance(dataframe, pd.DataFrame): + raise ValidationError( + f"Expected pandas DataFrame, got {type(dataframe).__name__}" + ) + + tracking_id = self.progress_tracker.start_tracking( + file="dataframe", + module="ingest", + submodule="PandasIngestor", + message="Ingesting pandas DataFrame", + ) + + try: + # Extract basic information + row_count = len(dataframe) + column_count = len(dataframe.columns) + columns = list(dataframe.columns) + dtypes = {col: str(dtype) for col, dtype in dataframe.dtypes.items()} + + # Build schema information + schema = { + "columns": columns, + "dtypes": dtypes, + "shape": (row_count, column_count), + "index_type": str(type(dataframe.index).__name__), + "has_nulls": dataframe.isnull().any().any(), + "null_counts": dataframe.isnull().sum().to_dict(), + } + + # Extract metadata + metadata = { + "memory_usage": dataframe.memory_usage(deep=True).sum(), + "index_name": dataframe.index.name, + } + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Ingested DataFrame: {row_count} rows, {column_count} columns", + ) + + self.logger.info( + f"DataFrame ingestion completed: {row_count} row(s), {column_count} column(s)" + ) + + return PandasData( + dataframe=dataframe, + row_count=row_count, + column_count=column_count, + columns=columns, + dtypes=dtypes, + schema=schema, + metadata=metadata, + ) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest DataFrame: {e}") + raise ProcessingError(f"Failed to ingest DataFrame: {e}") from e + + def from_csv( + self, + file_path: Union[str, Path], + **pandas_options, + ) -> PandasData: + """ + Ingest data from CSV file. + + This method reads a CSV file using pandas and ingests it as a DataFrame. + + Args: + file_path: Path to CSV file + **pandas_options: Additional options passed to pd.read_csv() + + Returns: + PandasData: Ingested data object + + Raises: + ProcessingError: If CSV reading fails + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise ValidationError(f"CSV file not found: {file_path}") + + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="ingest", + submodule="PandasIngestor", + message=f"Ingesting CSV: {file_path.name}", + ) + + try: + # Read CSV with pandas + dataframe = pd.read_csv(file_path, **pandas_options) + + self.progress_tracker.update_tracking( + tracking_id, message="CSV read successfully, processing DataFrame..." + ) + + # Ingest the DataFrame + return self.ingest_dataframe(dataframe, **pandas_options) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest CSV {file_path}: {e}") + raise ProcessingError(f"Failed to ingest CSV: {e}") from e + + def from_json( + self, + file_path: Union[str, Path], + **pandas_options, + ) -> PandasData: + """ + Ingest data from JSON file. + + This method reads a JSON file using pandas and ingests it as a DataFrame. + + Args: + file_path: Path to JSON file + **pandas_options: Additional options passed to pd.read_json() + + Returns: + PandasData: Ingested data object + + Raises: + ProcessingError: If JSON reading fails + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise ValidationError(f"JSON file not found: {file_path}") + + tracking_id = self.progress_tracker.start_tracking( + file=str(file_path), + module="ingest", + submodule="PandasIngestor", + message=f"Ingesting JSON: {file_path.name}", + ) + + try: + # Read JSON with pandas + dataframe = pd.read_json(file_path, **pandas_options) + + self.progress_tracker.update_tracking( + tracking_id, message="JSON read successfully, processing DataFrame..." + ) + + # Ingest the DataFrame + return self.ingest_dataframe(dataframe, **pandas_options) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest JSON {file_path}: {e}") + raise ProcessingError(f"Failed to ingest JSON: {e}") from e + + def from_dict( + self, + data: Union[Dict[str, List], List[Dict[str, Any]]], + **pandas_options, + ) -> PandasData: + """ + Ingest data from dictionary. + + This method converts a dictionary or list of dictionaries to a DataFrame + and ingests it. + + Args: + data: Dictionary or list of dictionaries + **pandas_options: Additional options passed to pd.DataFrame() + + Returns: + PandasData: Ingested data object + + Raises: + ValidationError: If data format is invalid + ProcessingError: If conversion fails + """ + if not isinstance(data, (dict, list)): + raise ValidationError( + f"Expected dict or list, got {type(data).__name__}" + ) + + tracking_id = self.progress_tracker.start_tracking( + file="dictionary", + module="ingest", + submodule="PandasIngestor", + message="Converting dictionary to DataFrame", + ) + + try: + # Convert to DataFrame + dataframe = pd.DataFrame(data, **pandas_options) + + self.progress_tracker.update_tracking( + tracking_id, message="Dictionary converted, processing DataFrame..." + ) + + # Ingest the DataFrame + return self.ingest_dataframe(dataframe, **pandas_options) + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + self.logger.error(f"Failed to ingest dictionary: {e}") + raise ProcessingError(f"Failed to ingest dictionary: {e}") from e +