graflo 1.3.3__py3-none-any.whl

graflo/data_source/factory.py

@@ -0,0 +1,298 @@
+"""Factory for creating data source instances.
+
+This module provides a factory for creating appropriate data source instances
+based on configuration. It supports file-based, API, and SQL data sources.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+from graflo.architecture.onto import EncodingType
+from graflo.data_source.api import APIConfig, APIDataSource
+from graflo.data_source.base import AbstractDataSource, DataSourceType
+from graflo.data_source.file import (
+    JsonFileDataSource,
+    JsonlFileDataSource,
+    TableFileDataSource,
+)
+from graflo.data_source.memory import InMemoryDataSource
+from graflo.data_source.sql import SQLConfig, SQLDataSource
+from graflo.util.chunker import ChunkerFactory, ChunkerType
+
+logger = logging.getLogger(__name__)
+
+
+class DataSourceFactory:
+    """Factory for creating data source instances.
+
+    This factory creates appropriate data source instances based on the
+    provided configuration. It supports file-based, API, and SQL data sources.
+    """
+
+    @staticmethod
+    def _guess_file_type(filename: Path) -> ChunkerType:
+        """Guess the file type based on file extension.
+
+        Args:
+            filename: Path to the file
+
+        Returns:
+            ChunkerType: Guessed file type
+
+        Raises:
+            ValueError: If file extension is not recognized
+        """
+        return ChunkerFactory._guess_chunker_type(filename)
+
+    @classmethod
+    def create_file_data_source(
+        cls,
+        path: Path | str,
+        file_type: str | ChunkerType | None = None,
+        encoding: EncodingType = EncodingType.UTF_8,
+        sep: str | None = None,
+    ) -> JsonFileDataSource | JsonlFileDataSource | TableFileDataSource:
+        """Create a file-based data source.
+
+        Args:
+            path: Path to the file
+            file_type: Type of file ('json', 'jsonl', 'table') or ChunkerType.
+                If None, will be guessed from file extension.
+            encoding: File encoding (default: UTF_8)
+            sep: Field separator for table files (default: ',').
+                Only used for table files.
+
+        Returns:
+            Appropriate file data source instance (JsonFileDataSource,
+            JsonlFileDataSource, or TableFileDataSource)
+
+        Raises:
+            ValueError: If file type cannot be determined
+        """
+        if isinstance(path, str):
+            path = Path(path)
+
+        # Determine file type
+        if file_type is None:
+            try:
+                file_type_enum = cls._guess_file_type(path)
+            except ValueError as e:
+                raise ValueError(
+                    f"Could not determine file type for {path}. "
+                    f"Please specify file_type explicitly. Error: {e}"
+                )
+        elif isinstance(file_type, str):
+            file_type_enum = ChunkerType(file_type.lower())
+        else:
+            file_type_enum = file_type
+
+        # Create appropriate data source
+        if file_type_enum == ChunkerType.JSON:
+            return JsonFileDataSource(path=path, encoding=encoding)
+        elif file_type_enum == ChunkerType.JSONL:
+            return JsonlFileDataSource(path=path, encoding=encoding)
+        elif file_type_enum == ChunkerType.TABLE:
+            # sep is only for table files
+            return TableFileDataSource(path=path, encoding=encoding, sep=sep or ",")
+        else:
+            raise ValueError(f"Unsupported file type: {file_type_enum}")
+
+    @classmethod
+    def create_api_data_source(cls, config: APIConfig) -> APIDataSource:
+        """Create an API data source.
+
+        Args:
+            config: API configuration
+
+        Returns:
+            APIDataSource instance
+        """
+        return APIDataSource(config=config)
+
+    @classmethod
+    def create_sql_data_source(cls, config: SQLConfig) -> SQLDataSource:
+        """Create a SQL data source.
+
+        Args:
+            config: SQL configuration
+
+        Returns:
+            SQLDataSource instance
+        """
+        return SQLDataSource(config=config)
+
+    @classmethod
+    def create_in_memory_data_source(
+        cls,
+        data: list[dict] | list[list] | pd.DataFrame,
+        columns: list[str] | None = None,
+    ) -> InMemoryDataSource:
+        """Create an in-memory data source.
+
+        Args:
+            data: Data to process (list[dict], list[list], or pd.DataFrame)
+            columns: Optional column names for list[list] data
+
+        Returns:
+            InMemoryDataSource instance
+        """
+        return InMemoryDataSource(data=data, columns=columns)
+
+    @classmethod
+    def create_data_source(
+        cls,
+        source_type: DataSourceType | str | None = None,
+        **kwargs: Any,
+    ) -> AbstractDataSource:
+        """Create a data source of the specified type.
+
+        This is a general factory method that routes to specific factory methods
+        based on the source type.
+
+        Args:
+            source_type: Type of data source to create. If None, will be inferred
+                from kwargs (e.g., 'path' -> FILE, 'data' -> IN_MEMORY, 'config' with url -> API)
+            **kwargs: Configuration parameters for the data source
+
+        Returns:
+            Data source instance
+
+        Raises:
+            ValueError: If source type is not supported or required parameters are missing
+        """
+        # Auto-detect source type if not provided
+        if source_type is None:
+            if "path" in kwargs or "file_type" in kwargs:
+                source_type = DataSourceType.FILE
+            elif "data" in kwargs:
+                source_type = DataSourceType.IN_MEMORY
+            elif "config" in kwargs:
+                config = kwargs["config"]
+                # Check if it's an API config (has 'url') or SQL config (has 'connection_string')
+                if isinstance(config, dict):
+                    if "url" in config:
+                        source_type = DataSourceType.API
+                    elif "connection_string" in config or "query" in config:
+                        source_type = DataSourceType.SQL
+                    else:
+                        # Try to create from dict
+                        if "source_type" in config:
+                            source_type = DataSourceType(config["source_type"].lower())
+                        else:
+                            raise ValueError(
+                                "Cannot determine source type from config. "
+                                "Please specify source_type or provide 'url' (API) "
+                                "or 'connection_string'/'query' (SQL) in config."
+                            )
+                elif hasattr(config, "url"):
+                    source_type = DataSourceType.API
+                elif hasattr(config, "connection_string") or hasattr(config, "query"):
+                    source_type = DataSourceType.SQL
+                else:
+                    raise ValueError(
+                        "Cannot determine source type from config. "
+                        "Please specify source_type explicitly."
+                    )
+            else:
+                raise ValueError(
+                    "Cannot determine source type. Please specify source_type or "
+                    "provide one of: path (FILE), data (IN_MEMORY), or config (API/SQL)."
+                )
+
+        if isinstance(source_type, str):
+            source_type = DataSourceType(source_type.lower())
+
+        if source_type == DataSourceType.FILE:
+            return cls.create_file_data_source(**kwargs)
+        elif source_type == DataSourceType.API:
+            if "config" not in kwargs:
+                # Create APIConfig from kwargs
+                from graflo.data_source.api import APIConfig, PaginationConfig
+
+                # Handle nested pagination config manually
+                api_kwargs = kwargs.copy()
+                pagination_dict = api_kwargs.pop("pagination", None)
+                pagination = None
+                if pagination_dict is not None:
+                    if isinstance(pagination_dict, dict):
+                        # Manually construct PaginationConfig to avoid dataclass_wizard issues
+                        pagination = PaginationConfig(**pagination_dict)
+                    else:
+                        pagination = pagination_dict
+                api_kwargs["pagination"] = pagination
+                config = APIConfig(**api_kwargs)
+                return cls.create_api_data_source(config=config)
+            config = kwargs["config"]
+            if isinstance(config, dict):
+                from graflo.data_source.api import APIConfig, PaginationConfig
+
+                # Handle nested pagination config manually
+                config_copy = config.copy()
+                pagination_dict = config_copy.pop("pagination", None)
+                pagination = None
+                if pagination_dict is not None:
+                    if isinstance(pagination_dict, dict):
+                        # Manually construct PaginationConfig to avoid dataclass_wizard issues
+                        pagination = PaginationConfig(**pagination_dict)
+                    else:
+                        pagination = pagination_dict
+                config_copy["pagination"] = pagination
+                config = APIConfig(**config_copy)
+            return cls.create_api_data_source(config=config)
+        elif source_type == DataSourceType.SQL:
+            if "config" not in kwargs:
+                # Create SQLConfig from kwargs
+                from graflo.data_source.sql import SQLConfig
+
+                config = SQLConfig.from_dict(kwargs)
+                return cls.create_sql_data_source(config=config)
+            config = kwargs["config"]
+            if isinstance(config, dict):
+                from graflo.data_source.sql import SQLConfig
+
+                config = SQLConfig.from_dict(config)
+            return cls.create_sql_data_source(config=config)
+        elif source_type == DataSourceType.IN_MEMORY:
+            if "data" not in kwargs:
+                raise ValueError("In-memory data source requires 'data' parameter")
+            return cls.create_in_memory_data_source(**kwargs)
+        else:
+            raise ValueError(f"Unsupported data source type: {source_type}")
+
+    @classmethod
+    def create_data_source_from_config(
+        cls, config: dict[str, Any]
+    ) -> AbstractDataSource:
+        """Create a data source from a configuration dictionary.
+
+        The configuration dict should contain:
+        - 'source_type': Type of data source (FILE, API, SQL, IN_MEMORY)
+        - Other parameters specific to the data source type
+
+        Examples:
+            File source:
+                {"source_type": "file", "path": "data.json"}
+            API source:
+                {"source_type": "api", "config": {"url": "https://api.example.com"}}
+            SQL source:
+                {"source_type": "sql", "config": {"connection_string": "...", "query": "..."}}
+            In-memory source:
+                {"source_type": "in_memory", "data": [...]}
+
+        Args:
+            config: Configuration dictionary
+
+        Returns:
+            Data source instance
+
+        Raises:
+            ValueError: If configuration is invalid
+        """
+        config = config.copy()
+        source_type = config.pop("source_type", None)
+        return cls.create_data_source(source_type=source_type, **config)

graflo/data_source/file.py

@@ -0,0 +1,133 @@
+"""File-based data source implementations.
+
+This module provides data source implementations for file-based data sources,
+including JSON, JSONL, and CSV/TSV files. It integrates with the existing
+chunker logic for efficient batch processing.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from pathlib import Path
+from typing import Iterator
+
+from graflo.architecture.onto import EncodingType
+from graflo.data_source.base import AbstractDataSource, DataSourceType
+from graflo.util.chunker import ChunkerFactory, ChunkerType
+
+
+@dataclasses.dataclass
+class FileDataSource(AbstractDataSource):
+    """Base class for file-based data sources.
+
+    This class provides a common interface for file-based data sources,
+    integrating with the existing chunker system for batch processing.
+
+    Attributes:
+        path: Path to the file
+        file_type: Type of file (json, jsonl, table)
+        encoding: File encoding (default: UTF_8)
+    """
+
+    path: Path | str
+    file_type: str | None = None
+    encoding: EncodingType = EncodingType.UTF_8
+
+    def __post_init__(self):
+        """Initialize the file data source."""
+        self.source_type = DataSourceType.FILE
+        if isinstance(self.path, str):
+            self.path = Path(self.path)
+
+    def iter_batches(
+        self, batch_size: int = 1000, limit: int | None = None
+    ) -> Iterator[list[dict]]:
+        """Iterate over file data in batches.
+
+        Args:
+            batch_size: Number of items per batch
+            limit: Maximum number of items to retrieve
+
+        Yields:
+            list[dict]: Batches of documents as dictionaries
+        """
+        # Determine chunker type
+        chunker_type = None
+        if self.file_type:
+            chunker_type = ChunkerType(self.file_type.lower())
+
+        # Create chunker using factory
+        chunker_kwargs = {
+            "resource": self.path,
+            "type": chunker_type,
+            "batch_size": batch_size,
+            "limit": limit,
+            "encoding": self.encoding,
+        }
+        # Only add sep for table files
+        if chunker_type == ChunkerType.TABLE and hasattr(self, "sep"):
+            chunker_kwargs["sep"] = self.sep
+
+        chunker = ChunkerFactory.create_chunker(**chunker_kwargs)
+
+        # Yield batches
+        for batch in chunker:
+            yield batch
+
+
+@dataclasses.dataclass
+class JsonFileDataSource(FileDataSource):
+    """Data source for JSON files.
+
+    JSON files are expected to contain hierarchical data structures,
+    similar to REST API responses. The chunker handles nested structures
+    and converts them to dictionaries.
+
+    Attributes:
+        path: Path to the JSON file
+        encoding: File encoding (default: UTF_8)
+    """
+
+    def __post_init__(self):
+        """Initialize the JSON file data source."""
+        super().__post_init__()
+        self.file_type = ChunkerType.JSON.value
+
+
+@dataclasses.dataclass
+class JsonlFileDataSource(FileDataSource):
+    """Data source for JSONL (JSON Lines) files.
+
+    JSONL files contain one JSON object per line, making them suitable
+    for streaming and batch processing.
+
+    Attributes:
+        path: Path to the JSONL file
+        encoding: File encoding (default: UTF_8)
+    """
+
+    def __post_init__(self):
+        """Initialize the JSONL file data source."""
+        super().__post_init__()
+        self.file_type = ChunkerType.JSONL.value
+
+
+@dataclasses.dataclass
+class TableFileDataSource(FileDataSource):
+    """Data source for CSV/TSV files.
+
+    Table files are converted to dictionaries with column headers as keys.
+    Each row becomes a dictionary.
+
+    Attributes:
+        path: Path to the CSV/TSV file
+        encoding: File encoding (default: UTF_8)
+        sep: Field separator (default: ',')
+    """
+
+    sep: str = ","
+
+    def __post_init__(self):
+        """Initialize the table file data source."""
+        super().__post_init__()
+        self.file_type = ChunkerType.TABLE.value

