maleo-database 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

maleo/database/config/__init__.py

@@ -1,5 +1,14 @@
+from elasticsearch import AsyncElasticsearch, Elasticsearch
+from motor.motor_asyncio import AsyncIOMotorClient
 from pydantic import BaseModel, Field
-from typing import Generic
+from pymongo import MongoClient
+from redis.asyncio import Redis as AsyncRedis
+from redis import Redis as SyncRedis
+from sqlalchemy.engine import create_engine as create_sync_engine, Engine
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
+from typing import Generic, Literal, TypeVar, Union, overload
+from maleo.types.base.dict import StringToAnyDict
+from ..enums import Connection
 from .additional import AdditionalConfigT, RedisAdditionalConfig
 from .connection import (
     ConnectionConfigT,
@@ -35,25 +44,61 @@ class BaseDatabaseConfig(
     additional: AdditionalConfigT = Field(..., description="Additional config")
 
 
-class PostgreSQLDatabaseConfig(
+class MySQLDatabaseConfig(
     BaseDatabaseConfig[
-        PostgreSQLConnectionConfig,
-        PostgreSQLPoolingConfig,
+        MySQLConnectionConfig,
+        MySQLPoolingConfig,
         None,
     ]
 ):
     additional: None = None
 
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return {
+            **self.connection.engine_kwargs,
+            **self.pooling.engine_kwargs,
+        }
 
-class MySQLDatabaseConfig(
+    @overload
+    def create_engine(self, connection: Literal[Connection.ASYNC]) -> AsyncEngine: ...
+    @overload
+    def create_engine(self, connection: Literal[Connection.SYNC]) -> Engine: ...
+    def create_engine(self, connection: Connection) -> Union[AsyncEngine, Engine]:
+        url = self.connection.make_url(connection)
+        if connection is Connection.ASYNC:
+            return create_async_engine(url, **self.engine_kwargs)
+        elif connection is Connection.SYNC:
+            return create_sync_engine(url, **self.engine_kwargs)
+
+
+class PostgreSQLDatabaseConfig(
     BaseDatabaseConfig[
-        MySQLConnectionConfig,
-        MySQLPoolingConfig,
+        PostgreSQLConnectionConfig,
+        PostgreSQLPoolingConfig,
         None,
     ]
 ):
     additional: None = None
 
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return {
+            **self.connection.engine_kwargs,
+            **self.pooling.engine_kwargs,
+        }
+
+    @overload
+    def create_engine(self, connection: Literal[Connection.ASYNC]) -> AsyncEngine: ...
+    @overload
+    def create_engine(self, connection: Literal[Connection.SYNC]) -> Engine: ...
+    def create_engine(self, connection: Connection) -> Union[AsyncEngine, Engine]:
+        url = self.connection.make_url(connection)
+        if connection is Connection.ASYNC:
+            return create_async_engine(url, **self.engine_kwargs)
+        elif connection is Connection.SYNC:
+            return create_sync_engine(url, **self.engine_kwargs)
+
 
 class SQLiteDatabaseConfig(
     BaseDatabaseConfig[
@@ -64,6 +109,24 @@ class SQLiteDatabaseConfig(
 ):
     additional: None = None
 
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return {
+            **self.connection.engine_kwargs,
+            **self.pooling.engine_kwargs,
+        }
+
+    @overload
+    def create_engine(self, connection: Literal[Connection.ASYNC]) -> AsyncEngine: ...
+    @overload
+    def create_engine(self, connection: Literal[Connection.SYNC]) -> Engine: ...
+    def create_engine(self, connection: Connection) -> Union[AsyncEngine, Engine]:
+        url = self.connection.make_url(connection)
+        if connection is Connection.ASYNC:
+            return create_async_engine(url, **self.engine_kwargs)
+        elif connection is Connection.SYNC:
+            return create_sync_engine(url, **self.engine_kwargs)
+
 
 class SQLServerDatabaseConfig(
     BaseDatabaseConfig[
@@ -74,6 +137,72 @@ class SQLServerDatabaseConfig(
 ):
     additional: None = None
 
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return {
+            **self.connection.engine_kwargs,
+            **self.pooling.engine_kwargs,
+        }
+
+    @overload
+    def create_engine(self, connection: Literal[Connection.ASYNC]) -> AsyncEngine: ...
+    @overload
+    def create_engine(self, connection: Literal[Connection.SYNC]) -> Engine: ...
+    def create_engine(self, connection: Connection) -> Union[AsyncEngine, Engine]:
+        url = self.connection.make_url(connection)
+        if connection is Connection.ASYNC:
+            return create_async_engine(url, **self.engine_kwargs)
+        elif connection is Connection.SYNC:
+            return create_sync_engine(url, **self.engine_kwargs)
+
+
+SQLConfigT = TypeVar(
+    "SQLConfigT",
+    PostgreSQLDatabaseConfig,
+    MySQLDatabaseConfig,
+    SQLiteDatabaseConfig,
+    SQLServerDatabaseConfig,
+)
+
+
+class ElasticsearchDatabaseConfig(
+    BaseDatabaseConfig[
+        ElasticsearchConnectionConfig,
+        ElasticsearchPoolingConfig,
+        None,
+    ]
+):
+    additional: None = None
+
+    @property
+    def client_kwargs(self) -> StringToAnyDict:
+        client_kwargs = {}
+
+        if self.connection.username and self.connection.password:
+            client_kwargs["http_auth"] = (
+                self.connection.username,
+                self.connection.password,
+            )
+
+        client_kwargs.update(self.pooling.client_kwargs)
+
+        return client_kwargs
+
+    @overload
+    def create_client(
+        self, connection: Literal[Connection.ASYNC]
+    ) -> AsyncElasticsearch: ...
+    @overload
+    def create_client(self, connection: Literal[Connection.SYNC]) -> Elasticsearch: ...
+    def create_client(
+        self, connection: Connection
+    ) -> Union[AsyncElasticsearch, Elasticsearch]:
+        hosts = [{"host": self.connection.host, "port": self.connection.port}]
+        if connection is Connection.ASYNC:
+            return AsyncElasticsearch(hosts, **self.client_kwargs)
+        else:
+            return Elasticsearch(hosts, **self.client_kwargs)
+
 
 class MongoDBDatabaseConfig(
     BaseDatabaseConfig[
@@ -84,6 +213,25 @@ class MongoDBDatabaseConfig(
 ):
     additional: None = None
 
+    @property
+    def client_kwargs(self) -> StringToAnyDict:
+        return self.pooling.client_kwargs
+
+    @overload
+    def create_client(
+        self, connection: Literal[Connection.ASYNC]
+    ) -> AsyncIOMotorClient: ...
+    @overload
+    def create_client(self, connection: Literal[Connection.SYNC]) -> MongoClient: ...
+    def create_client(
+        self, connection: Connection
+    ) -> Union[AsyncIOMotorClient, MongoClient]:
+        url = self.connection.make_url(connection)
+        if connection is Connection.ASYNC:
+            return AsyncIOMotorClient(url, **self.client_kwargs)
+        else:
+            return MongoClient(url, **self.client_kwargs)
+
 
 class RedisDatabaseConfig(
     BaseDatabaseConfig[
@@ -94,12 +242,37 @@ class RedisDatabaseConfig(
 ):
     additional: RedisAdditionalConfig = Field(..., description="Additional config")
 
+    @property
+    def client_kwargs(self) -> StringToAnyDict:
+        return self.pooling.client_kwargs
 
-class ElasticsearchDatabaseConfig(
-    BaseDatabaseConfig[
-        ElasticsearchConnectionConfig,
-        ElasticsearchPoolingConfig,
-        None,
-    ]
-):
-    additional: None = None
+    @overload
+    def create_client(self, connection: Literal[Connection.ASYNC]) -> AsyncRedis: ...
+    @overload
+    def create_client(self, connection: Literal[Connection.SYNC]) -> SyncRedis: ...
+    def create_client(self, connection: Connection) -> Union[AsyncRedis, SyncRedis]:
+        url = self.connection.make_url(connection)
+        if connection is Connection.ASYNC:
+            return AsyncRedis.from_url(url, **self.client_kwargs)
+        else:
+            return SyncRedis.from_url(url, **self.client_kwargs)
+
+
+NoSQLConfigT = TypeVar(
+    "NoSQLConfigT",
+    ElasticsearchDatabaseConfig,
+    MongoDBDatabaseConfig,
+    RedisDatabaseConfig,
+)
+
+
+ConfigT = TypeVar(
+    "ConfigT",
+    PostgreSQLDatabaseConfig,
+    MySQLDatabaseConfig,
+    SQLiteDatabaseConfig,
+    SQLServerDatabaseConfig,
+    MongoDBDatabaseConfig,
+    RedisDatabaseConfig,
+    ElasticsearchDatabaseConfig,
+)

maleo/database/config/connection.py

@@ -241,6 +241,41 @@ class BaseConnectionConfig(
 ConnectionConfigT = TypeVar("ConnectionConfigT", bound=BaseConnectionConfig)
 
 
+class MySQLConnectionConfig(
+    BaseConnectionConfig[Literal[Driver.MYSQL], str, str, str, int, str]
+):
+    driver: Literal[Driver.MYSQL] = Driver.MYSQL
+    port: int = Field(3306, description="MySQL port")
+    username: str = Field("root", description="MySQL username")
+
+    # MySQL-specific options
+    echo: bool = Field(False, description="Enable SQL statement logging")
+    charset: MySQLCharset = Field(MySQLCharset.UTF8MB4, description="Character set")
+
+    def _make_mysql_url(self, connection: Connection = Connection.ASYNC) -> str:
+        """MySQL database URL format with proper async/sync drivers."""
+        # Choose appropriate driver based on connection type
+        if connection is Connection.ASYNC:
+            driver_name = "mysql+aiomysql"
+        elif connection is Connection.SYNC:
+            driver_name = "mysql+pymysql"  # or mysqlclient
+        else:
+            driver_name = "mysql"  # Default
+
+        encoded_user = self._safe_encode_credential(self.username)
+        encoded_pass = self._safe_encode_credential(self.password)
+
+        base_url = f"{driver_name}://{encoded_user}:{encoded_pass}@{self.host}:{self.port}/{self.database}"
+
+        # Add MySQL-specific options
+        mysql_options = {"charset": self.charset}
+        return base_url + self._make_options_string(mysql_options)
+
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return self.model_dump(include={"echo"}, exclude_none=True)
+
+
 class PostgreSQLConnectionConfig(
     BaseConnectionConfig[Literal[Driver.POSTGRESQL], str, str, str, int, str]
 ):
@@ -279,36 +314,9 @@ class PostgreSQLConnectionConfig(
 
         return base_url + self._make_options_string(pg_options)
 
-
-class MySQLConnectionConfig(
-    BaseConnectionConfig[Literal[Driver.MYSQL], str, str, str, int, str]
-):
-    driver: Literal[Driver.MYSQL] = Driver.MYSQL
-    port: int = Field(3306, description="MySQL port")
-    username: str = Field("root", description="MySQL username")
-
-    # MySQL-specific options
-    echo: bool = Field(False, description="Enable SQL statement logging")
-    charset: MySQLCharset = Field(MySQLCharset.UTF8MB4, description="Character set")
-
-    def _make_mysql_url(self, connection: Connection = Connection.ASYNC) -> str:
-        """MySQL database URL format with proper async/sync drivers."""
-        # Choose appropriate driver based on connection type
-        if connection is Connection.ASYNC:
-            driver_name = "mysql+aiomysql"
-        elif connection is Connection.SYNC:
-            driver_name = "mysql+pymysql"  # or mysqlclient
-        else:
-            driver_name = "mysql"  # Default
-
-        encoded_user = self._safe_encode_credential(self.username)
-        encoded_pass = self._safe_encode_credential(self.password)
-
-        base_url = f"{driver_name}://{encoded_user}:{encoded_pass}@{self.host}:{self.port}/{self.database}"
-
-        # Add MySQL-specific options
-        mysql_options = {"charset": self.charset}
-        return base_url + self._make_options_string(mysql_options)
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return self.model_dump(include={"echo"}, exclude_none=True)
 
 
 class SQLiteConnectionConfig(
@@ -339,6 +347,10 @@ class SQLiteConnectionConfig(
         base_url = f"{driver_name}:///{self.database}"
         return base_url + self._make_options_string()
 
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return self.model_dump(include={"echo"}, exclude_none=True)
+
 
 class SQLServerConnectionConfig(
     BaseConnectionConfig[Literal[Driver.MSSQL], str, str, str, int, str]
@@ -392,6 +404,10 @@ class SQLServerConnectionConfig(
         }
         return base_url + self._make_options_string(mssql_options)
 
+    @property
+    def engine_kwargs(self) -> StringToAnyDict:
+        return self.model_dump(include={"echo"}, exclude_none=True)
+
 
 class MongoDBConnectionConfig(
     BaseConnectionConfig[