Flowfile 0.3.5__py3-none-any.whl → 0.3.6__py3-none-any.whl

flowfile_worker/external_sources/airbyte_sources/cache_manager.py

@@ -1,161 +0,0 @@
-from airbyte.caches import DuckDBCache
-
-import os
-import atexit
-from pathlib import Path
-import filelock
-from contextlib import contextmanager
-
-from flowfile_worker.configs import logger
-
-
-class DuckDBCacheManager:
-    """
-    Manages DuckDB cache instances with multiprocessing-safe lock handling.
-    Coordinates cache access across different processes using file-based locks.
-    """
-
-    def __init__(self,
-                 base_path: Path = Path.home() / ".flowfile/.tmp",
-                 max_retries: int = 3,
-                 retry_delay: float = 0.5):
-        logger.info(f"Initializing DuckDBCacheManager with base_path: {base_path}")
-
-        # Shared resources path
-        self.base_path = base_path
-        self.base_path.mkdir(parents=True, exist_ok=True)
-
-        # Central coordination directory
-        self.coordinator_dir = self.base_path / "coordinator"
-        self.coordinator_dir.mkdir(exist_ok=True)
-
-        # Process-specific information
-        self.process_id = os.getpid()
-        self.max_retries = max_retries
-        self.retry_delay = retry_delay
-
-        # Global lock for cache allocation
-        self.global_lock_file = self.coordinator_dir / "global.lock"
-        logger.debug(f"Process {self.process_id} initialized with global lock at {self.global_lock_file}")
-
-        # Register cleanup
-        atexit.register(self.cleanup_process_resources)
-
-    @contextmanager
-    def get_cache(self) -> 'DuckDBCache':
-        """
-        Get an available cache instance with proper locking.
-        """
-        logger.debug(f"Process {self.process_id} attempting to get cache")
-        cache_id = None
-        cache = None
-
-        # First, try to get any available existing cache
-        with filelock.FileLock(str(self.global_lock_file), timeout=0.1):
-            logger.debug(f"Process {self.process_id} acquired global lock")
-            for i in range(10):  # Check first 10 possible cache slots
-                try:
-                    cache_lock = filelock.FileLock(
-                        str(self.coordinator_dir / f"cache_{i}.lock"),
-                        timeout=0.1
-                    )
-                    cache_lock.acquire()
-
-                    # If we got the lock, use this cache slot
-                    cache_id = i
-                    cache_path = self.base_path / f"airbyte_cache_{i}"
-                    logger.info(f"Process {self.process_id} acquired cache slot {i} at {cache_path}")
-
-                    cache = DuckDBCache(
-                        db_path=cache_path,
-                        schema_name="main"
-                    )
-
-                    # Keep track of which process is using this cache
-                    with open(self.coordinator_dir / f"cache_{i}.pid", 'w') as f:
-                        f.write(str(self.process_id))
-
-                    break
-
-                except filelock.Timeout:
-                    logger.debug(f"Process {self.process_id} failed to acquire lock for cache slot {i}")
-                    continue
-
-        if cache is None:
-            logger.info(f"Process {self.process_id} couldn't acquire existing cache, creating new one")
-            # If no existing cache is available, create a new one
-            cache_id = self._create_new_cache_slot()
-            cache_path = self.base_path / f"airbyte_cache_{cache_id}"
-            cache = DuckDBCache(
-                db_path=cache_path,
-                schema_name="main"
-            )
-
-        try:
-            yield cache
-        finally:
-            # Cleanup
-            if cache_id is not None:
-                lock_file = self.coordinator_dir / f"cache_{cache_id}.lock"
-                pid_file = self.coordinator_dir / f"cache_{cache_id}.pid"
-                logger.debug(f"Process {self.process_id} cleaning up cache slot {cache_id}")
-
-                try:
-                    if lock_file.exists():
-                        lock_file.unlink()
-                    if pid_file.exists():
-                        pid_file.unlink()
-                except OSError as e:
-                    logger.error(f"Process {self.process_id} failed to cleanup files for cache {cache_id}: {str(e)}")
-
-    def _create_new_cache_slot(self) -> int:
-        """Create a new cache slot with proper locking."""
-        logger.debug(f"Process {self.process_id} creating new cache slot")
-        with filelock.FileLock(str(self.global_lock_file), timeout=10):
-            # Find the next available slot
-            existing_caches = set()
-            for cache_file in self.coordinator_dir.glob("cache_*.lock"):
-                try:
-                    slot = int(cache_file.stem.split('_')[1])
-                    existing_caches.add(slot)
-                except (ValueError, IndexError):
-                    continue
-
-            # Get the first available slot number
-            new_slot = 0
-            while new_slot in existing_caches:
-                new_slot += 1
-
-            # Create the lock file for this slot
-            lock_file = self.coordinator_dir / f"cache_{new_slot}.lock"
-            lock_file.touch()
-            logger.info(f"Process {self.process_id} created new cache slot {new_slot}")
-
-            return new_slot
-
-    def cleanup_process_resources(self):
-        """Clean up resources when the process exits."""
-        logger.debug(f"Process {self.process_id} starting cleanup")
-        # Clean up any cache slots owned by this process
-        for pid_file in self.coordinator_dir.glob("*.pid"):
-            try:
-                with open(pid_file, 'r') as f:
-                    pid = int(f.read().strip())
-
-                if pid == self.process_id:
-                    cache_id = pid_file.stem.split('_')[1]
-                    logger.info(f"Process {self.process_id} cleaning up cache slot {cache_id}")
-
-                    # Remove the cache files
-                    cache_path = self.base_path / f"airbyte_cache_{cache_id}"
-                    lock_file = self.coordinator_dir / f"cache_{cache_id}.lock"
-
-                    for file in [pid_file, lock_file, cache_path]:
-                        try:
-                            if file.exists():
-                                file.unlink()
-                                logger.debug(f"Process {self.process_id} removed file: {file}")
-                        except OSError as e:
-                            logger.error(f"Process {self.process_id} failed to remove file {file}: {str(e)}")
-            except (ValueError, OSError) as e:
-                logger.error(f"Process {self.process_id} encountered error during cleanup: {str(e)}")

