irides-core 0.1.0__py3-none-any.whl

core/__init__.py

@@ -0,0 +1,2 @@
+"""Irides core library."""
+

core/basic_test.py

@@ -0,0 +1,78 @@
+import sys
+
+from core.db_connector.connectors.mysql import MySQLConnector
+from core.db_connector.cache_manager import CacheManager
+from core.db_connector.models import TableDescription
+from loguru import logger
+
+# Configuration for MySQL Test (ADJUST THESE VALUES)
+MYSQL_CONFIG = {
+    "host": "<your_mysql_host>",  # e.g., "localhost" or "mysql_test_container"
+    "user": "<your_mysql_user>",  # e.g., "root"
+    "password": "<your_mysql_password>",  # e.g., "password"
+    "port": 3306,
+    # "database": "test_db" # Optional, for specific database operations
+}
+
+def run_basic_mysql_test():
+    logger.info("--- Running Basic MySQL Test ---")
+    connector = None
+    try:
+        cache_manager = CacheManager()
+        connector = MySQLConnector(connection_params=MYSQL_CONFIG, cache_manager=cache_manager)
+        logger.info("MySQL connection successful!")
+
+        logger.info("\n--- Listing Instances ---")
+        instances = connector.list_instances()
+        for instance in instances:
+            logger.info(f"Instance: {instance.name} (Version: {instance.version})")
+
+            logger.info(f"\n--- Listing Schemas for Instance: {instance.name} ---")
+            schemas = connector.list_schemas(instance_name=instance.name)
+            for schema in schemas:
+                logger.info(f"  Schema: {schema.name}")
+
+                logger.info(f"\n  --- Listing Tables for Schema: {schema.name} ---")
+                tables = connector.list_tables(instance_name=instance.name, schema_name=schema.name)
+                if not tables:
+                    logger.info(f"    No tables found in schema: {schema.name}")
+                for table in tables:
+                    logger.info(f"    Table: {table.name}")
+
+                    logger.info(f"\n      --- Describing Table: {table.name} in Schema: {schema.name} ---")
+                    table_desc: TableDescription = connector.describe_table(instance_name=instance.name, schema_name=schema.name, table_name=table.name)
+                    
+                    logger.info(f"        Columns:")
+                    if not table_desc.columns:
+                        logger.info(f"          No columns found for table: {table.name}")
+                    for col in table_desc.columns:
+                        logger.info(f"          - Name: {col.name}, Type: {col.data_type}, Nullable: {col.is_nullable}, Default: {col.default_value}")
+                    
+                    if table_desc.primary_key:
+                        logger.info(f"        Primary Key: {', '.join(table_desc.primary_key.column_names)}")
+                    
+                    if table_desc.foreign_keys:
+                        logger.info(f"        Foreign Keys:")
+                        for fk in table_desc.foreign_keys:
+                            logger.info(f"          - Column: {fk.column_name} -> {fk.referenced_table}.{fk.referenced_column} (Constraint: {fk.constraint_name})")
+                    
+                    if table_desc.indexes:
+                        logger.info(f"        Indexes:")
+                        for idx in table_desc.indexes:
+                            logger.info(f"          - Name: {idx.name}, Columns: {', '.join(idx.column_names)}, Unique: {idx.is_unique}, Primary: {idx.is_primary}, Type: {idx.type}")
+                    
+                    logger.info("-" * 40) # Separator for readability
+
+    except ConnectionError as e:
+        logger.error(f"MySQL connection failed: {e}")
+        logger.exception("Connection Error Traceback:")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"An unexpected error occurred: {e}")
+        logger.exception("Unexpected Error Traceback:")
+        sys.exit(1)
+    finally:
+        pass
+
+if __name__ == "__main__":
+    run_basic_mysql_test()

core/check_mysql_connection.py

@@ -0,0 +1,23 @@
+import sys
+from core.db_connector.connectors.mysql import MySQLConnector
+from core.db_connector.cache_manager import CacheManager
+from loguru import logger
+
+MYSQL_TEST_CONFIG = {
+    "host": "host.docker.internal",
+    "user": "root",
+    "password": "",
+    "port": 3306,
+}
+
+try:
+    logger.info("Attempting MySQL connection check...")
+    temp_conn = MySQLConnector(
+        connection_params=MYSQL_TEST_CONFIG,
+        cache_manager=CacheManager(),
+    )
+    logger.info("MySQL connection successful!")
+except Exception as e:
+    logger.error(f"MySQL connection failed: {e}")
+    logger.exception("Connection Error Traceback:")
+    sys.exit(1) # Exit with error code if connection fails

