memorisdk 1.0.2__py3-none-any.whl → 2.0.1__py3-none-any.whl

memori/database/auto_creator.py

@@ -0,0 +1,320 @@
+"""
+Database Auto-Creation System
+
+This module automatically creates databases if they don't exist, supporting
+PostgreSQL and MySQL with proper error handling and security validation.
+"""
+
+import ssl
+from typing import Dict
+from urllib.parse import parse_qs, urlparse
+
+from loguru import logger
+from sqlalchemy import create_engine, text
+from sqlalchemy.exc import OperationalError, ProgrammingError
+
+from .connection_utils import DatabaseConnectionUtils
+
+
+class DatabaseAutoCreator:
+    """Handles automatic database creation for PostgreSQL and MySQL"""
+
+    def __init__(self, schema_init: bool = True):
+        """
+        Initialize database auto-creator.
+
+        Args:
+            schema_init: Whether to enable automatic database creation
+        """
+        self.schema_init = schema_init
+        self.utils = DatabaseConnectionUtils()
+
+    def ensure_database_exists(self, connection_string: str) -> str:
+        """
+        Ensure target database exists, creating it if necessary.
+
+        Args:
+            connection_string: Original database connection string
+
+        Returns:
+            Connection string to use (may be unchanged if creation not needed)
+
+        Raises:
+            DatabaseCreationError: If database creation fails
+        """
+        if not self.schema_init:
+            logger.debug("Auto-creation disabled, using original connection string")
+            return connection_string
+
+        try:
+            # Parse connection string
+            components = self.utils.parse_connection_string(connection_string)
+
+            # SQLite doesn't need database creation
+            if not components["needs_creation"]:
+                logger.debug(
+                    f"Database engine {components['engine']} auto-creates, no action needed"
+                )
+                return connection_string
+
+            # Validate database name
+            if not self.utils.validate_database_name(components["database"]):
+                raise ValueError(f"Invalid database name: {components['database']}")
+
+            # Check if database exists
+            if self._database_exists(components):
+                logger.debug(f"Database '{components['database']}' already exists")
+                return connection_string
+
+            # Create database
+            self._create_database(components)
+            logger.info(f"Successfully created database '{components['database']}'")
+            return connection_string
+
+        except Exception as e:
+            logger.error(f"Database auto-creation failed: {e}")
+            # Don't raise exception - let the original connection attempt proceed
+            # This allows graceful degradation if user has manual setup
+            return connection_string
+
+    def _database_exists(self, components: Dict[str, str]) -> bool:
+        """Check if target database exists."""
+        try:
+            engine = components["engine"]
+
+            if engine == "postgresql":
+                return self._postgresql_database_exists(components)
+            elif engine == "mysql":
+                return self._mysql_database_exists(components)
+            else:
+                logger.warning(f"Database existence check not supported for {engine}")
+                return False
+
+        except Exception as e:
+            logger.error(f"Failed to check database existence: {e}")
+            return False
+
+    def _postgresql_database_exists(self, components: Dict[str, str]) -> bool:
+        """Check if PostgreSQL database exists."""
+        try:
+            # Connect to postgres system database
+            engine = create_engine(components["default_url"])
+
+            with engine.connect() as conn:
+                result = conn.execute(
+                    text("SELECT 1 FROM pg_database WHERE datname = :dbname"),
+                    {"dbname": components["database"]},
+                )
+                exists = result.fetchone() is not None
+
+            engine.dispose()
+            return exists
+
+        except Exception as e:
+            logger.error(f"PostgreSQL database existence check failed: {e}")
+            return False
+
+    def _get_mysql_connect_args(self, original_url: str) -> Dict:
+        """Get MySQL connection arguments with SSL support for system database connections."""
+        connect_args = {"charset": "utf8mb4"}
+
+        # Parse original URL for SSL parameters
+        parsed = urlparse(original_url)
+        if parsed.query:
+            query_params = parse_qs(parsed.query)
+
+            # Handle SSL parameters for PyMySQL - same logic as sqlalchemy_manager
+            if any(key in query_params for key in ["ssl", "ssl_disabled"]):
+                if query_params.get("ssl", ["false"])[0].lower() == "true":
+                    # Enable SSL with secure configuration for required secure transport
+                    connect_args["ssl"] = {
+                        "ssl_disabled": False,
+                        "check_hostname": False,
+                        "verify_mode": ssl.CERT_NONE,
+                    }
+                    # Also add ssl_disabled=False for PyMySQL
+                    connect_args["ssl_disabled"] = False
+                elif query_params.get("ssl_disabled", ["true"])[0].lower() == "false":
+                    # Enable SSL with secure configuration for required secure transport
+                    connect_args["ssl"] = {
+                        "ssl_disabled": False,
+                        "check_hostname": False,
+                        "verify_mode": ssl.CERT_NONE,
+                    }
+                    # Also add ssl_disabled=False for PyMySQL
+                    connect_args["ssl_disabled"] = False
+
+        return connect_args
+
+    def _mysql_database_exists(self, components: Dict[str, str]) -> bool:
+        """Check if MySQL database exists."""
+        try:
+            # Connect to mysql system database with SSL support
+            connect_args = self._get_mysql_connect_args(components["original_url"])
+            engine = create_engine(components["default_url"], connect_args=connect_args)
+
+            with engine.connect() as conn:
+                result = conn.execute(
+                    text(
+                        "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = :dbname"
+                    ),
+                    {"dbname": components["database"]},
+                )
+                exists = result.fetchone() is not None
+
+            engine.dispose()
+            return exists
+
+        except ModuleNotFoundError as e:
+            if "mysql" in str(e).lower():
+                logger.error(f"MySQL database existence check failed: {e}")
+                error_msg = (
+                    "❌ MySQL driver not found for database existence check. Install one of:\n"
+                    "- pip install mysql-connector-python\n"
+                    "- pip install PyMySQL\n"
+                    "- pip install memorisdk[mysql]"
+                )
+                logger.error(error_msg)
+            return False
+        except Exception as e:
+            logger.error(f"MySQL database existence check failed: {e}")
+            return False
+
+    def _create_database(self, components: Dict[str, str]) -> None:
+        """Create the target database."""
+        engine = components["engine"]
+
+        if engine == "postgresql":
+            self._create_postgresql_database(components)
+        elif engine == "mysql":
+            self._create_mysql_database(components)
+        else:
+            raise ValueError(f"Database creation not supported for {engine}")
+
+    def _create_postgresql_database(self, components: Dict[str, str]) -> None:
+        """Create PostgreSQL database."""
+        try:
+            logger.info(f"Creating PostgreSQL database '{components['database']}'...")
+
+            # Connect to postgres system database
+            engine = create_engine(components["default_url"])
+
+            # PostgreSQL requires autocommit for CREATE DATABASE
+            with engine.connect() as conn:
+                # Set autocommit mode
+                conn = conn.execution_options(isolation_level="AUTOCOMMIT")
+
+                # Create database (can't use parameters for database name)
+                # Database name is already validated, so this is safe
+                conn.execute(text(f'CREATE DATABASE "{components["database"]}"'))
+
+            engine.dispose()
+            logger.info(
+                f"PostgreSQL database '{components['database']}' created successfully"
+            )
+
+        except (OperationalError, ProgrammingError) as e:
+            error_msg = str(e)
+            if "already exists" in error_msg.lower():
+                logger.info(
+                    f"PostgreSQL database '{components['database']}' already exists"
+                )
+                return
+            elif "permission denied" in error_msg.lower():
+                raise PermissionError(
+                    f"Insufficient permissions to create database '{components['database']}'"
+                )
+            else:
+                raise RuntimeError(f"Failed to create PostgreSQL database: {e}")
+
+        except Exception as e:
+            raise RuntimeError(f"Unexpected error creating PostgreSQL database: {e}")
+
+    def _create_mysql_database(self, components: Dict[str, str]) -> None:
+        """Create MySQL database."""
+        try:
+            logger.info(f"Creating MySQL database '{components['database']}'...")
+
+            # Connect to mysql system database with SSL support
+            connect_args = self._get_mysql_connect_args(components["original_url"])
+            engine = create_engine(components["default_url"], connect_args=connect_args)
+
+            with engine.connect() as conn:
+                # Create database (can't use parameters for database name)
+                # Database name is already validated, so this is safe
+                conn.execute(
+                    text(
+                        f'CREATE DATABASE `{components["database"]}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci'
+                    )
+                )
+                conn.commit()
+
+            engine.dispose()
+            logger.info(
+                f"MySQL database '{components['database']}' created successfully"
+            )
+
+        except ModuleNotFoundError as e:
+            if "mysql" in str(e).lower():
+                error_msg = (
+                    "❌ MySQL driver not found for database creation. Install one of:\n"
+                    "- pip install mysql-connector-python\n"
+                    "- pip install PyMySQL\n"
+                    "- pip install memorisdk[mysql]"
+                )
+                logger.error(error_msg)
+                raise RuntimeError(error_msg)
+            else:
+                raise RuntimeError(f"Missing dependency for database creation: {e}")
+        except (OperationalError, ProgrammingError) as e:
+            error_msg = str(e)
+            if "database exists" in error_msg.lower():
+                logger.info(f"MySQL database '{components['database']}' already exists")
+                return
+            elif "access denied" in error_msg.lower():
+                raise PermissionError(
+                    f"Insufficient permissions to create database '{components['database']}'"
+                )
+            else:
+                raise RuntimeError(f"Failed to create MySQL database: {e}")
+
+        except Exception as e:
+            raise RuntimeError(f"Unexpected error creating MySQL database: {e}")
+
+    def get_database_info(self, connection_string: str) -> Dict[str, str]:
+        """
+        Get detailed information about database from connection string.
+
+        Args:
+            connection_string: Database connection URL
+
+        Returns:
+            Dictionary with database information
+        """
+        try:
+            components = self.utils.parse_connection_string(connection_string)
+
+            info = {
+                "engine": components["engine"],
+                "database": components["database"],
+                "host": components["host"],
+                "port": components["port"],
+                "needs_creation": components["needs_creation"],
+                "auto_creation_enabled": self.schema_init,
+            }
+
+            # Add existence check if auto-creation is enabled
+            if self.schema_init and components["needs_creation"]:
+                info["exists"] = self._database_exists(components)
+
+            return info
+
+        except Exception as e:
+            logger.error(f"Failed to get database info: {e}")
+            return {"error": str(e)}
+
+
+class DatabaseCreationError(Exception):
+    """Raised when database creation fails"""
+
+    pass