flowfile_worker/external_sources/airbyte_sources/main.py

@@ -1,89 +0,0 @@
-from ast import literal_eval
-import polars as pl
-from typing import Optional
-import airbyte as ab
-import os
-from airbyte import exceptions as exc
-
-
-from flowfile_worker.configs import logger
-from flowfile_worker.external_sources.airbyte_sources.models import AirbyteSettings
-from flowfile_worker.external_sources.airbyte_sources.cache_manager import DuckDBCacheManager
-
-
-class AirbyteGetter:
-    stream: str
-    source_name: str
-    _type: str
-    _airbyte_module = None
-    _enforce_full_refresh: Optional[bool] = True
-    version: Optional[str] = None
-
-    def __init__(self, airbyte_settings: AirbyteSettings):
-        self._airbyte_response = None
-        self.stream = airbyte_settings.stream
-        self.source_name = airbyte_settings.source_name
-        self._enforce_full_refresh = airbyte_settings.enforce_full_refresh
-        self.cache_manager = DuckDBCacheManager()
-        if hasattr(airbyte_settings, 'version'):
-            self.version = airbyte_settings.version
-
-        # Handle config
-        if airbyte_settings.config_ref and not airbyte_settings.config:
-            logger.info(f"Getting config from {airbyte_settings.config_ref}")
-            config = literal_eval(os.environ.get(airbyte_settings.config_ref))
-        else:
-            logger.info(f"Using provided config")
-            config = airbyte_settings.config
-
-        if config is None:
-            raise ValueError("Config must be provided")
-
-        self.config = config
-        self._type = 'airbyte'
-        self.read_result = None
-
-    def __call__(self) -> pl.DataFrame:
-        with self.cache_manager.get_cache() as cache:
-            if self.read_result is None:
-                # Lazy import airbyte
-
-                try:
-                    source = ab.get_source(
-                        name=self.source_name,
-                        config=self.config,
-                        streams=self.stream,
-                        docker_image=True,
-                        version=self.version
-                    )
-
-                    logger.debug(f'Starting to load data for {self.stream}')  # Changed to debug level
-
-                    self.read_result = source.read(
-                        cache=cache,
-                        force_full_refresh=self._enforce_full_refresh
-                    )
-
-                except Exception as e:
-                    logger.error(f"Error during source operation: {str(e)}")
-                    raise
-
-            df = self.read_result[self.stream].to_pandas()
-            drop_cols = [c for c in df.columns if c.startswith('_airbyte')]
-            df.drop(drop_cols, axis=1, inplace=True)
-            return pl.from_pandas(df)
-
-
-def read_airbyte_source(airbyte_settings: AirbyteSettings) -> pl.DataFrame:
-    """
-    Read data from an Airbyte source and return it as a Polars DataFrame.
-    Args:
-        airbyte_settings (): The settings for the Airbyte source.
-
-    Returns: The data as a Polars DataFrame.
-    """
-    airbyte_getter = AirbyteGetter(airbyte_settings)
-    logger.info('Getting data from Airbyte')
-    data = airbyte_getter()
-    logger.info('Data retrieved from Airbyte')
-    return data

