pg

Postgres database clients for sync and async operations.

AsyncPgDb

Bases: AsyncDbBase

Postgres database client for async operations.

Source code in src/embar/db/pg.py

@final
class AsyncPgDb(AsyncDbBase):
    """
    Postgres database client for async operations.
    """

    db_type = "postgres"
    conn_wrapper: AsyncConnectionWrapper[AsyncConnection | AsyncConnectionPool]
    _commit_after_execute: bool = True

    def __init__(self, connection_or_pool: AsyncConnection | AsyncConnectionPool):
        """
        Create a new AsyncPgDb instance.
        """
        self.conn_wrapper = AsyncConnectionWrapper(connection_or_pool)

    async def close(self):
        """
        Close the database connection.
        """
        if self.conn_wrapper:
            await self.conn_wrapper.close()

    def transaction(self) -> AsyncPgDbTransaction:
        """
        Start an isolated transaction.

        ```python notest
        from embar.db.pg import AsyncPgDb
        db = AsyncPgDb(None)

        async with db.transaction() as tx:
            ...
        ```
        """
        return AsyncPgDbTransaction(self)

    def select[M: BaseModel](self, model: type[M]) -> SelectQuery[M, Self]:
        """
        Create a SELECT query.
        """
        return SelectQuery[M, Self](db=self, model=model)

    def select_distinct[M: BaseModel](self, model: type[M]) -> SelectDistinctQuery[M, Self]:
        """
        Create a SELECT query.
        """
        return SelectDistinctQuery[M, Self](db=self, model=model)

    def insert[T: Table](self, table: type[T]) -> InsertQuery[T, Self]:
        """
        Create an INSERT query.
        """
        return InsertQuery[T, Self](table=table, db=self)

    def update[T: Table](self, table: type[T]) -> UpdateQuery[T, Self]:
        """
        Create an UPDATE query.
        """
        return UpdateQuery[T, Self](table=table, db=self)

    def delete[T: Table](self, table: type[T]) -> DeleteQueryReady[T, Self]:
        """
        Create an UPDATE query.
        """
        return DeleteQueryReady[T, Self](table=table, db=self)

    def sql(self, template: Template) -> DbSql[Self]:
        """
        Execute a raw SQL query using template strings.
        """
        return DbSql(template, self)

    def migrate(self, tables: Sequence[type[Table]], enums: Sequence[type[EnumBase]] | None = None) -> Migration[Self]:
        """
        Create a migration from a list of tables.
        """
        ddls = merge_ddls(MigrationDefs(tables, enums))
        return Migration(ddls, self)

    def migrates(self, schema: types.ModuleType) -> Migration[Self]:
        """
        Create a migration from a schema module.
        """
        defs = get_migration_defs(schema)
        return self.migrate(defs.tables, defs.enums)

    @override
    async def execute(self, query: QuerySingle) -> None:
        """
        Execute a query without returning results.
        """
        async with self.conn_wrapper as conn:
            await conn.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                await conn.commit()

    @override
    async def executemany(self, query: QueryMany):
        """
        Execute a query with multiple parameter sets.
        """
        params = _jsonify_dicts(query.many_params)
        async with self.conn_wrapper as conn:
            async with conn.cursor() as cur:
                await cur.executemany(query.sql, params)  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                await conn.commit()

    @override
    async def fetch(self, query: QuerySingle | QueryMany) -> list[dict[str, Any]]:
        """
        Execute a query and return results as a list of dicts.
        """
        async with self.conn_wrapper as conn:
            async with conn.cursor() as cur:
                if isinstance(query, QuerySingle):
                    await cur.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
                else:
                    await cur.executemany(query.sql, query.many_params, returning=True)  # pyright:ignore[reportArgumentType]

                if cur.description is None:
                    return []
                columns: list[str] = [desc[0] for desc in cur.description]
                results: list[dict[str, Any]] = []

                for row in await cur.fetchall():
                    data = dict(zip(columns, row))
                    results.append(data)
            if self._commit_after_execute:
                await conn.commit()
            return results

    @override
    async def truncate(self, schema: str | None = None):
        """
        Truncate all tables in the schema.
        """
        schema = schema if schema is not None else "public"
        tables = await self._get_live_table_names(schema)
        if tables is None:
            return
        table_names = ", ".join(tables)
        async with self.conn_wrapper as conn:
            async with conn.cursor() as cursor:
                await cursor.execute(f"TRUNCATE TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
                if self._commit_after_execute:
                    await conn.commit()

    @override
    async def drop_tables(self, schema: str | None = None):
        """
        Drop all tables in the schema.
        """
        schema = schema if schema is not None else "public"
        tables = await self._get_live_table_names(schema)
        if tables is None:
            return
        table_names = ", ".join(tables)
        async with self.conn_wrapper as conn:
            async with conn.cursor() as cursor:
                await cursor.execute(f"DROP TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
                if self._commit_after_execute:
                    await conn.commit()

    async def _get_live_table_names(self, schema: str) -> list[str] | None:
        async with self.conn_wrapper as conn:
            async with conn.cursor() as cursor:
                # Get all table names from public schema
                await cursor.execute(f"SELECT tablename FROM pg_tables WHERE schemaname = '{schema}'")  # pyright:ignore[reportArgumentType]
                tables = await cursor.fetchall()
                if not tables:
                    return None
                table_names = [f'"{table[0]}"' for table in tables]
            return table_names

__init__(connection_or_pool)

Create a new AsyncPgDb instance.

Source code in src/embar/db/pg.py

def __init__(self, connection_or_pool: AsyncConnection | AsyncConnectionPool):
    """
    Create a new AsyncPgDb instance.
    """
    self.conn_wrapper = AsyncConnectionWrapper(connection_or_pool)

close() async

Close the database connection.

Source code in src/embar/db/pg.py

async def close(self):
    """
    Close the database connection.
    """
    if self.conn_wrapper:
        await self.conn_wrapper.close()

delete(table)

Create an UPDATE query.

Source code in src/embar/db/pg.py

def delete[T: Table](self, table: type[T]) -> DeleteQueryReady[T, Self]:
    """
    Create an UPDATE query.
    """
    return DeleteQueryReady[T, Self](table=table, db=self)

drop_tables(schema=None) async

Drop all tables in the schema.

Source code in src/embar/db/pg.py

@override
async def drop_tables(self, schema: str | None = None):
    """
    Drop all tables in the schema.
    """
    schema = schema if schema is not None else "public"
    tables = await self._get_live_table_names(schema)
    if tables is None:
        return
    table_names = ", ".join(tables)
    async with self.conn_wrapper as conn:
        async with conn.cursor() as cursor:
            await cursor.execute(f"DROP TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                await conn.commit()

execute(query) async

Execute a query without returning results.

Source code in src/embar/db/pg.py

@override
async def execute(self, query: QuerySingle) -> None:
    """
    Execute a query without returning results.
    """
    async with self.conn_wrapper as conn:
        await conn.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
        if self._commit_after_execute:
            await conn.commit()

executemany(query) async

Execute a query with multiple parameter sets.

Source code in src/embar/db/pg.py

@override
async def executemany(self, query: QueryMany):
    """
    Execute a query with multiple parameter sets.
    """
    params = _jsonify_dicts(query.many_params)
    async with self.conn_wrapper as conn:
        async with conn.cursor() as cur:
            await cur.executemany(query.sql, params)  # pyright:ignore[reportArgumentType]
        if self._commit_after_execute:
            await conn.commit()

fetch(query) async

Execute a query and return results as a list of dicts.

Source code in src/embar/db/pg.py

@override
async def fetch(self, query: QuerySingle | QueryMany) -> list[dict[str, Any]]:
    """
    Execute a query and return results as a list of dicts.
    """
    async with self.conn_wrapper as conn:
        async with conn.cursor() as cur:
            if isinstance(query, QuerySingle):
                await cur.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
            else:
                await cur.executemany(query.sql, query.many_params, returning=True)  # pyright:ignore[reportArgumentType]

            if cur.description is None:
                return []
            columns: list[str] = [desc[0] for desc in cur.description]
            results: list[dict[str, Any]] = []

            for row in await cur.fetchall():
                data = dict(zip(columns, row))
                results.append(data)
        if self._commit_after_execute:
            await conn.commit()
        return results

insert(table)

Create an INSERT query.

Source code in src/embar/db/pg.py

def insert[T: Table](self, table: type[T]) -> InsertQuery[T, Self]:
    """
    Create an INSERT query.
    """
    return InsertQuery[T, Self](table=table, db=self)

migrate(tables, enums=None)

Create a migration from a list of tables.

Source code in src/embar/db/pg.py

def migrate(self, tables: Sequence[type[Table]], enums: Sequence[type[EnumBase]] | None = None) -> Migration[Self]:
    """
    Create a migration from a list of tables.
    """
    ddls = merge_ddls(MigrationDefs(tables, enums))
    return Migration(ddls, self)

migrates(schema)

Create a migration from a schema module.

Source code in src/embar/db/pg.py

def migrates(self, schema: types.ModuleType) -> Migration[Self]:
    """
    Create a migration from a schema module.
    """
    defs = get_migration_defs(schema)
    return self.migrate(defs.tables, defs.enums)

select(model)

Create a SELECT query.

Source code in src/embar/db/pg.py

def select[M: BaseModel](self, model: type[M]) -> SelectQuery[M, Self]:
    """
    Create a SELECT query.
    """
    return SelectQuery[M, Self](db=self, model=model)

select_distinct(model)

Create a SELECT query.

Source code in src/embar/db/pg.py

def select_distinct[M: BaseModel](self, model: type[M]) -> SelectDistinctQuery[M, Self]:
    """
    Create a SELECT query.
    """
    return SelectDistinctQuery[M, Self](db=self, model=model)

sql(template)

Execute a raw SQL query using template strings.

Source code in src/embar/db/pg.py

def sql(self, template: Template) -> DbSql[Self]:
    """
    Execute a raw SQL query using template strings.
    """
    return DbSql(template, self)

transaction()

Start an isolated transaction.

```python notest from embar.db.pg import AsyncPgDb db = AsyncPgDb(None)

async with db.transaction() as tx: ... ```

Source code in src/embar/db/pg.py

def transaction(self) -> AsyncPgDbTransaction:
    """
    Start an isolated transaction.

    ```python notest
    from embar.db.pg import AsyncPgDb
    db = AsyncPgDb(None)

    async with db.transaction() as tx:
        ...
    ```
    """
    return AsyncPgDbTransaction(self)

truncate(schema=None) async

Truncate all tables in the schema.

Source code in src/embar/db/pg.py

@override
async def truncate(self, schema: str | None = None):
    """
    Truncate all tables in the schema.
    """
    schema = schema if schema is not None else "public"
    tables = await self._get_live_table_names(schema)
    if tables is None:
        return
    table_names = ", ".join(tables)
    async with self.conn_wrapper as conn:
        async with conn.cursor() as cursor:
            await cursor.execute(f"TRUNCATE TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                await conn.commit()

update(table)

Create an UPDATE query.

Source code in src/embar/db/pg.py

def update[T: Table](self, table: type[T]) -> UpdateQuery[T, Self]:
    """
    Create an UPDATE query.
    """
    return UpdateQuery[T, Self](table=table, db=self)

AsyncPgDbTransaction

Transaction context manager for AsyncPgDb.

Source code in src/embar/db/pg.py

class AsyncPgDbTransaction:
    """
    Transaction context manager for AsyncPgDb.
    """

    _db: AsyncPgDb
    _conn_cm: AbstractAsyncContextManager[AsyncConnection] | None = None
    _tx: AbstractAsyncContextManager[AsyncTransaction] | None = None

    def __init__(self, db: AsyncPgDb):
        self._db = db

    async def __aenter__(self) -> AsyncPgDb:
        pool_or_conn = self._db.conn_wrapper.conn_or_pool

        if isinstance(pool_or_conn, AsyncConnectionPool):
            # Ensure pool is open
            await pool_or_conn.open()

            # Check out a dedicated connection for the transaction
            self._conn_cm = pool_or_conn.connection()
            conn = await self._conn_cm.__aenter__()
        else:
            conn = pool_or_conn

        # Create an AsyncPgDb that uses this single connection (no auto-commit)
        tx_db = AsyncPgDb(conn)
        tx_db._commit_after_execute = False  # pyright: ignore[reportPrivateUsage]
        self._db = tx_db

        self._tx = conn.transaction()
        await self._tx.__aenter__()
        return tx_db

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ):
        result = None
        if self._tx is not None:
            result = await self._tx.__aexit__(exc_type, exc_val, exc_tb)
        if self._conn_cm is not None:
            await self._conn_cm.__aexit__(exc_type, exc_val, exc_tb)
        return result

PgDb

Bases: DbBase

Postgres database client for synchronous operations.

Source code in src/embar/db/pg.py

@final
class PgDb(DbBase):
    """
    Postgres database client for synchronous operations.
    """

    db_type = "postgres"
    conn_wrapper: ConnectionWrapper[Connection | ConnectionPool]
    _commit_after_execute: bool = True

    def __init__(self, connection_or_pool: Connection | ConnectionPool):
        """
        Create a new PgDb instance.
        """
        self.conn_wrapper = ConnectionWrapper(connection_or_pool)

    def close(self):
        """
        Close the database connection.
        """
        if self.conn_wrapper:
            self.conn_wrapper.close()

    def transaction(self) -> PgDbTransaction:
        """
        Start an isolated transaction.

        ```python notest
        from embar.db.pg import PgDb
        db = PgDb(None)

        with db.transaction() as tx:
            ...
        ```
        """
        return PgDbTransaction(self)

    def select[M: BaseModel](self, model: type[M]) -> SelectQuery[M, Self]:
        """
        Create a SELECT query.
        """
        return SelectQuery[M, Self](db=self, model=model)

    def select_distinct[M: BaseModel](self, model: type[M]) -> SelectDistinctQuery[M, Self]:
        """
        Create a SELECT query.
        """
        return SelectDistinctQuery[M, Self](db=self, model=model)

    def insert[T: Table](self, table: type[T]) -> InsertQuery[T, Self]:
        """
        Create an INSERT query.
        """
        return InsertQuery[T, Self](table=table, db=self)

    def update[T: Table](self, table: type[T]) -> UpdateQuery[T, Self]:
        """
        Create an UPDATE query.
        """
        return UpdateQuery[T, Self](table=table, db=self)

    def delete[T: Table](self, table: type[T]) -> DeleteQueryReady[T, Self]:
        """
        Create an UPDATE query.
        """
        return DeleteQueryReady[T, Self](table=table, db=self)

    def sql(self, template: Template) -> DbSql[Self]:
        """
        Execute a raw SQL query using template strings.
        """
        return DbSql(template, self)

    def migrate(self, tables: Sequence[type[Table]], enums: Sequence[type[EnumBase]] | None = None) -> Migration[Self]:
        """
        Create a migration from a list of tables.
        """
        ddls = merge_ddls(MigrationDefs(tables, enums))
        return Migration(ddls, self)

    def migrates(self, schema: types.ModuleType) -> Migration[Self]:
        """
        Create a migration from a schema module.
        """
        defs = get_migration_defs(schema)
        return self.migrate(defs.tables, defs.enums)

    @override
    def execute(self, query: QuerySingle) -> None:
        """
        Execute a query without returning results.
        """
        with self.conn_wrapper as conn:
            conn.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                conn.commit()

    @override
    def executemany(self, query: QueryMany):
        """
        Execute a query with multiple parameter sets.
        """
        params = _jsonify_dicts(query.many_params)
        with self.conn_wrapper as conn:
            with conn.cursor() as cur:
                cur.executemany(query.sql, params)  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                conn.commit()

    @override
    def fetch(self, query: QuerySingle | QueryMany) -> list[dict[str, Any]]:
        """
        Execute a query and return results as a list of dicts.
        """
        with self.conn_wrapper as conn:
            with conn.cursor() as cur:
                if isinstance(query, QuerySingle):
                    cur.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
                else:
                    cur.executemany(query.sql, query.many_params, returning=True)  # pyright:ignore[reportArgumentType]

                if cur.description is None:
                    return []
                columns: list[str] = [desc[0] for desc in cur.description]
                results: list[dict[str, Any]] = []
                for row in cur.fetchall():
                    data = dict(zip(columns, row))
                    results.append(data)
            if self._commit_after_execute:
                conn.commit()  # Commit after SELECT
            return results

    @override
    def truncate(self, schema: str | None = None):
        """
        Truncate all tables in the schema.
        """
        schema = schema if schema is not None else "public"
        tables = self._get_live_table_names(schema)
        if tables is None:
            return
        table_names = ", ".join(tables)
        with self.conn_wrapper as conn:
            with conn.cursor() as cursor:
                cursor.execute(f"TRUNCATE TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
                if self._commit_after_execute:
                    conn.commit()

    @override
    def drop_tables(self, schema: str | None = None):
        """
        Drop all tables in the schema.
        """
        schema = schema if schema is not None else "public"
        tables = self._get_live_table_names(schema)
        if tables is None:
            return
        table_names = ", ".join(tables)
        with self.conn_wrapper as conn:
            with conn.cursor() as cursor:
                cursor.execute(f"DROP TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
                if self._commit_after_execute:
                    conn.commit()

    def _get_live_table_names(self, schema: str) -> list[str] | None:
        with self.conn_wrapper as conn:
            with conn.cursor() as cursor:
                # Get all table names from public schema
                cursor.execute(f"SELECT tablename FROM pg_tables WHERE schemaname = '{schema}'")  # pyright:ignore[reportArgumentType]
                tables = cursor.fetchall()
                if not tables:
                    return None
                table_names = [f'"{table[0]}"' for table in tables]
            return table_names

__init__(connection_or_pool)

Create a new PgDb instance.

Source code in src/embar/db/pg.py

def __init__(self, connection_or_pool: Connection | ConnectionPool):
    """
    Create a new PgDb instance.
    """
    self.conn_wrapper = ConnectionWrapper(connection_or_pool)

close()

Close the database connection.

Source code in src/embar/db/pg.py

def close(self):
    """
    Close the database connection.
    """
    if self.conn_wrapper:
        self.conn_wrapper.close()

delete(table)

Create an UPDATE query.

Source code in src/embar/db/pg.py

def delete[T: Table](self, table: type[T]) -> DeleteQueryReady[T, Self]:
    """
    Create an UPDATE query.
    """
    return DeleteQueryReady[T, Self](table=table, db=self)

drop_tables(schema=None)

Drop all tables in the schema.

Source code in src/embar/db/pg.py

@override
def drop_tables(self, schema: str | None = None):
    """
    Drop all tables in the schema.
    """
    schema = schema if schema is not None else "public"
    tables = self._get_live_table_names(schema)
    if tables is None:
        return
    table_names = ", ".join(tables)
    with self.conn_wrapper as conn:
        with conn.cursor() as cursor:
            cursor.execute(f"DROP TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                conn.commit()

execute(query)

Execute a query without returning results.

Source code in src/embar/db/pg.py

@override
def execute(self, query: QuerySingle) -> None:
    """
    Execute a query without returning results.
    """
    with self.conn_wrapper as conn:
        conn.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
        if self._commit_after_execute:
            conn.commit()

executemany(query)

Execute a query with multiple parameter sets.

Source code in src/embar/db/pg.py

@override
def executemany(self, query: QueryMany):
    """
    Execute a query with multiple parameter sets.
    """
    params = _jsonify_dicts(query.many_params)
    with self.conn_wrapper as conn:
        with conn.cursor() as cur:
            cur.executemany(query.sql, params)  # pyright:ignore[reportArgumentType]
        if self._commit_after_execute:
            conn.commit()

fetch(query)

Execute a query and return results as a list of dicts.

Source code in src/embar/db/pg.py

@override
def fetch(self, query: QuerySingle | QueryMany) -> list[dict[str, Any]]:
    """
    Execute a query and return results as a list of dicts.
    """
    with self.conn_wrapper as conn:
        with conn.cursor() as cur:
            if isinstance(query, QuerySingle):
                cur.execute(query.sql, query.params)  # pyright:ignore[reportArgumentType]
            else:
                cur.executemany(query.sql, query.many_params, returning=True)  # pyright:ignore[reportArgumentType]

            if cur.description is None:
                return []
            columns: list[str] = [desc[0] for desc in cur.description]
            results: list[dict[str, Any]] = []
            for row in cur.fetchall():
                data = dict(zip(columns, row))
                results.append(data)
        if self._commit_after_execute:
            conn.commit()  # Commit after SELECT
        return results

insert(table)

Create an INSERT query.

Source code in src/embar/db/pg.py

def insert[T: Table](self, table: type[T]) -> InsertQuery[T, Self]:
    """
    Create an INSERT query.
    """
    return InsertQuery[T, Self](table=table, db=self)

migrate(tables, enums=None)

Create a migration from a list of tables.

Source code in src/embar/db/pg.py

def migrate(self, tables: Sequence[type[Table]], enums: Sequence[type[EnumBase]] | None = None) -> Migration[Self]:
    """
    Create a migration from a list of tables.
    """
    ddls = merge_ddls(MigrationDefs(tables, enums))
    return Migration(ddls, self)

migrates(schema)

Create a migration from a schema module.

Source code in src/embar/db/pg.py

def migrates(self, schema: types.ModuleType) -> Migration[Self]:
    """
    Create a migration from a schema module.
    """
    defs = get_migration_defs(schema)
    return self.migrate(defs.tables, defs.enums)

select(model)

Create a SELECT query.

Source code in src/embar/db/pg.py

def select[M: BaseModel](self, model: type[M]) -> SelectQuery[M, Self]:
    """
    Create a SELECT query.
    """
    return SelectQuery[M, Self](db=self, model=model)

select_distinct(model)

Create a SELECT query.

Source code in src/embar/db/pg.py

def select_distinct[M: BaseModel](self, model: type[M]) -> SelectDistinctQuery[M, Self]:
    """
    Create a SELECT query.
    """
    return SelectDistinctQuery[M, Self](db=self, model=model)

sql(template)

Execute a raw SQL query using template strings.

Source code in src/embar/db/pg.py

def sql(self, template: Template) -> DbSql[Self]:
    """
    Execute a raw SQL query using template strings.
    """
    return DbSql(template, self)

transaction()

Start an isolated transaction.

```python notest from embar.db.pg import PgDb db = PgDb(None)

with db.transaction() as tx: ... ```

Source code in src/embar/db/pg.py

def transaction(self) -> PgDbTransaction:
    """
    Start an isolated transaction.

    ```python notest
    from embar.db.pg import PgDb
    db = PgDb(None)

    with db.transaction() as tx:
        ...
    ```
    """
    return PgDbTransaction(self)

truncate(schema=None)

Truncate all tables in the schema.

Source code in src/embar/db/pg.py

@override
def truncate(self, schema: str | None = None):
    """
    Truncate all tables in the schema.
    """
    schema = schema if schema is not None else "public"
    tables = self._get_live_table_names(schema)
    if tables is None:
        return
    table_names = ", ".join(tables)
    with self.conn_wrapper as conn:
        with conn.cursor() as cursor:
            cursor.execute(f"TRUNCATE TABLE {table_names} CASCADE")  # pyright:ignore[reportArgumentType]
            if self._commit_after_execute:
                conn.commit()

update(table)

Create an UPDATE query.

Source code in src/embar/db/pg.py

def update[T: Table](self, table: type[T]) -> UpdateQuery[T, Self]:
    """
    Create an UPDATE query.
    """
    return UpdateQuery[T, Self](table=table, db=self)

PgDbTransaction

Transaction context manager for PgDb.

Source code in src/embar/db/pg.py

class PgDbTransaction:
    """
    Transaction context manager for PgDb.
    """

    _db: PgDb
    _conn_cm: AbstractContextManager[Connection] | None = None
    _tx: AbstractContextManager[Transaction] | None = None

    def __init__(self, db: PgDb):
        self._db = db

    def __enter__(self) -> PgDb:
        pool_or_conn = self._db.conn_wrapper.conn_or_pool

        if isinstance(pool_or_conn, ConnectionPool):
            # Ensure pool is open (idempotent if already open)
            pool_or_conn.open()

            # Check out a dedicated connection for the transaction
            self._conn_cm = pool_or_conn.connection()
            conn = self._conn_cm.__enter__()
        else:
            conn = pool_or_conn

        # Create a PgDb that uses this single connection (no auto-commit)
        tx_db = PgDb(conn)
        tx_db._commit_after_execute = False  # pyright: ignore[reportPrivateUsage]
        self._db = tx_db

        self._tx = conn.transaction()
        self._tx.__enter__()
        return tx_db

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ):
        result = None
        if self._tx is not None:
            result = self._tx.__exit__(exc_type, exc_val, exc_tb)
        if self._conn_cm is not None:
            self._conn_cm.__exit__(exc_type, exc_val, exc_tb)
        return result