core/db_connector/__init__.py

@@ -0,0 +1,12 @@
+from .manager import ConnectorManager
+from .storage import BaseMetadataStore, FileMetadataStore, get_metadata_store
+from .ai_service import AIDocumentationService
+
+__all__ = [
+    "ConnectorManager",
+    "BaseMetadataStore",
+    "FileMetadataStore",
+    "get_metadata_store",
+    "AIDocumentationService",
+]
+

core/db_connector/ai_service.py

@@ -0,0 +1,133 @@
+from datetime import datetime, timezone
+import json
+import os
+import time
+from typing import Dict, Any, Optional
+from loguru import logger
+
+
+class AIDocumentationService:
+    """Service to generate basic AI documentation for database table schemas using LiteLLM."""
+
+    def __init__(self, model_name: Optional[str] = None):
+        self.model_name = model_name or os.getenv("LITELLM_MODEL") or "gpt-4o-mini"
+        self.timeout = int(os.getenv("LITELLM_TIMEOUT_SECONDS", "120"))
+        self.max_retries = int(os.getenv("LITELLM_MAX_RETRIES", "2"))
+        self.last_error: Optional[str] = None
+
+    def generate_table_documentation(
+        self, table_description: Dict[str, Any]
+    ) -> Optional[Dict[str, Any]]:
+        """Generates AI documentation for a table description payload using LiteLLM.
+
+        This method is non-blocking: if LiteLLM is not installed, if API keys are missing,
+        or if any API/network error occurs, it catches the exception, logs a warning,
+        and returns None without throwing an exception.
+        """
+        table_name = table_description.get("table_name", "unknown")
+        schema_name = table_description.get("schema_name", "unknown")
+        instance_name = table_description.get("instance_name", "unknown")
+        columns = table_description.get("columns", [])
+        self.last_error = None
+
+        prompt = (
+            f"Analyse the following database table schema and generate a concise business documentation.\n"
+            f"Instance Name: {instance_name}\n"
+            f"Schema Name: {schema_name}\n"
+            f"Table Name: {table_name}\n"
+            f"Columns: {json.dumps(columns, default=str)}\n"
+            f"Primary Key: {table_description.get('primary_key', [])}\n"
+            f"Foreign Keys: {json.dumps(table_description.get('foreign_keys', []), default=str)}\n\n"
+            f"Return a valid JSON object with the following keys:\n"
+            f"- 'summary': A high-level sentence describing the purpose of the table.\n"
+            f"- 'column_descriptions': A map of column_name -> brief functional description.\n"
+        )
+
+        try:
+            import litellm
+        except ImportError:
+            self.last_error = "litellm package is not installed"
+            logger.warning(f"AIDocumentationService: {self.last_error}. Skipping AI doc generation.")
+            return None
+
+        try:
+            logger.info(
+                f"AIDocumentationService: Generating AI docs for table '{table_name}' using model '{self.model_name}'"
+            )
+
+            kwargs = {}
+            if os.getenv("LITELLM_API_KEY"):
+                kwargs["api_key"] = os.getenv("LITELLM_API_KEY")
+            if os.getenv("LITELLM_API_BASE"):
+                kwargs["api_base"] = os.getenv("LITELLM_API_BASE")
+
+            messages = [
+                {
+                    "role": "system",
+                    "content": "You are an expert Data Architect. Produce clean, structured schema documentation in valid JSON format.",
+                },
+                {"role": "user", "content": prompt},
+            ]
+
+            last_attempt_error: Optional[Exception] = None
+            response = None
+            for attempt in range(self.max_retries):
+                try:
+                    response = litellm.completion(
+                        model=self.model_name,
+                        messages=messages,
+                        response_format={"type": "json_object"},
+                        timeout=self.timeout,
+                        **kwargs,
+                    )
+                    break
+                except Exception as attempt_exc:
+                    last_attempt_error = attempt_exc
+                    logger.warning(
+                        f"AIDocumentationService: attempt {attempt + 1}/{self.max_retries} failed for '{table_name}': {attempt_exc}"
+                    )
+                    if attempt < self.max_retries - 1:
+                        time.sleep(2)
+
+            if response is None:
+                raise last_attempt_error
+
+
+            content = response.choices[0].message.content
+            if not content:
+                self.last_error = f"Empty response content from model '{self.model_name}'"
+                logger.warning(f"AIDocumentationService: {self.last_error} for table '{table_name}'")
+                return None
+
+            # Strip markdown code fences if present (e.g. ```json ... ```)
+            stripped = content.strip()
+            if stripped.startswith("```"):
+                stripped = stripped.split("\n", 1)[1] if "\n" in stripped else stripped[3:]
+                if stripped.endswith("```"):
+                    stripped = stripped[:-3].strip()
+            else:
+                stripped = stripped
+
+            try:
+                parsed_doc = json.loads(stripped)
+            except json.JSONDecodeError:
+                self.last_error = f"Non-JSON response from model '{self.model_name}': {content!r:.200}"
+                logger.warning(f"AIDocumentationService: {self.last_error} for table '{table_name}'")
+                return None
+
+            logger.info(
+                f"AIDocumentationService: AI docs generated for table '{table_name}' using model '{self.model_name}'"
+            )
+
+            return {
+                "summary": parsed_doc.get("summary", ""),
+                "column_descriptions": parsed_doc.get("column_descriptions", {}),
+                "generated_at": datetime.now(timezone.utc).isoformat(),
+                "model": self.model_name,
+            }
+        except Exception as e:
+            self.last_error = str(e)
+            logger.warning(
+                f"AIDocumentationService: Failed or skipped AI doc generation for table '{table_name}': {e}"
+            )
+            return None