flowfile_worker/external_sources/airbyte_sources/models.py

@@ -1,133 +0,0 @@
-from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union
-from pydantic import BaseModel, field_validator, ConfigDict
-from flowfile_worker.configs import logger
-
-# Use TYPE_CHECKING to avoid circular imports
-if TYPE_CHECKING:
-    from airbyte import Source
-else:
-    Source = Any
-
-
-class LazyAirbyteSource:
-    """Lazy wrapper for airbyte Source class."""
-    _source_class = None
-
-    @classmethod
-    def get_source_class(cls):
-        if cls._source_class is None:
-            logger.info("Importing airbyte Source class")
-            from airbyte import Source
-            cls._source_class = Source
-        return cls._source_class
-
-
-class JsonSchema(BaseModel):
-    type: Optional[Union[str, List[str]]]
-    airbyte_type: Optional[Union[str, List[str]]] = None
-    format: Optional[str] = None
-
-
-class AirbyteProperty(BaseModel):
-    name: str
-    json_schema: JsonSchema
-
-
-class AirbyteResponse(BaseModel):
-    source: Any  # Using Any to avoid direct Source import
-    properties: list[AirbyteProperty]
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    @field_validator('source')
-    @classmethod
-    def validate_source(cls, v: Any) -> Any:
-        source_class = LazyAirbyteSource.get_source_class()
-        if not isinstance(v, source_class):
-            raise ValueError(f"Source must be an instance of airbyte.Source, got {type(v)}")
-        return v
-
-
-class GenericProperties(BaseModel):
-    type: str
-    title: Optional[str] = None
-    description: Optional[str] = None
-    order: Optional[int] = None
-    required: Optional[List[str]] = None
-    airbyte_secret: Optional[bool] = None
-    pattern: Optional[str] = None
-    pattern_descriptor: Optional[str] = None
-    format: Optional[str] = None
-    examples: Optional[List[Any]] = None
-    enum: Optional[List[str]] = None
-    minimum: Optional[float] = None
-    maximum: Optional[float] = None
-    items: Optional[Any] = None
-    properties: Optional[Dict[str, Any]] = None
-
-    @field_validator('items', 'properties')
-    @classmethod
-    def validate_nested(cls, value: Any) -> Any:
-        if isinstance(value, dict):
-            if 'type' in value:
-                return GenericProperties(**value)
-            return {k: GenericProperties(**v) if isinstance(v, dict) else v for k, v in value.items()}
-        return value
-
-
-class GenericSchema(BaseModel):
-    title: str
-    type: str
-    required: Optional[List[str]] = None
-    additionalProperties: Optional[bool] = None
-    properties: Dict[str, GenericProperties]
-
-
-class FieldProperty(BaseModel):
-    title: Optional[str] = None
-    type: str
-    key: str
-    description: Optional[str] = None
-    airbyte_secret: Optional[bool] = None
-    input_value: Optional[str] = None
-    default: Any
-
-
-class OverallFieldProperty(BaseModel):
-    title: Optional[str] = None
-    type: str
-    key: str
-    required: bool
-    properties: List[FieldProperty]
-    items: Optional[List[FieldProperty]]
-    isOpen: bool
-    description: Optional[str] = None
-    input_value: Optional[str] = None
-    airbyte_secret: Optional[bool] = None
-    default: Any
-
-
-class AirbyteConfigTemplate(BaseModel):
-    source_name: str
-    docs_url: Optional[str] = None
-    config_spec: Dict
-    available_streams: Optional[List[str]] = None
-
-
-class AirbyteSettings(BaseModel):
-
-    source_name: str
-    stream: str
-    config_ref: Optional[str] = None
-    config: Optional[Dict] = None
-    fields: Optional[List] = None
-    enforce_full_refresh: Optional[bool] = True
-    flowfile_flow_id: int = 1
-    flowfile_node_id: int | str = -1
-    version: Optional[str] = None
-
-
-def get_source_instance(*args, **kwargs) -> 'Source':
-    """Helper function to get a Source instance with lazy loading."""
-    source_class = LazyAirbyteSource.get_source_class()
-    return source_class(*args, **kwargs)