1.1.0

Author: krylosov-aa

Makefile

@@ -10,3 +10,9 @@ test:
 uv:
 	uv sync
 	source .venv/bin/activate
+
+build:
+	uv build
+
+publish:
+	uv publish

context_async_sqlalchemy/__init__.py

@@ -9,7 +9,8 @@
 )
 from .connect import (
     DBConnect,
-    db_connect,
+    master_connect,
+    replica_connect,
 )
 from .session import (
     db_session,
@@ -19,6 +20,8 @@
     commit_db_session,
     rollback_db_session,
     close_db_session,
+    new_non_ctx_atomic_session,
+    new_non_ctx_session,
 )
 from .auto_commit import auto_commit_by_status_code
 from .fastapi_utils.middleware import fastapi_db_session_middleware
@@ -32,7 +35,8 @@
     "pop_db_session_from_context",
     "run_in_new_ctx",
     "DBConnect",
-    "db_connect",
+    "master_connect",
+    "replica_connect",
     "db_session",
     "atomic_db_session",
     "run_with_new_db_session",
@@ -42,4 +46,6 @@
     "close_db_session",
     "auto_commit_by_status_code",
     "fastapi_db_session_middleware",
+    "new_non_ctx_atomic_session",
+    "new_non_ctx_session",
 ]

context_async_sqlalchemy/connect.py

@@ -1,35 +1,74 @@
 import asyncio
+from typing import Any, Callable, Coroutine
+
 from sqlalchemy.ext.asyncio import (
     async_sessionmaker,
     AsyncEngine,
     AsyncSession,
 )
 
+EngineCreatorFunc = Callable[[str], AsyncEngine]
+SessionMakerCreatorFunc = Callable[
+    [AsyncEngine], async_sessionmaker[AsyncSession]
+]
+AsyncFunc = Callable[["DBConnect"], Coroutine[Any, Any, None]]
+
 
 class DBConnect:
     """stores the database connection parameters"""
 
     def __init__(self) -> None:
-        self.engine: AsyncEngine | None = None
-        self.session_maker: async_sessionmaker[AsyncSession] | None = None
+        self.host: str | None = None
+        self._engine: AsyncEngine | None = None
+        self._session_maker: async_sessionmaker[AsyncSession] | None = None
+
+        self.engine_creator: EngineCreatorFunc | None = None
+        self.session_maker_creator: SessionMakerCreatorFunc | None = None
+        self.before_create_session_handler: AsyncFunc | None = None
         self._lock = asyncio.Lock()
 
-    async def close(self) -> None:
-        if self.engine:
-            await self.engine.dispose()
-        self.engine = None
+    async def connect(self, host: str) -> None:
+        """initiates engine and session maker"""
+        assert host
+        async with self._lock:
+            await self._connect(host)
 
+    async def change_host(self, host: str) -> None:
+        """Renews the connection if a host needs to be changed"""
+        assert host
+        async with self._lock:
+            if host != self.host:
+                await self._connect(host)
 
-db_connect = DBConnect()
+    async def create_session(self) -> AsyncSession:
+        """Creates a new session"""
+        if self.before_create_session_handler:
+            await self.before_create_session_handler(self)
+        maker = await self.get_session_maker()
+        return maker()
 
+    async def get_session_maker(self) -> async_sessionmaker[AsyncSession]:
+        """Gets the session maker"""
+        if not self._session_maker:
+            assert self.host
+            await self.connect(self.host)
+
+        assert self._session_maker
+        return self._session_maker
+
+    async def close(self) -> None:
+        if self._engine:
+            await self._engine.dispose()
+        self._engine = None
 
-def get_session_maker() -> async_sessionmaker[AsyncSession]:
-    """Gets the session maker"""
-    assert db_connect.session_maker
-    return db_connect.session_maker
+    async def _connect(self, host: str) -> None:
+        self.host = host
+        await self.close()
+        assert self.engine_creator
+        self._engine = self.engine_creator(host)
+        assert self.session_maker_creator
+        self._session_maker = self.session_maker_creator(self._engine)
 
 
-def create_session() -> AsyncSession:
-    """Creates a new session"""
-    assert db_connect.session_maker
-    return db_connect.session_maker()
+master_connect = DBConnect()
+replica_connect = DBConnect()

context_async_sqlalchemy/context.py

@@ -3,7 +3,7 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from .connect import create_session
+from .connect import master_connect
 
 
 def init_db_session_ctx() -> Token[dict[str, AsyncSession] | None]:
@@ -104,6 +104,7 @@ async def _new_ctx_wrapper(
     **kwargs: Any,
 ) -> AsyncCallableResult:
     _init_db_session_ctx()
-    async with create_session() as session:
+    session_maker = await master_connect.get_session_maker()
+    async with session_maker() as session:
         put_db_session_to_context(session)
         return await callable_func(*args, **kwargs)

context_async_sqlalchemy/session.py