core/db_connector/cache_manager.py

@@ -0,0 +1,69 @@
+import redis
+import json
+from typing import Optional, Any
+from loguru import logger
+import os
+
+class CacheManager:
+    def __init__(
+        self,
+        host: str = 'localhost',
+        port: int = 6379,
+        db: int = 0,
+        ttl_seconds: int = 86400,
+        project_prefix: Optional[str] = None,
+        socket_connect_timeout: float = 2.0,
+        socket_timeout: float = 2.0,
+    ):
+        self.redis_client = None
+        self.ttl_seconds = ttl_seconds # Default 1 day
+        self.project_prefix = project_prefix if project_prefix is not None else os.getenv("CACHE_KEY_PREFIX", "multi-db-connector")
+        try:
+            self.redis_client = redis.Redis(
+                host=host,
+                port=port,
+                db=db,
+                decode_responses=True,
+                socket_connect_timeout=socket_connect_timeout,
+                socket_timeout=socket_timeout,
+                health_check_interval=30,
+            )
+            self.redis_client.ping()
+            logger.info(f"Successfully connected to Redis at {host}:{port}/{db}")
+        except redis.exceptions.RedisError as e:
+            logger.error(f"Could not connect to Redis cache at {host}:{port}/{db}: {e}")
+            self.redis_client = None # Ensure client is None if connection fails
+
+    def _serialize(self, data: Any) -> str:
+        """Serializes data to JSON string."""
+        return json.dumps(data)
+
+    def _deserialize(self, data_str: str) -> Any:
+        """Deserializes JSON string to data."""
+        return json.loads(data_str)
+
+    def get_cached_data(self, key: str, no_cache: bool = False) -> Optional[Any]:
+        prefixed_key = f"{self.project_prefix}:{key}" # Add prefix
+        if not self.redis_client or no_cache:
+            return None
+        try:
+            cached_data = self.redis_client.get(prefixed_key)
+            if cached_data:
+                logger.debug(f"Cache hit for key: {prefixed_key}")
+                return self._deserialize(cached_data)
+            logger.debug(f"Cache miss for key: {prefixed_key}")
+            return None
+        except Exception as e:
+            logger.error(f"Error retrieving data from Redis for key {prefixed_key}: {e}")
+            return None
+
+    def set_cached_data(self, key: str, data: Any, ttl: Optional[int] = None):
+        prefixed_key = f"{self.project_prefix}:{key}" # Add prefix
+        if not self.redis_client:
+            return
+        try:
+            ttl_to_use = ttl if ttl is not None else self.ttl_seconds
+            self.redis_client.setex(prefixed_key, ttl_to_use, self._serialize(data))
+            logger.debug(f"Data set in cache for key: {prefixed_key} with TTL: {ttl_to_use}s")
+        except Exception as e:
+            logger.error(f"Error setting data in Redis for key {prefixed_key}: {e}")

core/db_connector/config_service.py

