pyfect-dev 0.1.0__py3-none-any.whl

pyfect/sql/mysql/_session.py

@@ -0,0 +1,262 @@
+"""One asyncmy connection, used from the SQL loop thread (``pyfect.sql._loop``).
+
+Statements take turns on it (a lock, as mysql2 queues commands), and a
+stream holds the turn until it closes. A prepared statement goes through
+``COM_STMT_PREPARE`` / ``COM_STMT_EXECUTE`` (binary protocol), kept
+prepared per connection; any other statement is sent as text with its
+parameters escaped in (``_bind``), several statements at once allowed.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from collections import OrderedDict
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, cast
+
+from pyfect import _json_text as json_text
+from pyfect.sql.mysql import _asyncmy
+from pyfect.sql.mysql._asyncmy import Cursor, PreparedStatement, RawConnection
+from pyfect.sql.mysql._bind import bind
+
+_STATEMENT_CACHE_SIZE = 256
+_STREAM_CHUNK = 128
+_KILL_TIMEOUT = 5.0
+_NO_BACKSLASH_ESCAPES = 0x200  # the protocol's SERVER_STATUS_NO_BACKSLASH_ESCAPES
+_DRAIN_GRACE = 0.1
+
+_log = logging.getLogger("pyfect.sql.mysql")
+
+
+@dataclass(frozen=True, slots=True)
+class ResultSetHeader:
+    """What a statement without rows returns (mysql2's ``ResultSetHeader``)."""
+
+    affected_rows: int
+    insert_id: int
+
+
+type Row = dict[str, Any]
+type ResultSet = list[Row] | list[tuple[Any, ...]] | ResultSetHeader
+
+
+def _header(affected_rows: object, insert_id: int) -> ResultSetHeader:
+    affected = affected_rows if isinstance(affected_rows, int) else 0
+    return ResultSetHeader(affected_rows=affected, insert_id=insert_id)
+
+
+def _parameter(value: object) -> object:
+    """A dict or list as its JSON text (mysql2's ``execute``); asyncmy would
+    send its ``repr``."""
+    if isinstance(value, dict | list):
+        return json_text.stringify(cast("object", value))
+    return value
+
+
+def _parameters(params: Sequence[object]) -> tuple[object, ...]:
+    return tuple(_parameter(value) for value in params)
+
+
+class Session:
+    """A connection and its prepared statements."""
+
+    __slots__ = ("_lock", "_options", "_statements", "broken", "raw")
+
+    def __init__(self, raw: RawConnection, options: Mapping[str, Any]) -> None:
+        self.raw = raw
+        self._options = options
+        self._lock = asyncio.Lock()
+        self._statements: OrderedDict[str, PreparedStatement] = OrderedDict()
+        self.broken = False
+
+    @staticmethod
+    async def open(options: Mapping[str, Any]) -> Session:
+        return Session(await _asyncmy.connect(options), options)
+
+    @property
+    def usable(self) -> bool:
+        return not self.broken and self.raw.connected
+
+    def _cut_short(self) -> None:
+        # A read cut short leaves the protocol out of step.
+        self.broken = True
+        self.raw.close()
+
+    # --- statements -----------------------------------------------------------
+
+    async def _prepared(self, sql: str) -> PreparedStatement:
+        statements = self._statements
+        found = statements.get(sql)
+        if found is not None:
+            statements.move_to_end(sql)
+            return found
+        prepared = await self.raw.prepare(sql)
+        statements[sql] = prepared
+        if len(statements) > _STATEMENT_CACHE_SIZE:
+            _, oldest = statements.popitem(last=False)
+            await oldest.close()
+        return prepared
+
+    async def _execute(
+        self, sql: str, params: Sequence[object], as_tuples: bool
+    ) -> list[ResultSet]:
+        prepared = await self._prepared(sql)
+        result = await prepared.execute(_parameters(params))
+        description = result.description
+        if description is None:
+            return [_header(result.affected_rows, result.insert_id)]
+        rows = result.rows or ()
+        if as_tuples:
+            return [list(rows)]
+        names = [column[0] for column in description]
+        return [[dict(zip(names, row, strict=True)) for row in rows]]
+
+    def _bound(self, sql: str, params: Sequence[object]) -> str:
+        # The server reports the session's mode with every result, and the
+        # connection's ``escape`` follows the same flag.
+        no_backslash_escapes = bool(self.raw.server_status & _NO_BACKSLASH_ESCAPES)
+        return bind(
+            sql, _parameters(params), self.raw.escape, no_backslash_escapes=no_backslash_escapes
+        )
+
+    async def _query(self, sql: str, params: Sequence[object], as_tuples: bool) -> list[ResultSet]:
+        text = self._bound(sql, params)
+        cursor = _asyncmy.text_cursor(self.raw, as_tuples=as_tuples)
+        results: list[ResultSet] = []
+        try:
+            await cursor.execute(text)
+            while True:
+                if cursor.description is None:
+                    results.append(_header(cursor.rowcount, cursor.lastrowid))
+                else:
+                    results.append(list(await cursor.fetchall()))
+                if not await cursor.nextset():
+                    break
+        except Exception:
+            # A failed statement leaves the connection in step, so what
+            # follows it can be read past.
+            with contextlib.suppress(Exception):
+                await cursor.close()
+            raise
+        await cursor.close()
+        return results
+
+    async def run(
+        self, sql: str, params: Sequence[object], *, prepared: bool, as_tuples: bool
+    ) -> list[ResultSet]:
+        """Each statement's result: its rows, or a header."""
+        async with self._lock:
+            try:
+                if prepared:
+                    return await self._execute(sql, params, as_tuples)
+                return await self._query(sql, params, as_tuples)
+            except asyncio.CancelledError:
+                self._cut_short()
+                raise
+
+    # --- streams --------------------------------------------------------------
+
+    async def open_stream(self, sql: str, params: Sequence[object]) -> Cursor:
+        """An unbuffered cursor on ``sql``; the session is this stream's
+        until ``close_stream``."""
+        await self._lock.acquire()
+        try:
+            cursor = _asyncmy.stream_cursor(self.raw)
+            await cursor.execute(self._bound(sql, params))
+        except asyncio.CancelledError:
+            self._cut_short()
+            self._lock.release()
+            raise
+        except BaseException:
+            self._lock.release()
+            raise
+        return cursor
+
+    async def next_rows(self, cursor: Cursor) -> list[Row]:
+        """Up to a chunk of rows; ``[]`` at the end."""
+        try:
+            return list(await cursor.fetchmany(_STREAM_CHUNK))
+        except asyncio.CancelledError:
+            self._cut_short()
+            raise
+        except Exception as error:
+            if _asyncmy.is_server_error(error):
+                _asyncmy.end_unbuffered_result(self.raw)
+            else:
+                self._cut_short()
+            raise
+
+    async def close_stream(self, cursor: Cursor) -> None:
+        """End the stream, then give the turn back: what it left unread is read past, and a
+        statement still sending rows after a moment is killed, as it may send for hours.
+        Only the statement: a transaction goes on, if the connection stays in step."""
+        try:
+            if self.usable:
+                await self._read_past(cursor)
+        finally:
+            self._lock.release()
+
+    async def _read_past(self, cursor: Cursor) -> None:
+        reading = asyncio.ensure_future(cursor.close())
+        try:
+            done, _ = await asyncio.wait([reading], timeout=_DRAIN_GRACE)
+            if not done:
+                await self.kill(only_statement=True)
+                done, _ = await asyncio.wait([reading], timeout=_KILL_TIMEOUT)
+            if not done:
+                reading.cancel()
+                await asyncio.wait([reading])
+                self._cut_short()
+                await self.kill()
+            elif (error := reading.exception()) is not None:
+                if _asyncmy.is_server_error(error):
+                    _asyncmy.end_unbuffered_result(self.raw)
+                else:
+                    self._cut_short()
+        except asyncio.CancelledError:
+            reading.cancel()
+            await asyncio.wait([reading])
+            self._cut_short()
+            raise
+
+    # --- ending ---------------------------------------------------------------
+
+    async def kill(self, *, only_statement: bool = False) -> None:
+        """End this connection (or only its statement) on the server, from another
+        connection, so a statement cut short stops holding its locks. A kill that cannot be
+        sent fails nobody, but the statement stays on the server, so it is logged."""
+        thread_id = self.raw.thread_id()
+        try:
+            async with asyncio.timeout(_KILL_TIMEOUT):
+                side = await _asyncmy.connect(self._options)
+                try:
+                    await _asyncmy.kill(side, thread_id, only_statement=only_statement)
+                finally:
+                    side.close()
+        except Exception as error:
+            _log.warning(
+                "KILL %d failed, so its statement stays on the server: %r", thread_id, error
+            )
+
+    async def stop(self, task: asyncio.Task[Any]) -> None:
+        """Interrupt ``task``'s statement: a statement on the wire is cut
+        short, which closes this connection, and the server is told to
+        end it."""
+        if task.done():
+            return
+        task.cancel()
+        await asyncio.wait([task])
+        if self.broken:
+            await self.kill()
+
+    async def close(self) -> None:
+        with contextlib.suppress(Exception):
+            if self.raw.connected:
+                await self.raw.ensure_closed()
+        self.raw.close()
+
+
+__all__ = ["ResultSet", "ResultSetHeader", "Row", "Session"]

pyfect/sql/mysql/client.py

@@ -0,0 +1,244 @@
+"""The MySQL client (Effect's ``@effect/sql-mysql2`` ``MysqlClient``): an
+``SqlClient`` on a pool of asyncmy connections.
+
+``make`` opens the pool and checks it with ``SELECT 1`` (within five
+seconds); statements outside a transaction each take a connection for
+themselves, a transaction or a stream keeps one until it ends::
+
+    program.pipe(fx.provide(mysql.layer(mysql.MysqlClientConfig(url=redacted.make(url)))))
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any, Never, Unpack
+
+from pyfect._core import Effect
+from pyfect._effect_args import require_effect
+from pyfect.config import Config, ConfigError, unwrap
+from pyfect.context import Context, Tag
+from pyfect.duration import to_seconds
+from pyfect.effect import fail
+from pyfect.layer import Layer, provide
+from pyfect.layer import effect_context as layer_effect_context
+from pyfect.reactivity import Reactivity
+from pyfect.reactivity import layer as reactivity_layer
+from pyfect.redacted import value as redacted_value
+from pyfect.scope import Scope, scoped
+from pyfect.sql.client import ClientFields, SqlClient, SqlClientService
+from pyfect.sql.client import make as make_sql_client
+from pyfect.sql.connection import Connection
+from pyfect.sql.error import ConnectionError, SqlError
+from pyfect.sql.mysql import _pool
+from pyfect.sql.mysql._asyncmy import connect_keywords
+from pyfect.sql.mysql._connection import MysqlConnection
+from pyfect.sql.mysql.config import (
+    DEFAULT_IDLE_TIMEOUT,
+    DEFAULT_MAX_CONNECTIONS,
+    MysqlClientConfig,
+    parse_url,
+)
+from pyfect.sql.mysql.error import classify
+from pyfect.sql.statement import Compiled, Compiler, default_escape, default_transforms
+from pyfect.sql.statement import make_compiler as make_statement_compiler
+
+_escape = default_escape("`")
+_CONNECT_TIMEOUT = 5.0
+
+
+class MysqlClientService(SqlClientService):
+    """A MySQL client: an ``SqlClientService`` and its config."""
+
+    __slots__ = ("config",)
+
+    config: MysqlClientConfig
+
+    def __init__(self, *, config: MysqlClientConfig, **fields: Unpack[ClientFields]) -> None:
+        super().__init__(**fields)
+        self.config = config
+
+
+class MysqlClient(Tag[MysqlClientService], key="@effect/sql-mysql2/MysqlClient"):
+    """The MySQL client service."""
+
+
+def _nothing(*_args: object) -> Compiled:
+    return ("", ())
+
+
+def make_compiler(transform: Callable[[str], str] | None = None) -> Compiler:
+    """MySQL's compiler: ``?`` placeholders and `` `quoted` `` names (renamed
+    with ``transform`` unless a statement runs without transforms)."""
+
+    def on_identifier(value: str, without_transform: bool) -> str:
+        if transform is None or without_transform:
+            return _escape(value)
+        return _escape(transform(value))
+
+    return make_statement_compiler(
+        dialect="mysql",
+        placeholder=lambda _index, _value: "?",
+        on_identifier=on_identifier,
+        on_record_update=_nothing,
+        on_custom=_nothing,
+    )
+
+
+def connect_options(config: MysqlClientConfig) -> dict[str, Any]:
+    """``asyncmy.connect``'s keywords for ``config``."""
+    if config.url is not None:
+        url = parse_url(redacted_value(config.url))
+        fields: dict[str, Any] = {
+            "host": url.host,
+            "port": url.port or 3306,
+            "user": url.username,
+            "password": url.password or "",
+            "database": url.database,
+        }
+        return connect_keywords(fields)
+    fields = {
+        **config.connect_options,
+        "host": config.host or config.connect_options.get("host", "localhost"),
+        "port": config.port or config.connect_options.get("port", 3306),
+    }
+    if config.database is not None:
+        fields["database"] = config.database
+    if config.username is not None:
+        fields["user"] = config.username
+    if config.password is not None:
+        fields["password"] = redacted_value(config.password)
+    return connect_keywords(fields)
+
+
+def _connect_failed(error: SqlError) -> SqlError:
+    cause = error.reason.cause
+    return SqlError(
+        reason=classify(cause, message="MysqlClient: Failed to connect", operation="connect")
+    )
+
+
+def _connect_timed_out() -> Effect[Never, SqlError]:
+    return fail(
+        SqlError(
+            reason=ConnectionError(
+                cause=Exception("connection timeout"),
+                message="MysqlClient: Connection timeout",
+                operation="connect",
+            )
+        )
+    )
+
+
+def _span_attributes(config: MysqlClientConfig) -> list[tuple[str, object]]:
+    attributes: list[tuple[str, object]] = [
+        *config.span_attributes.items(),
+        ("db.system.name", "mysql"),
+        ("server.address", config.host or "localhost"),
+        ("server.port", config.port or 3306),
+    ]
+    if config.database:
+        attributes.append(("db.namespace", config.database))
+    return attributes
+
+
+def make(config: MysqlClientConfig) -> Effect[MysqlClientService, SqlError, Scope | Reactivity]:
+    """A client on a pool of connections; it closes with the scope."""
+    prepare = not config.disable_prepared_statements
+    idle = (
+        DEFAULT_IDLE_TIMEOUT if config.connection_ttl is None else to_seconds(config.connection_ttl)
+    )
+    pool = _pool.make(
+        connect_options(config),
+        max_connections=config.max_connections or DEFAULT_MAX_CONNECTIONS,
+        idle_timeout=idle,
+    )
+
+    def build(
+        check_out: Effect[Any, SqlError, Scope],
+    ) -> Effect[MysqlClientService, SqlError, Reactivity]:
+        acquirer: Effect[Connection, SqlError, Scope] = check_out.map(
+            lambda session: MysqlConnection(session, prepare=prepare)
+        )
+
+        def borrow[A, E, R](
+            f: Callable[[Connection], Effect[A, E, R]],
+        ) -> Effect[A, E | SqlError, R]:
+            return scoped(acquirer.flat_map(f))
+
+        connected = (
+            scoped(
+                acquirer.flat_map(
+                    lambda conn: conn.execute_unprepared("SELECT 1", (), None)
+                ).map_error(_connect_failed)
+            )
+            .timeout_or_else(_CONNECT_TIMEOUT, _connect_timed_out)
+            .as_void()
+        )
+        transform_rows = (
+            None
+            if config.transform_result_names is None
+            else default_transforms(config.transform_result_names).array
+        )
+        client = make_sql_client(
+            driver="mysql2",
+            acquirer=acquirer,
+            borrower=borrow,
+            transaction_acquirer=acquirer,
+            release_savepoint=lambda name: f"RELEASE SAVEPOINT {name}",
+            compiler=make_compiler(config.transform_query_names),
+            span_attributes=_span_attributes(config),
+            transform_rows=transform_rows,
+        )
+        return connected.and_then(
+            client.map(lambda c: MysqlClientService(config=config, **c.fields()))
+        )
+
+    return pool.flat_map(build)
+
+
+def _contexts(client: MysqlClientService) -> Context[MysqlClient | SqlClient]:
+    return Context.empty().add(MysqlClient, client).add(SqlClient, client)
+
+
+def layer_from[E](
+    acquire: Effect[MysqlClientService, E, Scope | Reactivity],
+) -> Layer[MysqlClient | SqlClient, E]:
+    """Provides the client ``acquire`` makes as ``MysqlClient`` and
+    ``SqlClient``, with its own ``Reactivity``."""
+    require_effect(acquire, "acquire")
+    return layer_effect_context(acquire.map(_contexts)).pipe(provide(reactivity_layer))
+
+
+def layer(config: MysqlClientConfig) -> Layer[MysqlClient | SqlClient, SqlError]:
+    """Provides a client (``make``)."""
+    return layer_from(make(config))
+
+
+def layer_config(
+    config: Config[MysqlClientConfig] | Mapping[str, Any],
+) -> Layer[MysqlClient | SqlClient, ConfigError | SqlError]:
+    """``layer`` with the config read from the ``ConfigProvider``: a
+    ``Config`` of the whole config, or a mapping from field names to their
+    ``Config`` (Effect's ``Config.Wrap``)::
+
+        mysql.layer_config({"url": fx.config.redacted("DATABASE_URL")})
+    """
+    read: Effect[MysqlClientConfig, ConfigError] = (
+        config
+        if isinstance(config, Config)
+        else unwrap(config).map(lambda fields: MysqlClientConfig(**fields))
+    )
+    return layer_from(read.flat_map(make))
+
+
+__all__ = [
+    "MysqlClient",
+    "MysqlClientService",
+    "SqlError",
+    "connect_options",
+    "layer",
+    "layer_config",
+    "layer_from",
+    "make",
+    "make_compiler",
+]