@@ -3,7 +3,7 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from .connect import create_session
+from .connect import master_connect
 from .context import (
     AsyncCallable,
     AsyncCallableResult,
@@ -24,7 +24,7 @@ async def db_session() -> AsyncSession:
     """
     session = get_db_session_from_context()
     if not session:
-        session = create_session()
+        session = await master_connect.create_session()
         put_db_session_to_context(session)
     return session
 
@@ -128,6 +128,22 @@ async def close_db_session() -> None:
         await session.close()
 
 
+@asynccontextmanager
+async def new_non_ctx_session() -> AsyncGenerator[AsyncSession, None]:
+    """Creating a new session without using a context"""
+    session_maker = await master_connect.get_session_maker()
+    async with session_maker() as session:
+        yield session
+
+
+@asynccontextmanager
+async def new_non_ctx_atomic_session() -> AsyncGenerator[AsyncSession, None]:
+    """Creating a new session with transaction without using a context"""
+    async with new_non_ctx_session() as session:
+        async with session.begin():
+            yield session
+
+
 async def _atomic_wrapper(
     callable_func: AsyncCallable[AsyncCallableResult],
     *args: Any,

exmaples/fastapi_example/database.py

@@ -10,7 +10,7 @@
 )
 
 
-def create_engine() -> AsyncEngine:
+def create_engine(host: str) -> AsyncEngine:
     """
     database connection parameters.
     In production code, you will probably take these parameters from
@@ -19,7 +19,6 @@ def create_engine() -> AsyncEngine:
     pg_user = "krylosov-aa"
     pg_password = ""
     pg_port = 6432
-    host = "localhost"
     pg_db = "test"
     return create_async_engine(
         f"postgresql+asyncpg://"

exmaples/fastapi_example/setup_app.py

@@ -6,7 +6,11 @@
 from fastapi import FastAPI
 from starlette.middleware.base import BaseHTTPMiddleware
 
-from context_async_sqlalchemy import db_connect, fastapi_db_session_middleware
+from context_async_sqlalchemy import (
+    master_connect,
+    replica_connect,
+    fastapi_db_session_middleware,
+)
 
 from .database import create_engine, create_session_maker
 from .routes.atomic_usage import handler_with_db_session_and_atomic
@@ -32,19 +36,17 @@ def setup_app() -> FastAPI:
     )
     setup_middlewares(app)
     setup_routes(app)
-    setup_database()
     return app
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]:
-    """
-    It is important to clean up resources at the end of an application's
-        life.
-    """
+    """Database connection lifecycle management"""
+    await setup_database()
     yield
     await asyncio.gather(
-        db_connect.close(),  # Close the engine if it was open
+        master_connect.close(),  # Close the engine if it was open
+        replica_connect.close(),  # Close the engine if it was open
     )
 
 
@@ -58,13 +60,14 @@ def setup_middlewares(app: FastAPI) -> None:
     )
 
 
-def setup_database() -> None:
+async def setup_database() -> None:
     """
     Here you pass the database connection parameters to the library.
     More specifically, the engine and session maker.
     """
-    db_connect.engine = create_engine()
-    db_connect.session_maker = create_session_maker(db_connect.engine)
+    master_connect.engine_creator = create_engine
+    master_connect.session_maker_creator = create_session_maker
+    await master_connect.connect("127.0.0.1")
 
 
 def setup_routes(app: FastAPI) -> None:

exmaples/fastapi_example/tests/conftest.py

@@ -2,28 +2,30 @@
 Basic settings and fixtures for testing
 """
 
-from typing import AsyncGenerator, Generator
+from typing import AsyncGenerator
 
 import pytest_asyncio
-import pytest
 from fastapi import FastAPI
 from httpx import AsyncClient, ASGITransport
 from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
 
+from context_async_sqlalchemy import master_connect, replica_connect
 from exmaples.fastapi_example.database import (
     create_engine,
     create_session_maker,
 )
-from exmaples.fastapi_example.setup_app import setup_app
+from exmaples.fastapi_example.setup_app import lifespan, setup_app
 
 
-@pytest.fixture
-def app() -> Generator[FastAPI]:
+@pytest_asyncio.fixture
+async def app() -> AsyncGenerator[FastAPI]:
     """
     A new application for each test allows for complete isolation between
         tests.
     """
-    yield setup_app()
+    app = setup_app()
+    async with lifespan(app):
+        yield app
 
 
 @pytest_asyncio.fixture
@@ -51,7 +53,14 @@ async def db_session_test(
 async def session_maker_test() -> AsyncGenerator[
     async_sessionmaker[AsyncSession]
 ]:
-    engine = create_engine()
+    engine = create_engine("127.0.0.1")
     session_maker = create_session_maker(engine)
     yield session_maker
     await engine.dispose()
+
+
+@pytest_asyncio.fixture(autouse=True)
+async def close_connect() -> AsyncGenerator[None]:
+    yield
+    await master_connect.close()
+    await replica_connect.close()

exmaples/fastapi_example/tests/non_transactional/conftest.py

@@ -52,7 +52,7 @@ async def cleanup_tables_before() -> None:
 
 
 async def _cleanup_tables() -> None:
-    engine = create_engine()
+    engine = create_engine("127.0.0.1")
     session_maker = create_session_maker(engine)
     async with session_maker() as session:
         await session.execute(

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "context-async-sqlalchemy"
-version = "1.0.1"
+version = "1.1.0"
 description = "A convenient way to configure and interact with a async sqlalchemy session through context in asynchronous applications"
 readme = "README.md"
 authors = [