@@ -0,0 +1,116 @@
+from typing import Dict, Any, List, Optional
+from loguru import logger
+
+from core.db_connector.manager import ConnectorManager
+from core.db_connector.models import Instance
+from core.db_connector.configurations import get_db_configurations
+
+
+class ConfigService:
+    def __init__(self, connector_manager: ConnectorManager, config_file: Optional[str] = None):
+        self.connector_manager = connector_manager
+        self.db_configurations: Dict[str, Dict[str, Any]] = get_db_configurations(config_file=config_file)
+        logger.info(f"ConfigService: Initialized with {len(self.db_configurations)} database configurations.")
+
+    def get_available_configurations(self) -> List[str]:
+        logger.info("ConfigService: Fetching available configurations.")
+        return list(self.db_configurations.keys())
+
+    def _get_connector_details(self, config_name: str) -> Dict[str, Any]:
+        config = self.db_configurations.get(config_name)
+        if not config:
+            logger.error(f"ConfigService: Configuration name '{config_name}' not found.")
+            raise ValueError(f"Configuration name '{config_name}' not found.")
+
+        connector_type = config.get("connector_type")
+        connection_params = config.get("connection_params")
+
+        if not connector_type or not connection_params:
+            raise ValueError(
+                f"ConfigService: Configuration '{config_name}' is missing "
+                "'connector_type' or 'connection_params'."
+            )
+        return {"connector_type": connector_type, "connection_params": connection_params}
+
+    def _get_hosts(self, config_name: str) -> List[str]:
+        """Return explicitly configured connection hosts, when present.
+
+        Configurations may use either a multi-host ``hosts`` collection or a
+        flat ``host`` connection parameter. Connectors such as Athena and
+        DynamoDB discover instances remotely and expose neither.
+        """
+        details = self._get_connector_details(config_name)
+        connection_params = details["connection_params"]
+        hosts = [item["host"] for item in connection_params.get("hosts", [])]
+        if hosts:
+            return hosts
+        return [connection_params["host"]] if connection_params.get("host") else []
+
+    def configuration_matches_instance(
+        self, config_name: str, instance_name: str, no_cache: bool = False
+    ) -> bool:
+        """Return whether an instance belongs to a database configuration."""
+        configured_hosts = self._get_hosts(config_name)
+        if configured_hosts:
+            return instance_name in configured_hosts
+        return instance_name in {
+            instance.name for instance in self.list_instances(config_name, no_cache=no_cache)
+        }
+
+    def resolve_configurations_for_instance(
+        self, instance_name: str, no_cache: bool = False
+    ) -> List[str]:
+        return [
+            config_name
+            for config_name in self.get_available_configurations()
+            if self.configuration_matches_instance(config_name, instance_name, no_cache)
+        ]
+
+    def _get_connector_for_host(self, config_name: str, host: str):
+        details = self._get_connector_details(config_name)
+        host_params = next(
+            (h for h in details["connection_params"].get("hosts", []) if h["host"] == host),
+            None,
+        )
+        if host_params:
+            return self.connector_manager.get_connector(details["connector_type"], host_params)
+
+        # Flat connection parameters identify one connection, while ``host`` is
+        # an instance/catalog/region discovered through that connection.
+        if not details["connection_params"].get("hosts"):
+            return self.connector_manager.get_connector(
+                details["connector_type"], details["connection_params"]
+            )
+        raise ValueError(f"Host '{host}' not found in config '{config_name}'")
+
+    def list_instances(self, config_name: str, no_cache: bool = False) -> List[Instance]:
+        """List instances for both multi-host and flat connector configurations."""
+        configured_hosts = self._get_hosts(config_name)
+        if configured_hosts:
+            instances: List[Instance] = []
+            for host in configured_hosts:
+                connector = self._get_connector_for_host(config_name, host)
+                instances.extend(connector.list_instances(no_cache=no_cache))
+            return instances
+
+        details = self._get_connector_details(config_name)
+        connector = self.connector_manager.get_connector(
+            details["connector_type"], details["connection_params"]
+        )
+        return connector.list_instances(no_cache=no_cache)
+
+    def resolve_instance_names(
+        self, config_name: str, instance_name: Optional[str] = None, no_cache: bool = False
+    ) -> List[str]:
+        if instance_name:
+            return [instance_name]
+        return [instance.name for instance in self.list_instances(config_name, no_cache)]
+
+    def test_connection(self, config_name: str):
+        logger.info(f"ConfigService: Attempting to test connection for config: {config_name}")
+        for instance in self.list_instances(config_name, no_cache=True):
+            logger.info(
+                f"ConfigService: Successfully tested connection to {instance.name} "
+                f"for config: {config_name}"
+            )
+        return {"message": f"Successfully connected to all hosts in {config_name}."}