mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-09-15 04:00:33 +00:00
feat(ingest): Add 7 new data source ingestors
- Add PandasIngestor for DataFrame, CSV, JSON ingestion - Add DuckDBIngestor for CSV, Parquet, Excel with SQL queries - Add MongoIngestor for MongoDB document databases - Add ElasticIngestor for Elasticsearch indices - Add RESTIngestor for generic REST API endpoints - Add HuggingFaceIngestor for ML datasets from HuggingFace Hub - Add GDriveIngestor for Google Drive files and folders - Update registry, methods, and config for new ingestors - Add comprehensive documentation and code examples - Add optional dependencies to pyproject.toml
This commit is contained in:
KaifAhmad1
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user