graflo/data_source/memory.py

@@ -0,0 +1,72 @@
+"""In-memory data source implementations.
+
+This module provides data source implementations for in-memory data structures,
+including lists of dictionaries, lists of lists, and Pandas DataFrames.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Iterator
+
+import pandas as pd
+
+from graflo.data_source.base import AbstractDataSource, DataSourceType
+from graflo.util.chunker import ChunkerFactory
+
+
+@dataclasses.dataclass
+class InMemoryDataSource(AbstractDataSource):
+    """Data source for in-memory data structures.
+
+    This class provides a data source for Python objects that are already
+    in memory, including lists of dictionaries, lists of lists, and Pandas DataFrames.
+
+    Attributes:
+        data: Data to process (list[dict], list[list], or pd.DataFrame)
+        columns: Optional column names for list[list] data
+    """
+
+    data: list[dict] | list[list] | pd.DataFrame
+    columns: list[str] | None = None
+
+    def __post_init__(self):
+        """Initialize the in-memory data source."""
+        self.source_type = DataSourceType.IN_MEMORY
+
+    def iter_batches(
+        self, batch_size: int = 1000, limit: int | None = None
+    ) -> Iterator[list[dict]]:
+        """Iterate over in-memory data in batches.
+
+        Args:
+            batch_size: Number of items per batch
+            limit: Maximum number of items to retrieve
+
+        Yields:
+            list[dict]: Batches of documents as dictionaries
+        """
+        # Normalize data: convert list[list] to list[dict] if needed
+        data = self.data
+        if isinstance(data, list) and len(data) > 0 and isinstance(data[0], list):
+            # list[list] - convert to list[dict] using columns
+            if self.columns is None:
+                raise ValueError(
+                    "columns parameter is required when data is list[list]"
+                )
+            data = [{k: v for k, v in zip(self.columns, item)} for item in data]
+
+        # Create chunker using factory (only pass columns if it's a DataFrame)
+        chunker_kwargs = {
+            "resource": data,
+            "batch_size": batch_size,
+            "limit": limit,
+        }
+        # Note: columns is not passed to chunker - we handle list[list] conversion above
+        # DataFrame chunker doesn't need columns either
+
+        chunker = ChunkerFactory.create_chunker(**chunker_kwargs)
+
+        # Yield batches
+        for batch in chunker:
+            yield batch

graflo/data_source/registry.py

@@ -0,0 +1,82 @@
+"""Data source registry for mapping data sources to resources.
+
+This module provides a registry for mapping data sources to resource names.
+Many data sources can map to the same resource, allowing flexible data
+ingestion from multiple sources.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING
+
+from graflo.onto import BaseDataclass
+
+if TYPE_CHECKING:
+    from graflo.data_source.base import AbstractDataSource
+
+
+@dataclasses.dataclass
+class DataSourceRegistry(BaseDataclass):
+    """Registry for mapping data sources to resource names.
+
+    This class maintains a mapping from resource names to lists of data sources.
+    Many data sources can map to the same resource, allowing data to be ingested
+    from multiple sources and combined.
+
+    Attributes:
+        sources: Dictionary mapping resource names to lists of data sources
+    """
+
+    sources: dict[str, list[AbstractDataSource]] = dataclasses.field(
+        default_factory=dict
+    )
+
+    def register(self, data_source: AbstractDataSource, resource_name: str) -> None:
+        """Register a data source for a resource.
+
+        Args:
+            data_source: Data source to register
+            resource_name: Name of the resource to map to
+        """
+        if resource_name not in self.sources:
+            self.sources[resource_name] = []
+        self.sources[resource_name].append(data_source)
+        data_source.resource_name = resource_name
+
+    def get_data_sources(self, resource_name: str) -> list[AbstractDataSource]:
+        """Get all data sources for a resource.
+
+        Args:
+            resource_name: Name of the resource
+
+        Returns:
+            List of data sources for the resource (empty list if none found)
+        """
+        return self.sources.get(resource_name, [])
+
+    def get_all_data_sources(self) -> list[AbstractDataSource]:
+        """Get all registered data sources.
+
+        Returns:
+            List of all registered data sources
+        """
+        all_sources = []
+        for sources_list in self.sources.values():
+            all_sources.extend(sources_list)
+        return all_sources
+
+    def has_resource(self, resource_name: str) -> bool:
+        """Check if a resource has any data sources.
+
+        Args:
+            resource_name: Name of the resource
+
+        Returns:
+            True if the resource has data sources, False otherwise
+        """
+        return resource_name in self.sources and len(self.sources[resource_name]) > 0
+
+    def clear(self) -> None:
+        """Clear all registered data sources."""
+        self.sources.clear()