memori/database/connection_utils.py

@@ -0,0 +1,207 @@
+"""
+Database Connection Utilities for Auto-Creation System
+
+This module handles parsing connection strings, creating databases automatically,
+and managing multi-database scenarios for memori instances.
+"""
+
+import re
+from typing import Dict, Tuple
+from urllib.parse import urlparse
+
+from loguru import logger
+
+
+class DatabaseConnectionUtils:
+    """Utilities for parsing and managing database connections"""
+
+    # Default system databases for each engine
+    DEFAULT_DATABASES = {
+        "postgresql": "postgres",
+        "mysql": "mysql",
+        "sqlite": None,  # SQLite doesn't need default DB
+    }
+
+    @classmethod
+    def parse_connection_string(cls, connection_string: str) -> Dict[str, str]:
+        """
+        Parse database connection string and extract components.
+
+        Args:
+            connection_string: Database connection URL
+
+        Returns:
+            Dictionary with parsed components
+
+        Example:
+            mysql://root:pass@localhost:3306/memori_dev
+            -> {
+                'engine': 'mysql',
+                'user': 'root',
+                'password': 'pass',
+                'host': 'localhost',
+                'port': 3306,
+                'database': 'memori_dev',
+                'base_url': 'mysql://root:pass@localhost:3306',
+                'default_url': 'mysql://root:pass@localhost:3306/mysql'
+            }
+        """
+        try:
+            parsed = urlparse(connection_string)
+
+            # Extract components
+            engine = parsed.scheme.split("+")[
+                0
+            ].lower()  # Handle postgresql+psycopg2, mysql+pymysql
+            driver = parsed.scheme.split("+")[1] if "+" in parsed.scheme else None
+            user = parsed.username or ""
+            password = parsed.password or ""
+            host = parsed.hostname or "localhost"
+            port = parsed.port
+            database = parsed.path.lstrip("/") if parsed.path else ""
+
+            # Build base URL without database, preserving the driver
+            auth = f"{user}:{password}@" if user or password else ""
+            if user and not password:
+                auth = f"{user}@"
+
+            port_str = f":{port}" if port else ""
+            scheme = f"{engine}+{driver}" if driver else engine
+            base_url = f"{scheme}://{auth}{host}{port_str}"
+
+            # Create default database URL for system operations
+            default_db = cls.DEFAULT_DATABASES.get(engine)
+
+            # Preserve query parameters (especially SSL settings) for system database connections
+            query_string = f"?{parsed.query}" if parsed.query else ""
+            default_url = (
+                f"{base_url}/{default_db}{query_string}"
+                if default_db
+                else f"{base_url}{query_string}"
+            )
+
+            return {
+                "engine": engine,
+                "user": user,
+                "password": password,
+                "host": host,
+                "port": port,
+                "database": database,
+                "base_url": base_url,
+                "default_url": default_url,
+                "original_url": connection_string,
+                "needs_creation": engine
+                in ["postgresql", "mysql"],  # SQLite auto-creates
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to parse connection string: {e}")
+            raise ValueError(f"Invalid connection string format: {connection_string}")
+
+    @classmethod
+    def build_connection_string(
+        cls, components: Dict[str, str], target_database: str
+    ) -> str:
+        """
+        Build connection string with specific database name.
+
+        Args:
+            components: Parsed connection components
+            target_database: Database name to connect to
+
+        Returns:
+            Complete connection string
+        """
+        return f"{components['base_url']}/{target_database}"
+
+    @classmethod
+    def validate_database_name(cls, database_name: str) -> bool:
+        """
+        Validate database name for security and compatibility.
+
+        Args:
+            database_name: Name to validate
+
+        Returns:
+            True if valid, False otherwise
+        """
+        if not database_name:
+            return False
+
+        # Basic SQL injection prevention
+        if any(
+            char in database_name.lower()
+            for char in [";", "'", '"', "\\", "/", "*", "?", "<", ">", "|"]
+        ):
+            return False
+
+        # Check length (most databases have limits)
+        if len(database_name) > 64:  # MySQL limit
+            return False
+
+        # Must start with letter or underscore, contain only alphanumeric, underscore, hyphen
+        if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_-]*$", database_name):
+            return False
+
+        # Reserved words check (only for database creation, not connection)
+        # Note: 'postgres' is a valid system database to connect to in PostgreSQL
+        reserved_words = ["mysql", "information_schema", "performance_schema", "sys"]
+        if database_name.lower() in reserved_words:
+            return False
+
+        return True
+
+    @classmethod
+    def generate_database_name(
+        cls, base_name: str = "memori", suffix: str = None, prefix: str = None
+    ) -> str:
+        """
+        Generate a valid database name with optional prefix/suffix.
+
+        Args:
+            base_name: Base database name
+            suffix: Optional suffix (e.g., "dev", "prod", "test")
+            prefix: Optional prefix (e.g., "company", "project")
+
+        Returns:
+            Generated database name
+
+        Examples:
+            generate_database_name() -> "memori"
+            generate_database_name(suffix="dev") -> "memori_dev"
+            generate_database_name(prefix="acme", suffix="prod") -> "acme_memori_prod"
+        """
+        parts = []
+
+        if prefix:
+            parts.append(prefix)
+
+        parts.append(base_name)
+
+        if suffix:
+            parts.append(suffix)
+
+        database_name = "_".join(parts)
+
+        if not cls.validate_database_name(database_name):
+            raise ValueError(f"Generated database name is invalid: {database_name}")
+
+        return database_name
+
+    @classmethod
+    def extract_database_info(cls, connection_string: str) -> Tuple[str, str, bool]:
+        """
+        Extract database engine, name, and creation requirement.
+
+        Args:
+            connection_string: Database connection URL
+
+        Returns:
+            Tuple of (engine, database_name, needs_creation)
+        """
+        components = cls.parse_connection_string(connection_string)
+        return (
+            components["engine"],
+            components["database"],
+            components["needs_creation"],
+        )