encore-dev 2.0.0a7__py3-none-any.whl

encore/storage/postgres/__init__.py

@@ -0,0 +1,737 @@
+"""Postgres declarations, explicit access handles, and streaming queries.
+
+Declare databases at module scope and use their query methods to execute SQL.
+Use begin() for transactions and acquire() for dedicated pooled connections.
+
+Typical usage example:
+
+    from encore.storage import postgres
+
+    db = postgres.Database("orders", migrations="./migrations")
+
+    async def get_user(user_id: int):
+        return await db.query("SELECT * FROM users WHERE id = $1", user_id).one()
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncGenerator, Callable, Generator
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, Literal
+
+from ...internal.lazy import Lazy
+from ...internal.runtime import EncoreError, _native, runtime
+from ...tracing._context import active_ctx_bytes
+
+__all__ = [
+    "Database",
+    "DatabaseRef",
+    "Queryable",
+    "Cursor",
+    "Connection",
+    "Transaction",
+    "ExecResult",
+    "Row",
+]
+
+logger = logging.getLogger(__name__)
+Row = dict[str, Any]
+_Access = Literal["reader", "writer", "admin"]
+
+
+@dataclass(frozen=True)
+class ExecResult:
+    """An immutable result of a completed SQL statement.
+
+    Attributes:
+        rows_affected: The number of rows affected, as reported by Postgres.
+    """
+
+    rows_affected: int
+
+
+class Cursor:
+    """A SQL cursor consumed once by iteration, one(), or all().
+
+    Execution starts when iteration advances or a helper is awaited. Starting
+    one consumption mode prevents any other consumption of this query, including
+    after an error or partial iteration. Obtain cursors from query() on a
+    database, reference, connection, or transaction.
+
+    Exhaust a stream or call aclose() to release its connection. For early exits,
+    use the query as an async context manager. Entering the context does not
+    execute or consume the query. Leaving it closes the stream. The one() and
+    all() helpers close their cursors automatically.
+    """
+
+    def __init__(self, rows: AsyncGenerator[Row, None]) -> None:
+        """Initialize a cursor with a deferred row stream.
+
+        Args:
+            rows: The row stream that executes the query when advanced.
+        """
+        self._rows = rows
+        self._consumed = False
+
+    def _consume(self) -> AsyncGenerator[Row, None]:
+        if self._consumed:
+            raise EncoreError("query has already been consumed")
+        self._consumed = True
+        return self._rows
+
+    def __aiter__(self) -> AsyncGenerator[Row, None]:
+        """Return the row iterator and reserve this query for streaming.
+
+        Returns:
+            An async iterator yielding dictionaries of column names to values.
+        """
+        return self._consume()
+
+    async def one(self) -> Row | None:
+        """Return zero or one row and close the cursor.
+
+        Returns:
+            A dictionary of column names to values, or None if no rows match.
+
+        Raises:
+            EncoreError: The query fails or returns more than one row.
+        """
+        rows = self._consume()
+        try:
+            first = None
+            async for row in rows:
+                if first is not None:
+                    raise EncoreError("query returned more than one row")
+                first = row
+            return first
+        finally:
+            await rows.aclose()
+
+    async def all(self) -> list[Row]:
+        """Collect all rows in memory and close the cursor.
+
+        Returns:
+            The result rows as dictionaries of column names to values.
+            An empty list if no rows match.
+
+        Raises:
+            EncoreError: The query fails.
+        """
+        rows = self._consume()
+        try:
+            return [row async for row in rows]
+        finally:
+            await rows.aclose()
+
+    async def aclose(self) -> None:
+        """Close the row stream and prevent further consumption of this query.
+
+        Closing an unstarted query does not execute SQL. Repeated closes have
+        no effect. Wait for any active read or helper to finish before closing.
+        """
+        self._consumed = True
+        await self._rows.aclose()
+
+    async def __aenter__(self) -> Cursor:
+        """Return this cursor without executing or consuming it."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Close the query without suppressing an exception from the block.
+
+        Args:
+            exc_type: The exception type, or None on a clean exit.
+            exc: The exception instance, or None on a clean exit.
+            tb: The traceback, or None on a clean exit.
+        """
+        await self.aclose()
+
+
+class _Queries:
+    """A query executor shared by the public Postgres wrappers.
+
+    Attributes:
+        get_handle: A callable that resolves the native execution handle on
+            demand.
+    """
+
+    def __init__(self, get_handle: Callable[[], _native.PostgresHandle]) -> None:
+        """Initialize the executor without resolving its native handle.
+
+        Args:
+            get_handle: A callable that returns the native execution handle.
+        """
+        self.get_handle = get_handle
+
+    async def query(self, sql: str, *params: Any) -> AsyncGenerator[Row, None]:
+        """Yield rows and close the native cursor when iteration ends.
+
+        Args:
+            sql: A SQL statement with positional placeholders such as $1 and $2.
+            *params: Values bound to the placeholders in order.
+
+        Yields:
+            Dictionaries mapping column names to decoded Python values.
+
+        Raises:
+            EncoreError: Postgres cannot execute the query or retrieve a row.
+        """
+        cursor = await self.get_handle().query(
+            sql, list(params), source=active_ctx_bytes()
+        )
+        try:
+            while (row := await cursor.next()) is not None:
+                yield row
+        finally:
+            await cursor.close()
+
+    async def exec(self, sql: str, *params: Any) -> ExecResult:
+        """Execute a statement and wait for its affected-row count.
+
+        Args:
+            sql: A SQL statement with positional placeholders such as $1 and $2.
+            *params: Values bound to the placeholders in order.
+
+        Returns:
+            An ExecResult containing the affected-row count reported by
+            Postgres. Any rows returned by the statement are discarded.
+
+        Raises:
+            EncoreError: Postgres cannot complete the statement.
+        """
+        count = await self.get_handle().exec(
+            sql, list(params), source=active_ctx_bytes()
+        )
+        return ExecResult(rows_affected=count)
+
+
+class Queryable:
+    """A query interface for database references, connections, and transactions.
+
+    SQL parameters use $1, $2, and subsequent positional placeholders.
+    Connections and transactions must be open and have no other active operation
+    or row stream when starting a query.
+    """
+
+    def __init__(self, queries: _Queries) -> None:
+        """Initialize the interface with its query executor.
+
+        Args:
+            queries: The executor used by this handle's query methods.
+        """
+        self._queries = queries
+
+    def query(self, sql: str, *params: Any) -> Cursor:
+        """Create a query for streaming or fetching rows.
+
+        Consume the query once with async iteration, one(), or all(). No SQL runs
+        until iteration advances or a helper is awaited. Exhaust the stream or
+        close it before reusing its connection or transaction. Use
+        the query as an async context manager when leaving iteration early.
+
+        Args:
+            sql: A SQL statement with positional placeholders such as $1 and $2.
+            *params: Values bound to the placeholders in order.
+
+        Returns:
+            A Cursor whose rows map column names to decoded Python values.
+
+        Raises:
+            EncoreError: During iteration, acquiring a connection, executing the
+                query, or retrieving a row fails.
+        """
+        return Cursor(self._queries.query(sql, *params))
+
+    async def exec(self, sql: str, *params: Any) -> ExecResult:
+        """Execute a statement and discard any returned rows.
+
+        Args:
+            sql: A SQL statement with positional placeholders such as $1 and $2.
+            *params: Values bound to the placeholders in order.
+
+        Returns:
+            An ExecResult containing the affected-row count reported by
+            Postgres.
+
+        Raises:
+            EncoreError: Acquiring a connection or executing the statement fails.
+        """
+        return await self._queries.exec(sql, *params)
+
+
+class DatabaseRef(Queryable):
+    """A database handle that can be passed to other functions.
+
+    Obtain a reference from Database.read(), Database.write(), or
+    Database.admin() so Encore can infer the service's required permissions.
+    References share the database's pool and connection privileges.
+    """
+
+    def __init__(
+        self,
+        get_handle: Callable[[], _native.PostgresHandle],
+        *,
+        access: _Access,
+        primary: bool = True,
+    ) -> None:
+        """Initialize a reference without resolving its native handle.
+
+        Application code should use Database.read(), Database.write(), or
+        Database.admin() so Encore can discover the required permissions.
+
+        Args:
+            get_handle: A callable that returns the database's native handle.
+            access: The selected access level: "reader", "writer", or "admin".
+            primary: The stored primary-server preference. Defaults to True.
+                Does not affect connection selection.
+        """
+        super().__init__(_Queries(get_handle))
+        self._access = access
+        self._primary = primary
+
+    @property
+    def connection_string(self) -> str:
+        """The proxy connection string for connecting with another client.
+
+        All references return the same URL, regardless of their
+        selected access level or primary-server preference.
+        """
+        return self._queries.get_handle().conn_string()
+
+    def begin(self) -> _BeginTransaction:
+        """Create an awaitable that begins a pooled transaction.
+
+        Returns:
+            An awaitable that produces a Transaction, also usable as an async
+            context manager. No transaction starts until the result is awaited
+            or entered. A clean context exit commits. An exceptional exit rolls
+            back. When awaiting directly, commit or roll back the transaction
+            explicitly.
+
+        Raises:
+            EncoreError: When awaited or entered, the runtime cannot acquire a
+                connection or begin the transaction.
+        """
+        return _BeginTransaction(self._queries.get_handle())
+
+    def acquire(self) -> _AcquireConnection:
+        """Create an awaitable that acquires a pooled connection.
+
+        Returns:
+            An awaitable that produces a Connection, also usable as an async
+            context manager. No connection is acquired until the result is
+            awaited or entered. Leaving the context closes the connection. When
+            awaiting directly, close the connection explicitly after use.
+
+        Raises:
+            EncoreError: When awaited or entered, the runtime cannot acquire a
+                connection.
+        """
+        return _AcquireConnection(self._queries.get_handle())
+
+
+class Database:
+    """A module-level Postgres database declaration.
+
+    Encore infers read-write permissions from direct database operations.
+    Use read(), write(), or admin() to obtain a reference that can be passed to
+    other functions.
+    """
+
+    def __init__(self, name: str, /, *, migrations: str | None = None) -> None:
+        """Declare a new Postgres database.
+
+        Encore discovers this declaration at module scope and provisions the
+        database. Constructing it does not acquire a connection.
+
+        Args:
+            name: The logical database name, supplied as a positional string
+                literal so Encore can resolve it statically.
+            migrations: A slash-separated path to a subdirectory relative to the
+                declaring file, such as "./migrations". Must be a string literal
+                or None. Defaults to None, which declares no Encore migration
+                directory.
+        """
+        self._name = name
+        self._migrations = migrations
+        self._handle = Lazy(lambda: runtime().sql_database(name))
+        self._writer = DatabaseRef(self._handle.get, access="writer")
+
+    @classmethod
+    def named(cls, name: str, /) -> Database:
+        """Reference an existing database without declaring it again.
+
+        Use this at module scope. The database is configured where it is
+        declared.
+
+        Args:
+            name: The existing database's logical name, supplied as a positional
+                string literal so Encore can resolve it statically.
+
+        Returns:
+            A Database that refers to the existing resource.
+        """
+        return cls(name)
+
+    @property
+    def name(self) -> str:
+        """The logical name of the database."""
+        return self._name
+
+    def admin(self) -> DatabaseRef:
+        """Create a handle that records admin access for this service.
+
+        Returns:
+            A reference sharing this database's pool and connection privileges.
+        """
+        return DatabaseRef(self._handle.get, access="admin")
+
+    def write(self) -> DatabaseRef:
+        """Return a handle that records read-write access for this service.
+
+        Returns:
+            A reference sharing this database's pool and connection privileges.
+        """
+        return self._writer
+
+    def read(self, primary: bool = False) -> DatabaseRef:
+        """Create a handle that records reader access for this service.
+
+        Args:
+            primary: The stored primary-server preference. Defaults to False.
+                Does not affect connection selection.
+
+        Returns:
+            A reference sharing this database's pool and connection privileges.
+        """
+        return DatabaseRef(self._handle.get, access="reader", primary=primary)
+
+    def query(self, sql: str, *params: Any) -> Cursor:
+        """Create a query for streaming or fetching rows.
+
+        Consume the query once with async iteration, one(), or all(). No SQL runs
+        until iteration advances or a helper is awaited. See Queryable.query
+        for stream cleanup requirements.
+
+        Args:
+            sql: A SQL statement with positional placeholders such as $1 and $2.
+            *params: Values bound to the placeholders in order.
+
+        Returns:
+            A Cursor whose rows map column names to decoded Python values.
+
+        Raises:
+            EncoreError: During iteration, the runtime cannot acquire a
+                connection or Postgres cannot execute the query or retrieve a
+                row.
+        """
+        return self._writer.query(sql, *params)
+
+    async def exec(self, sql: str, *params: Any) -> ExecResult:
+        """Execute a statement and discard any returned rows.
+
+        Args:
+            sql: A SQL statement with positional placeholders such as $1 and $2.
+            *params: Values bound to the placeholders in order.
+
+        Returns:
+            An ExecResult containing the affected-row count reported by
+            Postgres. Any rows returned by the statement are discarded.
+
+        Raises:
+            EncoreError: The runtime cannot acquire a connection or the
+                statement fails.
+        """
+        return await self._writer.exec(sql, *params)
+
+    @property
+    def connection_string(self) -> str:
+        """The proxy connection string for connecting with another client."""
+        return self._writer.connection_string
+
+    def begin(self) -> _BeginTransaction:
+        """Create an awaitable that begins a pooled transaction.
+
+        Returns:
+            An awaitable that produces a Transaction, also usable as an async
+            context manager. No transaction starts until the result is awaited
+            or entered. A clean context exit commits. An exceptional exit rolls
+            back. When awaiting directly, commit or roll back the transaction
+            explicitly.
+
+        Raises:
+            EncoreError: When awaited or entered, the runtime cannot acquire a
+                connection or begin the transaction.
+        """
+        return self._writer.begin()
+
+    def acquire(self) -> _AcquireConnection:
+        """Create an awaitable that acquires a pooled connection.
+
+        Returns:
+            An awaitable that produces a Connection, also usable as an async
+            context manager. No connection is acquired until the result is
+            awaited or entered. Leaving the context closes the connection. When
+            awaiting directly, close the connection explicitly after use.
+
+        Raises:
+            EncoreError: When awaited or entered, the runtime cannot acquire a
+                connection.
+        """
+        return self._writer.acquire()
+
+
+class Connection(Queryable):
+    """An acquired connection returned to its pool when closed or destroyed.
+
+    Obtain a connection with Database.acquire() or DatabaseRef.acquire(). Close
+    it explicitly or use it as an async context manager. Operations require an
+    open connection with no other active operation, transaction, or row stream.
+    Finish the transaction or close the stream before reusing the connection.
+    """
+
+    def __init__(self, impl: _native.PostgresHandle) -> None:
+        """Wrap an acquired native connection.
+
+        Args:
+            impl: An acquired native connection handle. Application code should
+                obtain a Connection by awaiting or entering acquire().
+        """
+        self._impl = impl
+        super().__init__(_Queries(lambda: self._impl))
+
+    def begin(self) -> _BeginTransaction:
+        """Create an awaitable that begins a transaction on this connection.
+
+        Returns:
+            An awaitable that produces a Transaction, also usable as an async
+            context manager. The transaction exclusively borrows this connection
+            until it finishes. A clean context exit commits. An exceptional exit
+            rolls back. The connection stays acquired afterwards.
+
+        Raises:
+            EncoreError: When awaited or entered, Postgres cannot begin the
+                transaction.
+        """
+        return _BeginTransaction(self._impl)
+
+    async def close(self) -> None:
+        """Return this connection to its pool.
+
+        Repeated calls have no effect after the connection is closed. Close any
+        active cursor and finish any transaction before closing the connection.
+
+        """
+        await self._impl.close()
+
+    async def __aenter__(self) -> Connection:
+        """Return this open connection for use in an async context."""
+        if self._impl.is_closed():
+            raise EncoreError("database handle is closed")
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Close the connection without suppressing an exception from the block.
+
+        Args:
+            exc_type: The exception type, or None on a clean exit.
+            exc: The exception instance, or None on a clean exit.
+            tb: The traceback, or None on a clean exit.
+
+        """
+        await self.close()
+
+
+class Transaction(Queryable):
+    """An open transaction that commits or rolls back on context exit.
+
+    Obtain a transaction by awaiting or entering begin(). A clean context exit
+    commits. An exceptional exit rolls back. Without a context manager, commit
+    or roll back explicitly. Abandoned transactions are rolled back when their
+    native handles are destroyed.
+
+    Once commit or rollback starts, it consumes the handle even if Postgres
+    reports an error. Operations require an open transaction with no other
+    active operation or row stream. Close row streams before committing or
+    rolling back.
+    """
+
+    def __init__(self, impl: _native.PostgresHandle) -> None:
+        """Wrap an open native transaction.
+
+        Args:
+            impl: An open native transaction handle. Application code should
+                obtain a Transaction by awaiting or entering begin().
+        """
+        self._impl = impl
+        super().__init__(_Queries(lambda: self._impl))
+
+    async def commit(self) -> None:
+        """Commit the transaction and close its handle.
+
+        Close active row streams before committing. Once completion starts, it
+        continues if the awaiting task is cancelled, and a borrowed connection
+        remains unavailable until completion finishes. A Postgres error still
+        consumes the transaction handle.
+
+        Raises:
+            EncoreError: Postgres cannot commit the transaction.
+        """
+        await self._impl.commit()
+
+    async def rollback(self) -> None:
+        """Roll back the transaction and close its handle.
+
+        Close active row streams before rolling back. Once completion starts, it
+        continues if the awaiting task is cancelled, and a borrowed connection
+        remains unavailable until completion finishes. A Postgres error still
+        consumes the transaction handle.
+
+        Raises:
+            EncoreError: Postgres cannot roll back the transaction.
+        """
+        await self._impl.rollback()
+
+    async def __aenter__(self) -> Transaction:
+        """Return this open transaction for use in an async context."""
+        if self._impl.is_closed():
+            raise EncoreError("database handle is closed")
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Commit on a clean exit or roll back while an exception unwinds.
+
+        An already completed transaction needs no further action. If rollback
+        fails while handling an exception, the failure is logged and the
+        original exception propagates.
+
+        Args:
+            exc_type: The exception type, or None on a clean exit.
+            exc: The exception instance, or None on a clean exit.
+            tb: The traceback, or None on a clean exit.
+
+        Raises:
+            EncoreError: Committing on a clean exit fails.
+        """
+        if self._impl.is_closed():
+            return
+        if exc_type is None:
+            await self.commit()
+            return
+        try:
+            await self.rollback()
+        except Exception:
+            logger.warning(
+                "rolling back the transaction failed while unwinding an error",
+                exc_info=True,
+            )
+
+
+class _BeginTransaction:
+    """An awaitable and async context manager that starts a transaction.
+
+    No transaction starts until the result is awaited or entered.
+    """
+
+    def __init__(self, impl: _native.PostgresHandle) -> None:
+        """Initialize the deferred transaction operation.
+
+        Args:
+            impl: A native database or acquired connection handle.
+        """
+        self._impl = impl
+        self._tx: Transaction | None = None
+
+    async def _begin(self) -> Transaction:
+        """Return a transaction after beginning it on the native handle."""
+        return Transaction(await self._impl.begin(source=active_ctx_bytes()))
+
+    def __await__(self) -> Generator[Any, Any, Transaction]:
+        """Return the iterator that drives transaction creation when awaited."""
+        return self._begin().__await__()
+
+    async def __aenter__(self) -> Transaction:
+        """Return a new transaction for the async context."""
+        self._tx = await self._begin()
+        return self._tx
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Commit on a clean exit or roll back on an exceptional exit.
+
+        Do nothing if the transaction has already finished. A rollback failure
+        is logged without suppressing the exception from the block.
+
+        Args:
+            exc_type: The exception type, or None on a clean exit.
+            exc: The exception instance, or None on a clean exit.
+            tb: The traceback, or None on a clean exit.
+        """
+        if self._tx is not None:
+            await self._tx.__aexit__(exc_type, exc, tb)
+
+
+class _AcquireConnection:
+    """An awaitable and async context manager that acquires a connection.
+
+    No connection is acquired until the result is awaited or entered.
+    """
+
+    def __init__(self, impl: _native.PostgresHandle) -> None:
+        """Initialize the deferred connection acquisition.
+
+        Args:
+            impl: A native database handle whose pool supplies the connection.
+        """
+        self._impl = impl
+        self._conn: Connection | None = None
+
+    async def _acquire(self) -> Connection:
+        """Return a connection after acquiring it from the native pool."""
+        return Connection(await self._impl.acquire(source=active_ctx_bytes()))
+
+    def __await__(self) -> Generator[Any, Any, Connection]:
+        """Return the iterator that acquires the connection when awaited."""
+        return self._acquire().__await__()
+
+    async def __aenter__(self) -> Connection:
+        """Return a newly acquired connection for the async context."""
+        self._conn = await self._acquire()
+        return self._conn
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Close the connection when the context exits.
+
+        Args:
+            exc_type: The exception type, or None on a clean exit.
+            exc: The exception instance, or None on a clean exit.
+            tb: The traceback, or None on a clean exit.
+        """
+        if self._conn is not None:
+            await self._conn.__aexit__(exc_type, exc, tb)