pyfect/sql/mysql/config.py

@@ -0,0 +1,108 @@
+"""Settings of a MySQL client (Effect's ``@effect/sql-mysql2``
+``MysqlClientConfig``).
+
+``url`` (``mysql://user:password@host:port/database``) wins over every
+other connection field and over ``connect_options``, as in Effect; it is
+checked when the config is made. ``connect_options`` (Effect's
+``poolConfig``) are passed to ``asyncmy.connect`` as they are, for TLS
+(``ssl``), a ``unix_socket``, ``init_command`` and the rest.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from typing import Any
+from urllib.parse import unquote, urlsplit
+
+from pyfect.duration import Duration
+from pyfect.redacted import Redacted
+from pyfect.redacted import value as redacted_value
+
+DEFAULT_MAX_CONNECTIONS = 10
+"""mysql2's ``connectionLimit``."""
+
+DEFAULT_IDLE_TIMEOUT = 60.0
+"""mysql2's ``idleTimeout``, in seconds."""
+
+
+@dataclass(frozen=True, slots=True)
+class UrlParts:
+    """What a ``mysql://`` URL names."""
+
+    host: str
+    port: int | None
+    database: str | None
+    username: str | None
+    password: str | None
+
+
+def parse_url(url: str) -> UrlParts:
+    """The parts of a ``mysql://`` URL; ``ValueError`` when it is not
+    one."""
+    parts = urlsplit(url)
+    if parts.scheme != "mysql":
+        raise ValueError(f"expected a mysql:// URL, got scheme {parts.scheme!r}")
+    if parts.query or parts.fragment:
+        raise ValueError("MySQL URLs take no query or fragment; use connect_options")
+    if not parts.hostname:
+        raise ValueError("the MySQL URL names no host")
+    database = unquote(parts.path.removeprefix("/")) or None
+    return UrlParts(
+        host=parts.hostname,
+        port=parts.port,
+        database=database,
+        username=None if parts.username is None else unquote(parts.username),
+        password=None if parts.password is None else unquote(parts.password),
+    )
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class MysqlClientConfig:
+    """How to connect, and what the client does with rows and names.
+
+    - ``url``: a ``mysql://`` URL, overriding the connection fields below
+      and ``connect_options``
+    - ``host`` (``localhost``) / ``port`` (3306) / ``database`` /
+      ``username`` / ``password``
+    - ``max_connections``: the pool's size (10)
+    - ``connection_ttl``: an unused connection closes after this long (60
+      seconds)
+    - ``connect_options``: more ``asyncmy.connect`` keywords (Effect's
+      ``poolConfig``)
+    - ``disable_prepared_statements``: send statements as text with their
+      parameters escaped in, for proxies without ``COM_STMT_PREPARE``
+      (Cloudflare Hyperdrive)
+    - ``span_attributes``: added to every statement's span
+    - ``transform_result_names`` / ``transform_query_names``: rename result
+      columns / identifiers (``fx.case`` has the usual ones)
+    """
+
+    url: Redacted[str] | None = None
+    host: str | None = None
+    port: int | None = None
+    database: str | None = None
+    username: str | None = None
+    password: Redacted[str] | None = None
+    max_connections: int | None = None
+    connection_ttl: Duration | None = None
+    connect_options: Mapping[str, Any] = field(default_factory=dict[str, Any])
+    disable_prepared_statements: bool = False
+    span_attributes: Mapping[str, object] = field(default_factory=dict[str, object])
+    transform_result_names: Callable[[str], str] | None = None
+    transform_query_names: Callable[[str], str] | None = None
+
+    def __post_init__(self) -> None:
+        if self.url is not None:
+            parse_url(redacted_value(self.url))
+        if self.max_connections is not None and self.max_connections < 1:
+            raise ValueError(f"max_connections must be >= 1, got {self.max_connections}")
+
+
+__all__ = [
+    "DEFAULT_IDLE_TIMEOUT",
+    "DEFAULT_MAX_CONNECTIONS",
+    "MysqlClientConfig",
+    "UrlParts",
+    "parse_url",
+]