diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..56ef81c --- /dev/null +++ b/.dockerignore @@ -0,0 +1,13 @@ +.git +.venv +venv +dist +build +*.egg-info +**/__pycache__ +**/*.pyc +.pytest_cache +docs/.build +.vscode +.idea +.DS_Store diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml new file mode 100644 index 0000000..519d028 --- /dev/null +++ b/.github/workflows/slo.yml @@ -0,0 +1,194 @@ +name: SLO + +on: + pull_request: + types: [labeled] + +permissions: + contents: read + pull-requests: write + checks: write + +jobs: + workload: + if: github.event.label.name == 'SLO' + + name: SLO workload (${{ matrix.workload.name }}) + runs-on: "large-runner-sqlalchemy" + + strategy: + fail-fast: false + matrix: + workload: + - name: core + command: "--read-rps 1000 --write-rps 100" + thresholds: "" + - name: orm + command: "--read-rps 1000 --write-rps 100" + thresholds: "" + - name: tx + command: "--read-rps 1000 --write-rps 100" + # Interactive transactions legitimately retry (chaos / lock conflicts), + # so retry_attempts is informational for this scenario only. + thresholds: | + metrics: + - name: read_retry_attempts + direction: lower_is_better + warning_change_percent: 100000000 + critical_change_percent: 100000000 + - name: write_retry_attempts + direction: lower_is_better + warning_change_percent: 100000000 + critical_change_percent: 100000000 + + concurrency: + group: slo-${{ github.ref }}-${{ matrix.workload.name }} + cancel-in-progress: true + + steps: + - name: Install tooling + run: | + set -euxo pipefail + YQ_VERSION=v4.48.2 + BUILDX_VERSION=0.30.1 + COMPOSE_VERSION=2.40.3 + + # Resilient download: retry transient network errors (e.g. curl 28 + # SSL connection timeout) instead of failing the whole job. + dl() { + sudo curl -fL \ + --retry 5 --retry-all-errors --retry-delay 3 \ + --connect-timeout 30 --max-time 300 \ + -o "$1" "$2" + } + + dl /usr/local/bin/yq \ + "https://github.com/mikefarah/yq/releases/download/${YQ_VERSION}/yq_linux_amd64" + sudo chmod +x /usr/local/bin/yq + + sudo mkdir -p /usr/local/lib/docker/cli-plugins + + dl /usr/local/lib/docker/cli-plugins/docker-buildx \ + "https://github.com/docker/buildx/releases/download/v${BUILDX_VERSION}/buildx-v${BUILDX_VERSION}.linux-amd64" + sudo chmod +x /usr/local/lib/docker/cli-plugins/docker-buildx + + dl /usr/local/lib/docker/cli-plugins/docker-compose \ + "https://github.com/docker/compose/releases/download/v${COMPOSE_VERSION}/docker-compose-linux-x86_64" + sudo chmod +x /usr/local/lib/docker/cli-plugins/docker-compose + + yq --version + docker --version + docker buildx version + docker compose version + + - name: Checkout current dialect version + uses: actions/checkout@v5 + with: + path: current + fetch-depth: 0 + + - name: Determine baseline commit + id: baseline + working-directory: current + run: | + set -euo pipefail + BASELINE=$(git merge-base HEAD origin/main) + echo "sha=${BASELINE}" >> "$GITHUB_OUTPUT" + if [ "$(git rev-parse origin/main)" = "${BASELINE}" ]; then + echo "ref=main" >> "$GITHUB_OUTPUT" + else + echo "ref=${BASELINE:0:7}" >> "$GITHUB_OUTPUT" + fi + + - name: Checkout baseline dialect version + uses: actions/checkout@v5 + with: + ref: ${{ steps.baseline.outputs.sha }} + path: baseline + fetch-depth: 1 + + - name: Build workload images (current + baseline) + run: | + set -euxo pipefail + + # Current: PR dialect + PR workload runner. + docker build \ + -f "$GITHUB_WORKSPACE/current/tests/slo/Dockerfile" \ + -t "ydb-slo-current" \ + "$GITHUB_WORKSPACE/current" + + # Baseline: baseline dialect with the CURRENT workload runner so the + # runner-side contract (entrypoint, metric names) is identical for both. + # main has no tests/ dir, so create the parent before copying the runner. + rm -rf "$GITHUB_WORKSPACE/baseline/tests/slo" + mkdir -p "$GITHUB_WORKSPACE/baseline/tests" + cp -r "$GITHUB_WORKSPACE/current/tests/slo" "$GITHUB_WORKSPACE/baseline/tests/slo" + docker build \ + -f "$GITHUB_WORKSPACE/baseline/tests/slo/Dockerfile" \ + -t "ydb-slo-baseline" \ + "$GITHUB_WORKSPACE/baseline" + + - name: Run SLO workload + uses: ydb-platform/ydb-slo-action/init@feat/per-scenario-thresholds + timeout-minutes: 30 + with: + github_issue: ${{ github.event.pull_request.number }} + github_token: ${{ secrets.GITHUB_TOKEN }} + workload_name: ${{ matrix.workload.name }} + workload_duration: "600" + workload_current_ref: ${{ github.head_ref || github.ref_name }} + workload_current_image: ydb-slo-current + workload_current_command: ${{ matrix.workload.command }} + workload_baseline_ref: ${{ steps.baseline.outputs.ref }} + workload_baseline_image: ydb-slo-baseline + workload_baseline_command: ${{ matrix.workload.command }} + # The workload emits only real retries (attempts beyond the first), so + # count them directly instead of the default "attempts - operations" + # (which would read 0 for every scenario under the new emission). + metrics_yaml: | + metrics: + - name: read_retry_attempts + query: | + sum by(ref) (increase(sdk_retry_attempts_total{operation_type="read"}[5s])) + - name: write_retry_attempts + query: | + sum by(ref) (increase(sdk_retry_attempts_total{operation_type="write"}[5s])) + thresholds_yaml: ${{ matrix.workload.thresholds }} + + report: + # Runs in the same workflow run as `workload`, so it sees the freshly + # uploaded metrics artifacts and can compare current vs baseline and gate + # the PR. Same-repo PRs have the pull-requests:write token this needs. + if: ${{ always() && github.event.label.name == 'SLO' }} + + name: SLO report + needs: workload + runs-on: ubuntu-latest + + steps: + - name: Publish SLO report + uses: ydb-platform/ydb-slo-action/report@feat/per-scenario-thresholds + with: + github_issue: ${{ github.event.pull_request.number }} + github_token: ${{ secrets.GITHUB_TOKEN }} + github_run_id: ${{ github.run_id }} + + remove-label: + # SLO runs once per label: drop the "SLO" label after the run so ordinary + # pushes don't re-trigger the (heavy) SLO suite. Re-add the label to re-run. + if: ${{ always() && github.event.label.name == 'SLO' }} + + name: Remove SLO label + needs: [workload, report] + runs-on: ubuntu-latest + permissions: + pull-requests: write + + steps: + - name: Remove SLO label + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + gh api -X DELETE \ + "repos/${{ github.repository }}/issues/${{ github.event.pull_request.number }}/labels/SLO" \ + && echo "Removed SLO label" || echo "SLO label already absent" diff --git a/docs/advanced.rst b/docs/advanced.rst index 9df13b5..35d801d 100644 --- a/docs/advanced.rst +++ b/docs/advanced.rst @@ -40,3 +40,52 @@ Pass ``_add_declare_for_yql_stmt_vars=True`` to :func:`sqlalchemy.create_engine` ) with engine.connect() as conn: conn.execute(sa.text("SELECT :id"), {"id": 1}) # runs as "DECLARE `$id` as Int64;\nSELECT $id" with param + +Retrying operations +------------------- + +YDB returns retryable errors (``Unavailable``, ``Overloaded``, ``Aborted``, +``BadSession``, ...) that the SDK knows how to retry. SQLAlchemy, however, wraps +the underlying driver error in :class:`sqlalchemy.exc.DBAPIError`, which the SDK +retry logic does not recognise. The helpers in ``ydb_sqlalchemy`` translate the +SQLAlchemy error back to the original ``ydb.Error`` and run the operation under +the SDK's retry policy: + +.. code-block:: python + + import sqlalchemy as sa + from ydb_sqlalchemy import retry_ydb_operation + + engine = sa.create_engine("yql+ydb://localhost:2136/local") + + def read_modify_write(): + with engine.begin() as conn: + value = conn.execute(sa.text("SELECT v FROM t WHERE id = 1")).scalar() + conn.execute(sa.text("UPDATE t SET v = :v WHERE id = 1"), {"v": value + 1}) + + # idempotent=True also retries ambiguous errors; set it only when re-running + # the whole callable is safe. + retry_ydb_operation(read_modify_write, max_retries=10, idempotent=True) + +The callable is re-run from scratch on every attempt, so it must open its own +connection/transaction and not rely on earlier state. Only ``ydb.Error`` is +retried; any other error propagates unchanged. + +There is also a decorator form that works on both sync and async functions, and +an async call form :func:`ydb_sqlalchemy.retry_ydb_operation_async`: + +.. code-block:: python + + from ydb_sqlalchemy import retry_ydb + + @retry_ydb(idempotent=True) + def read_modify_write(): + ... + +.. note:: + + In the default ``AUTOCOMMIT`` isolation level each statement is its own + transaction and is already retried inside ``ydb-dbapi``, so the helper is + mainly useful for **interactive transactions** (a multi-statement + read-modify-write under ``SERIALIZABLE``/``SNAPSHOT``), where the whole + transaction must be retried as a unit. diff --git a/setup.cfg b/setup.cfg index 650a160..5fec16e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,7 +3,7 @@ addopts= --tb native -v -r fxX -p no:warnings [sqla_testing] requirement_cls=ydb_sqlalchemy.sqlalchemy.requirements:Requirements -profile_file=test/profiles.txt +profile_file=tests/integration/profiles.txt [db] default=yql+ydb://localhost:2136/local diff --git a/test/__init__.py b/tests/integration/__init__.py similarity index 100% rename from test/__init__.py rename to tests/integration/__init__.py diff --git a/test/conftest.py b/tests/integration/conftest.py similarity index 100% rename from test/conftest.py rename to tests/integration/conftest.py diff --git a/test/test_core.py b/tests/integration/test_core.py similarity index 100% rename from test/test_core.py rename to tests/integration/test_core.py diff --git a/test/test_inspect.py b/tests/integration/test_inspect.py similarity index 100% rename from test/test_inspect.py rename to tests/integration/test_inspect.py diff --git a/test/test_orm.py b/tests/integration/test_orm.py similarity index 100% rename from test/test_orm.py rename to tests/integration/test_orm.py diff --git a/test/test_suite.py b/tests/integration/test_suite.py similarity index 100% rename from test/test_suite.py rename to tests/integration/test_suite.py diff --git a/tests/slo/Dockerfile b/tests/slo/Dockerfile new file mode 100644 index 0000000..5276efd --- /dev/null +++ b/tests/slo/Dockerfile @@ -0,0 +1,39 @@ +# syntax=docker/dockerfile:1 +# +# Packages the SQLAlchemy SLO workload runner for ydb-slo-action v2. +# +# Build context is the repository root. The action builds this image twice — for +# the `current` (PR) and `baseline` (merge-base) versions of the dialect — and +# runs both in parallel against the same YDB cluster, distinguishing them by the +# WORKLOAD_REF metric label (see tests/slo/docker-entrypoint.sh). +# +# Notes: +# - OpenTelemetry 1.39.x requires Python >= 3.9. +# - gcc/libc6-dev are needed to build hdrhistogram from sdist. + +FROM python:3.11-slim + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends gcc libc6-dev \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /src + +# 1. ydb-sqlalchemy dialect (current or baseline source tree). +COPY setup.py pyproject.toml requirements.txt README.md ./ +COPY ydb_sqlalchemy/ ydb_sqlalchemy/ +RUN pip install --no-cache-dir . + +# 2. SLO runner deps. +COPY tests/slo/requirements.txt tests/slo/requirements.txt +RUN pip install --no-cache-dir -r tests/slo/requirements.txt + +# 3. Workload source + entrypoint. +COPY tests/slo/src /src/tests/slo/src +COPY tests/slo/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh +RUN chmod +x /usr/local/bin/docker-entrypoint.sh + +ENTRYPOINT ["/usr/local/bin/docker-entrypoint.sh"] diff --git a/tests/slo/README.md b/tests/slo/README.md new file mode 100644 index 0000000..8600979 --- /dev/null +++ b/tests/slo/README.md @@ -0,0 +1,89 @@ +# SQLAlchemy SLO workload + +A load generator that exercises the `ydb-sqlalchemy` dialect under +[`ydb-platform/ydb-slo-action`](https://github.com/ydb-platform/ydb-slo-action). +It keeps reader and writer threads busy against a single key/value table and +reports latency / throughput / availability to Prometheus over OTLP, so the +action can compare the current (PR) dialect against a baseline and gate the PR +on regressions. + +## What it does + +Reader and writer threads run in parallel from dedicated pools. Three execution +modes (selected by `WORKLOAD_NAME` / `--mode`) exercise the ways real +applications use the dialect: + +| mode | read path | write path | +|--------|-----------------------------------|--------------------------------------------------| +| `core` | `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | +| `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | +| `tx` | `select()` in a SERIALIZABLE tx | read-modify-write (`select()` + `upsert()`) in a SERIALIZABLE tx | + +`core` and `orm` issue **single autocommit statements**, which ydb-dbapi already +retries internally, so no app-level retry is used. `tx` runs **interactive +transactions**, which ydb-dbapi does *not* retry on its own, so each transaction +is wrapped in `ydb_sqlalchemy.retry_ydb_operation` — the dialect helper that +translates a SQLAlchemy error back to the underlying `ydb.Error` and retries the +transient ones; the workload counts the attempts, so `sdk_retry_attempts_total` +reflects real retries in that mode. + +`core` writes use a monotonic id and UPSERT (idempotent, safe for the shared +table); `orm` writes use a random id so each `session.add` INSERT is collision-free +and avoids a hot last partition; `tx` writes increment an existing row, which is +why the read-modify-write must be one transaction (and why it is retried with +`idempotent=False`). + +## Layout + +``` +tests/slo/ +├── Dockerfile # image used by ydb-slo-action (build context = repo root) +├── docker-entrypoint.sh # create (idempotent) then run, honouring injected env +├── requirements.txt # hdrhistogram + opentelemetry +└── src/ + ├── __main__.py # entrypoint: python ./tests/slo/src ... + ├── options.py # argparse (create / run / cleanup) + ├── models.py # engine factory, table + imperatively-mapped ORM entity + ├── generator.py # row payload generator + ├── metrics.py # OTLP metrics (names match the action's metrics.yaml) + └── workload.py # rate limiting, retries, parallel read/write jobs +``` + +## CLI + +```bash +# Create the table and fill it with initial rows. +python ./tests/slo/src create grpc://localhost:2136 /local --mode core + +# Run the parallel read/write load for 60s. +python ./tests/slo/src run grpc://localhost:2136 /local \ + --mode core --time 60 --read-rps 500 --write-rps 50 \ + --otlp-endpoint http://localhost:9090/api/v1/otlp/v1/metrics + +# Drop the table. +python ./tests/slo/src cleanup grpc://localhost:2136 /local +``` + +When run by the action, `YDB_ENDPOINT`, `YDB_DATABASE`, `WORKLOAD_DURATION`, +`WORKLOAD_NAME`, `WORKLOAD_REF` and `OTEL_EXPORTER_OTLP_METRICS_ENDPOINT` are +injected automatically; the metrics endpoint is picked up from the environment. + +## Metrics + +Emitted via OTLP/HTTP and consumed by the action's default `metrics.yaml`: + +| Prometheus name | type | labels | +|-------------------------------------|---------|------------------------------------------| +| `sdk_operations_total` | counter | `ref`, `operation_type`, `operation_status` | +| `sdk_operations_success_total` | counter | `ref`, `operation_type` | +| `sdk_operations_failure_total` | counter | `ref`, `operation_type` | +| `sdk_retry_attempts_total` | counter | `ref`, `operation_type` | +| `sdk_operation_latency_p{50,95,99}_seconds` | gauge | `ref`, `operation_type`, `operation_status` | + +## CI + +`.github/workflows/slo.yml` runs the workload on pull requests labelled `SLO`. +It builds the current and baseline images, hands them to `ydb-slo-action/init`, +then publishes a comparison report with `ydb-slo-action/report`. The workload +job runs on the `large-runner-sqlalchemy` self-hosted runner with the full YDB +cluster; the report job runs on `ubuntu-latest`. diff --git a/tests/slo/docker-entrypoint.sh b/tests/slo/docker-entrypoint.sh new file mode 100644 index 0000000..9b926fb --- /dev/null +++ b/tests/slo/docker-entrypoint.sh @@ -0,0 +1,34 @@ +#!/bin/sh +# Workload entrypoint used by ydb-slo-action v2. +# +# The action launches the same image for both `current` and `baseline` workloads +# in parallel; both prepare the schema (idempotent) and then run the load. +# +# Env vars injected by the action: +# WORKLOAD_NAME core | orm (selects the SQLAlchemy layer) +# WORKLOAD_REF current / (used as the `ref` metric label) +# WORKLOAD_DURATION run duration in seconds +# YDB_ENDPOINT grpc://ydb:2136 +# YDB_DATABASE /Root/testdb +# OTEL_EXPORTER_OTLP_METRICS_ENDPOINT Prometheus OTLP receiver +# +# Anything passed after the script name is appended to the `run` command — this +# is how tuning flags from `workload_current_command` (e.g. --read-rps) arrive. + +set -e + +ENDPOINT="${YDB_ENDPOINT:-grpc://localhost:2136}" +DATABASE="${YDB_DATABASE:-/local}" +DURATION="${WORKLOAD_DURATION:-60}" +MODE="${WORKLOAD_NAME:-core}" + +echo "SLO workload: mode=${MODE} ref=${WORKLOAD_REF:-current} endpoint=${ENDPOINT} db=${DATABASE} duration=${DURATION}s" + +# Tolerate a parallel container winning the create race. +python ./tests/slo/src create "$ENDPOINT" "$DATABASE" --mode "$MODE" \ + || echo "WARN: create exited non-zero (treated as already-prepared)" >&2 + +exec python ./tests/slo/src run "$ENDPOINT" "$DATABASE" \ + --mode "$MODE" \ + --time "$DURATION" \ + "$@" diff --git a/tests/slo/requirements.txt b/tests/slo/requirements.txt new file mode 100644 index 0000000..f7ace45 --- /dev/null +++ b/tests/slo/requirements.txt @@ -0,0 +1,9 @@ +# Extra deps for the SLO workload runner. +# SQLAlchemy, ydb and ydb-dbapi come from installing the ydb-sqlalchemy dialect. + +hdrhistogram==0.10.3 + +# OpenTelemetry (OTLP/HTTP exporter). 1.39.x requires Python >= 3.9. +opentelemetry-api==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-exporter-otlp-proto-http==1.39.1 diff --git a/tests/slo/src/__main__.py b/tests/slo/src/__main__.py new file mode 100644 index 0000000..ef9959c --- /dev/null +++ b/tests/slo/src/__main__.py @@ -0,0 +1,12 @@ +import logging + +from options import parse_options +from workload import run_from_args + +if __name__ == "__main__": + args = parse_options() + logging.basicConfig( + level=logging.DEBUG if args.debug else logging.INFO, + format="%(asctime)s %(levelname)-8s %(threadName)s %(message)s", + ) + run_from_args(args) diff --git a/tests/slo/src/generator.py b/tests/slo/src/generator.py new file mode 100644 index 0000000..9137b71 --- /dev/null +++ b/tests/slo/src/generator.py @@ -0,0 +1,68 @@ +"""Row payload generator shared by the create/write paths.""" + +import random +import string +import threading +from dataclasses import dataclass +from datetime import datetime, timezone + +MAX_UINT64 = 2**64 - 1 + + +def _random_string(min_len: int, max_len: int) -> str: + length = random.randint(min_len, max_len) + return "".join(random.choices(string.ascii_lowercase, k=length)) + + +@dataclass +class Row: + object_id: int + payload_str: str + payload_double: float + payload_timestamp: datetime + + def as_params(self) -> dict: + return { + "object_id": self.object_id, + "payload_str": self.payload_str, + "payload_double": self.payload_double, + "payload_timestamp": self.payload_timestamp, + } + + +class RowGenerator: + """Thread-safe generator of rows with monotonically increasing ids.""" + + def __init__(self, start_id: int = 0) -> None: + self._id = start_id + self._lock = threading.Lock() + + def _next_id(self) -> int: + with self._lock: + self._id += 1 + if self._id >= MAX_UINT64: + self._id = 1 + return self._id + + def get(self) -> Row: + return Row( + object_id=self._next_id(), + payload_str=_random_string(20, 40), + payload_double=random.random(), + payload_timestamp=datetime.now(timezone.utc), + ) + + +def random_row() -> Row: + """A row with a random 63-bit id. + + Used for ORM inserts: a random id keeps each ``session.add`` collision-free + without cross-thread coordination, and avoids the hot last partition that a + monotonically increasing primary key would create in YDB. + """ + return Row( + object_id=random.getrandbits(63) + 1, + payload_str=_random_string(20, 40), + payload_double=random.random(), + payload_timestamp=datetime.now(timezone.utc), + ) diff --git a/tests/slo/src/metrics.py b/tests/slo/src/metrics.py new file mode 100644 index 0000000..f184dc8 --- /dev/null +++ b/tests/slo/src/metrics.py @@ -0,0 +1,240 @@ +""" +OTLP metrics for the SQLAlchemy SLO workload. + +The instrument names here are chosen so that, once exported through the +Prometheus OTLP receiver, they line up with the queries shipped in +ydb-slo-action's default ``metrics.yaml``: + + * ``sdk_operations_total`` (counter, labels: ref, operation_type, operation_status) + * ``sdk_operations_success_total`` (counter) + * ``sdk_operations_failure_total`` (counter) + * ``sdk_retry_attempts_total`` (counter, labels: ref, operation_type) + * ``sdk_operation_latency_p50_seconds`` / ``_p95_`` / ``_p99_`` (gauge) + +Latency percentiles are computed client-side per push window via HdrHistogram +and emitted as gauges; the histogram is reset after every push so each sample +describes only the most recent window. +""" + +import logging +import threading +import time +from abc import ABC, abstractmethod +from contextlib import contextmanager +from os import environ +from typing import Optional + +OP_TYPE_READ, OP_TYPE_WRITE = "read", "write" +OP_STATUS_SUCCESS, OP_STATUS_FAILURE = "success", "error" + +REF = environ.get("WORKLOAD_REF") or environ.get("REF") or "current" +WORKLOAD = environ.get("WORKLOAD_NAME") or environ.get("WORKLOAD") or "core" + +logger = logging.getLogger(__name__) + + +def _sdk_version() -> str: + try: + from importlib.metadata import version + + return version("ydb-sqlalchemy") + except Exception: + return "0.0.0" + + +class BaseMetrics(ABC): + @abstractmethod + def start(self, op_type: str) -> float: ... + + @abstractmethod + def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Optional[BaseException] = None) -> None: ... + + @abstractmethod + def push(self) -> None: ... + + @abstractmethod + def reset(self) -> None: ... + + @contextmanager + def measure(self, op_type: str): + start_ts = self.start(op_type) + error = None + try: + yield self + except Exception as err: + error = err + raise + finally: + self.stop(op_type, start_ts, error=error) + + +class DummyMetrics(BaseMetrics): + """No-op metrics used when no OTLP endpoint is configured (local runs).""" + + def start(self, op_type: str) -> float: + return time.time() + + def stop(self, op_type, start_time, attempts=1, error=None) -> None: + return None + + def push(self) -> None: + return None + + def reset(self) -> None: + return None + + +class OtlpMetrics(BaseMetrics): + _HDR_MIN_US = 1 + _HDR_MAX_US = 60_000_000 # 60s + _HDR_SIG_FIGS = 3 + _PERCENTILES = (("p50", 50.0), ("p95", 95.0), ("p99", 99.0)) + + def __init__(self, otlp_metrics_endpoint: str): + from hdrh.histogram import HdrHistogram + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + from opentelemetry.sdk.resources import Resource + + self._HdrHistogram = HdrHistogram + + resource = Resource.create( + { + "service.name": f"workload-{WORKLOAD}", + "service.instance.id": environ.get("SLO_INSTANCE_ID", f"{REF}-{WORKLOAD}"), + "ref": REF, + "sdk": "ydb-sqlalchemy", + "sdk_version": _sdk_version(), + "workload": WORKLOAD, + } + ) + + exporter = OTLPMetricExporter(endpoint=otlp_metrics_endpoint) + reader = PeriodicExportingMetricReader(exporter) + self._provider = MeterProvider(resource=resource, metric_readers=[reader]) + self._meter = self._provider.get_meter("ydb-sqlalchemy-slo") + + self._errors = self._meter.create_counter( + name="sdk.errors.total", + description="Total number of errors, categorized by error type.", + ) + self._operations_total = self._meter.create_counter( + name="sdk.operations.total", + description="Total number of operations attempted.", + ) + self._operations_success_total = self._meter.create_counter( + name="sdk.operations.success.total", + description="Total number of successful operations.", + ) + self._operations_failure_total = self._meter.create_counter( + name="sdk.operations.failure.total", + description="Total number of failed operations.", + ) + self._retry_attempts_total = self._meter.create_counter( + name="sdk.retry.attempts.total", + description="Total number of retries (attempts beyond the first one).", + ) + self._pending = self._meter.create_up_down_counter( + name="sdk.pending.operations", + description="Current number of in-flight operations.", + ) + self._latency_gauges = { + name: self._meter.create_gauge( + name=f"sdk.operation.latency.{name}.seconds", + unit="s", + description=f"Operation latency {name} over the last push window.", + ) + for name, _ in self._PERCENTILES + } + + self._lock = threading.Lock() + self._hdr: dict = {} + + def _get_hdr(self, op_type: str, op_status: str): + key = (op_type, op_status) + hist = self._hdr.get(key) + if hist is None: + hist = self._HdrHistogram(self._HDR_MIN_US, self._HDR_MAX_US, self._HDR_SIG_FIGS) + self._hdr[key] = hist + return hist + + def start(self, op_type: str) -> float: + self._pending.add(1, attributes={"ref": REF, "operation_type": op_type}) + return time.time() + + def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Optional[BaseException] = None) -> None: + duration = time.time() - start_time + duration_us = min(max(int(duration * 1_000_000), self._HDR_MIN_US), self._HDR_MAX_US) + + op_status = OP_STATUS_SUCCESS if error is None else OP_STATUS_FAILURE + base_attrs = {"ref": REF, "operation_type": op_type} + op_attrs = {**base_attrs, "operation_status": op_status} + + # Count retries (attempts beyond the first), not total attempts: this is + # exactly 0 when nothing was retried, so the *_retry_attempts metric is a + # clean 0 for autocommit workloads instead of PromQL noise around zero. + self._retry_attempts_total.add(max(int(attempts) - 1, 0), attributes=base_attrs) + self._pending.add(-1, attributes=base_attrs) + self._operations_total.add(1, attributes=op_attrs) + + if error is not None: + self._errors.add(1, attributes={**base_attrs, "error_type": type(error).__name__}) + self._operations_failure_total.add(1, attributes=base_attrs) + else: + self._operations_success_total.add(1, attributes=base_attrs) + + with self._lock: + self._get_hdr(op_type, op_status).record_value(duration_us) + + def push(self) -> None: + with self._lock: + for (op_type, op_status), hist in self._hdr.items(): + if hist.get_total_count() == 0: + continue + attrs = {"ref": REF, "operation_type": op_type, "operation_status": op_status} + for name, percentile in self._PERCENTILES: + value_s = hist.get_value_at_percentile(percentile) / 1_000_000 + self._latency_gauges[name].set(value_s, attributes=attrs) + for hist in self._hdr.values(): + hist.reset() + self._provider.force_flush() + + def reset(self) -> None: + with self._lock: + for hist in self._hdr.values(): + hist.reset() + self._provider.force_flush() + + +def _resolve_metrics_endpoint(cli_endpoint: Optional[str]) -> str: + """ + Resolution order: + 1. OTEL_EXPORTER_OTLP_METRICS_ENDPOINT (used as-is) + 2. OTEL_EXPORTER_OTLP_ENDPOINT + /v1/metrics suffix + 3. CLI --otlp-endpoint + """ + metrics_env = environ.get("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", "").strip() + if metrics_env: + return metrics_env + + base_env = environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "").strip() + if base_env: + base = base_env.rstrip("/") + return base if base.endswith("/v1/metrics") else f"{base}/v1/metrics" + + return (cli_endpoint or "").strip() + + +def create_metrics(otlp_endpoint: Optional[str]) -> BaseMetrics: + endpoint = _resolve_metrics_endpoint(otlp_endpoint) + if not endpoint: + logger.info("Metrics disabled (no OTLP endpoint); using DummyMetrics") + return DummyMetrics() + + logger.info("Exporting metrics via OTLP to: %s", endpoint) + try: + return OtlpMetrics(endpoint) + except Exception: + logger.exception("Failed to init OTLP metrics; falling back to DummyMetrics") + return DummyMetrics() diff --git a/tests/slo/src/models.py b/tests/slo/src/models.py new file mode 100644 index 0000000..63d13d3 --- /dev/null +++ b/tests/slo/src/models.py @@ -0,0 +1,93 @@ +""" +SQLAlchemy plumbing for the SLO workload: engine factory, the key/value table +definition and an imperatively-mapped ORM class for the ``orm`` workload mode. +""" + +import re + +import sqlalchemy as sa +from sqlalchemy.engine import URL +from sqlalchemy.orm import registry + +from ydb_sqlalchemy import types as ydb_types + +_ENDPOINT_RE = re.compile(r"^(?Pgrpcs?|grpc)://(?P[^:/]+):(?P\d+)") + + +def build_url(endpoint: str, database: str) -> URL: + """ + Turn a YDB endpoint (``grpc://ydb:2136``) and database path (``/Root/testdb``) + into a ``yql+ydb`` SQLAlchemy URL. + """ + match = _ENDPOINT_RE.match(endpoint.strip()) + if not match: + raise ValueError(f"Cannot parse YDB endpoint: {endpoint!r}") + + scheme = match.group("scheme") + host = match.group("host") + port = int(match.group("port")) + + # The dialect re-adds the leading slash to the database, so strip it here to + # avoid a doubled slash in the rendered URL. + db = database.strip().lstrip("/") + + query = {"protocol": "grpcs"} if scheme == "grpcs" else {} + return URL.create("yql+ydb", host=host, port=port, database=db, query=query) + + +def build_engine(endpoint: str, database: str, pool_size: int) -> sa.engine.Engine: + url = build_url(endpoint, database) + # No pool_pre_ping: it would add a SELECT 1 round-trip on every checkout. + # ydb-dbapi already retries transient errors and re-acquires sessions from + # its internal session pool, so broken connections recover without a ping. + return sa.create_engine( + url, + pool_size=pool_size, + max_overflow=max(4, pool_size), + ) + + +def build_table( + metadata: sa.MetaData, + table_name: str, + *, + min_partitions: int = 6, + max_partitions: int = 100, + partition_size_mb: int = 100, +) -> sa.Table: + return sa.Table( + table_name, + metadata, + sa.Column("object_id", ydb_types.UInt64, primary_key=True), + sa.Column("payload_str", sa.Unicode), + sa.Column("payload_double", sa.Float), + sa.Column("payload_timestamp", sa.TIMESTAMP), + ydb_auto_partitioning_by_size=True, + ydb_auto_partitioning_by_load=True, + ydb_auto_partitioning_min_partitions_count=min_partitions, + ydb_auto_partitioning_max_partitions_count=max_partitions, + ydb_auto_partitioning_partition_size_mb=partition_size_mb, + ) + + +class KeyValueRow: + """Plain class mapped imperatively onto the workload table for ``orm`` mode.""" + + def __init__(self, object_id, payload_str=None, payload_double=None, payload_timestamp=None): + self.object_id = object_id + self.payload_str = payload_str + self.payload_double = payload_double + self.payload_timestamp = payload_timestamp + + +_mapper_registry = registry() +_mapped = False + + +def ensure_mapped(table: sa.Table) -> type: + """Map :class:`KeyValueRow` onto ``table`` exactly once.""" + global _mapped + if not _mapped: + _mapper_registry.map_imperatively(KeyValueRow, table) + _mapped = True + return KeyValueRow diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py new file mode 100644 index 0000000..013431a --- /dev/null +++ b/tests/slo/src/options.py @@ -0,0 +1,64 @@ +import argparse +from os import environ + +_MODES = ("core", "orm", "tx") +_DEFAULT_MODE = (environ.get("WORKLOAD_NAME") or "core").strip().lower() +if _DEFAULT_MODE not in _MODES: + _DEFAULT_MODE = "core" + + +def _add_common(parser): + parser.add_argument("endpoint", help="YDB endpoint, e.g. grpc://localhost:2136") + parser.add_argument("database", help="YDB database path, e.g. /local") + parser.add_argument("-t", "--table-name", default="slo_sqlalchemy", help="Workload table name") + parser.add_argument( + "--mode", + choices=_MODES, + default=_DEFAULT_MODE, + help="Access pattern: core (Connection), orm (Session) or " + "tx (interactive read-modify-write transactions). Defaults to WORKLOAD_NAME.", + ) + parser.add_argument("--debug", action="store_true", help="Enable debug logging") + + +def _add_create(subparsers): + parser = subparsers.add_parser("create", help="Create the table and fill it with initial rows") + _add_common(parser) + parser.add_argument("-c", "--initial-data-count", default=1000, type=int, help="Initial row count") + parser.add_argument("--batch-size", default=100, type=int, help="Rows per insert batch") + parser.add_argument("-p-min", "--min-partitions-count", default=6, type=int, help="Min partitions") + parser.add_argument("-p-max", "--max-partitions-count", default=100, type=int, help="Max partitions") + parser.add_argument("-p-size", "--partition-size", default=100, type=int, help="Partition size [MB]") + + +def _add_run(subparsers): + parser = subparsers.add_parser("run", help="Run the parallel read/write SLO load") + _add_common(parser) + parser.add_argument("--read-rps", default=500, type=int, help="Target read RPS") + parser.add_argument("--write-rps", default=50, type=int, help="Target write RPS") + parser.add_argument("--read-threads", default=8, type=int, help="Reader threads") + parser.add_argument("--write-threads", default=4, type=int, help="Writer threads") + parser.add_argument("--initial-data-count", default=1000, type=int, help="Fallback id space when table is empty") + parser.add_argument("--time", default=60, type=int, help="Run duration [s]") + parser.add_argument("--shutdown-time", default=10, type=int, help="Graceful shutdown time [s]") + parser.add_argument( + "--otlp-endpoint", + default="", + type=str, + help="OTLP metrics endpoint; empty disables metrics unless OTEL_* env vars are set", + ) + parser.add_argument("--report-period", default=1000, type=int, help="Metrics push period [ms]") + + +def _add_cleanup(subparsers): + parser = subparsers.add_parser("cleanup", help="Drop the workload table") + _add_common(parser) + + +def parse_options(): + parser = argparse.ArgumentParser(description="YDB SQLAlchemy SLO workload") + subparsers = parser.add_subparsers(dest="command", required=True) + _add_create(subparsers) + _add_run(subparsers) + _add_cleanup(subparsers) + return parser.parse_args() diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py new file mode 100644 index 0000000..0f64d22 --- /dev/null +++ b/tests/slo/src/workload.py @@ -0,0 +1,320 @@ +""" +Core of the SQLAlchemy SLO workload. + +A single workload run keeps several reader and writer threads busy against one +YDB table through the ``ydb_sqlalchemy`` dialect, while a metrics thread ships +latency/availability samples to Prometheus over OTLP. Two execution modes are +supported, selected by ``WORKLOAD_NAME`` / ``--mode``: + + * ``core`` — SQLAlchemy Core (``Connection.execute``) + * ``orm`` — SQLAlchemy ORM (``Session`` + imperatively mapped entity) + +Each operation is wrapped in ``ydb_sqlalchemy.retry_operation``, which unwraps a +SQLAlchemy error back to the underlying ``ydb.Error`` and retries the transient +ones through the SDK's retry policy. We count the attempts (so ``retry_attempts`` +is a real signal) and record the final exception, if any, as an SLO failure. +""" + +import logging +import random +import threading +import time + +import sqlalchemy as sa +from sqlalchemy.orm import sessionmaker + +from ydb_sqlalchemy import retry_ydb_operation +from ydb_sqlalchemy import upsert as ydb_upsert + +from generator import RowGenerator, random_row +from metrics import OP_TYPE_READ, OP_TYPE_WRITE, create_metrics +from models import KeyValueRow, build_engine, build_table, ensure_mapped + +logger = logging.getLogger(__name__) + + +class SyncRateLimiter: + """Thread-safe limiter enforcing a minimum interval between permits.""" + + def __init__(self, min_interval_s: float) -> None: + self._min_interval_s = max(0.0, float(min_interval_s)) + self._lock = threading.Lock() + self._next_allowed_ts = 0.0 + + def __enter__(self): + if self._min_interval_s <= 0.0: + return self + while True: + with self._lock: + now = time.monotonic() + if now >= self._next_allowed_ts: + self._next_allowed_ts = now + self._min_interval_s + return self + sleep_for = self._next_allowed_ts - now + if sleep_for > 0: + time.sleep(sleep_for) + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + +def _limiter_for_rps(rps: int) -> SyncRateLimiter: + return SyncRateLimiter(0.0 if rps <= 0 else 1.0 / rps) + + +def _measure(metrics, op_type, fn, *, retry=False, idempotent=False): + """Time ``fn``, count its attempts and record the outcome. + + Autocommit single statements (``core``/``orm``) are already retried inside + ydb-dbapi, so they run once (``retry=False``). Interactive transactions + (``tx``) are not, so they run under ``ydb_sqlalchemy.retry_ydb_operation``, + which translates the SQLAlchemy error back to the underlying ``ydb.Error`` + and retries the transient ones — and we report how many attempts it took, so + ``retry_attempts`` is a real signal. The final exception, if any, is recorded + as an SLO failure but not re-raised, so the worker thread keeps going. + """ + start = metrics.start(op_type) + attempts = 0 + error = None + + def counted(): + nonlocal attempts + attempts += 1 + return fn() + + try: + if retry: + retry_ydb_operation(counted, idempotent=idempotent) + else: + counted() + except Exception as err: # noqa: BLE001 - record as a failed SLO op, keep going + error = err + metrics.stop(op_type, start, attempts=max(attempts, 1), error=error) + + +class Workload: + def __init__(self, args): + self.args = args + self.mode = getattr(args, "mode", "core") + self.table_name = args.table_name + self._session_factory = None + + # -- schema lifecycle --------------------------------------------------- + + def create(self): + engine = build_engine(self.args.endpoint, self.args.database, pool_size=4) + try: + metadata = sa.MetaData() + table = build_table( + metadata, + self.table_name, + min_partitions=self.args.min_partitions_count, + max_partitions=self.args.max_partitions_count, + partition_size_mb=self.args.partition_size, + ) + + # current + baseline containers create the table in parallel; the + # loser of the race sees "already exists" and we just move on. + try: + table.create(engine, checkfirst=True) + logger.info("Table %s is ready", self.table_name) + except Exception: + logger.warning("Create table %s reported an error (assuming it exists)", self.table_name, exc_info=True) + + self._fill_initial_data(engine, table) + finally: + engine.dispose() + + def _fill_initial_data(self, engine, table): + total = self.args.initial_data_count + batch_size = self.args.batch_size + generator = RowGenerator(start_id=0) + logger.info("Filling %s with %s initial rows", self.table_name, total) + + inserted = 0 + while inserted < total: + size = min(batch_size, total - inserted) + batch = [generator.get().as_params() for _ in range(size)] + self._upsert_batch(engine, table, batch) + inserted += size + + logger.info("Inserted %s rows into %s", inserted, self.table_name) + + @staticmethod + def _upsert_batch(engine, table, batch): + with engine.begin() as conn: + conn.execute(ydb_upsert(table), batch) + + def cleanup(self): + engine = build_engine(self.args.endpoint, self.args.database, pool_size=2) + try: + metadata = sa.MetaData() + table = build_table(metadata, self.table_name) + table.drop(engine, checkfirst=True) + logger.info("Dropped table %s", self.table_name) + finally: + engine.dispose() + + # -- load --------------------------------------------------------------- + + def run(self): + args = self.args + metrics = create_metrics(args.otlp_endpoint) + pool_size = args.read_threads + args.write_threads + 2 + engine = build_engine(args.endpoint, args.database, pool_size=pool_size) + + metadata = sa.MetaData() + table = build_table(metadata, self.table_name) + if self.mode == "orm": + ensure_mapped(table) + self._session_factory = sessionmaker(engine) + + max_id = self._max_id(engine, table) + logger.info("Starting '%s' SLO load on %s (max_id=%s)", self.mode, self.table_name, max_id) + + read_stmt = sa.select(table).where(table.c.object_id == sa.bindparam("object_id")) + read_limiter = _limiter_for_rps(args.read_rps) + write_limiter = _limiter_for_rps(args.write_rps) + row_generator = RowGenerator(start_id=max_id) + end_time = time.monotonic() + args.time + + threads = [] + for i in range(args.read_threads): + threads.append( + threading.Thread( + name=f"slo_read_{i}", + target=self._reader_loop, + args=(engine, read_stmt, metrics, read_limiter, max_id, end_time), + ) + ) + for i in range(args.write_threads): + threads.append( + threading.Thread( + name=f"slo_write_{i}", + target=self._writer_loop, + args=(engine, table, metrics, write_limiter, row_generator, max_id, end_time), + ) + ) + metrics_thread = threading.Thread( + name="slo_metrics", + target=self._metrics_loop, + args=(metrics, end_time, args.report_period), + ) + + for t in threads: + t.start() + metrics_thread.start() + + for t in threads: + t.join() + metrics_thread.join() + + metrics.push() + metrics.reset() + engine.dispose() + logger.info("Finished '%s' SLO load", self.mode) + + def _max_id(self, engine, table) -> int: + try: + with engine.connect() as conn: + value = conn.execute(sa.select(sa.func.max(table.c.object_id))).scalar() + return int(value) if value else self.args.initial_data_count + except Exception: + logger.warning("Could not read max(object_id); falling back to initial-data-count", exc_info=True) + return max(1, self.args.initial_data_count) + + # -- per-operation work ------------------------------------------------- + + def _reader_loop(self, engine, read_stmt, metrics, limiter, max_id, end_time): + while time.monotonic() < end_time: + with limiter: + object_id = random.randint(1, max(1, max_id)) + + if self.mode == "orm": + # Typical ORM read: fetch one mapped entity by primary key. + def do_read(oid=object_id): + with self._session_factory() as session: + session.get(KeyValueRow, oid) + + _measure(metrics, OP_TYPE_READ, do_read) + + elif self.mode == "tx": + # Read inside an interactive (SERIALIZABLE) transaction. + def do_read(oid=object_id): + with engine.connect().execution_options(isolation_level="SERIALIZABLE") as conn: + with conn.begin(): + conn.execute(read_stmt, {"object_id": oid}).fetchall() + + _measure(metrics, OP_TYPE_READ, do_read, retry=True, idempotent=True) + + else: # core + + def do_read(oid=object_id): + with engine.connect() as conn: + conn.execute(read_stmt, {"object_id": oid}).fetchall() + + _measure(metrics, OP_TYPE_READ, do_read) + + def _writer_loop(self, engine, table, metrics, limiter, row_generator, max_id, end_time): + while time.monotonic() < end_time: + with limiter: + if self.mode == "orm": + # Typical ORM write: add a new mapped object and commit the + # unit of work (a real INSERT through the session). + row = random_row() + + def do_write(r=row): + with self._session_factory() as session: + session.add(KeyValueRow(**r.as_params())) + session.commit() + + _measure(metrics, OP_TYPE_WRITE, do_write) + + elif self.mode == "tx": + # Read-modify-write in one interactive (SERIALIZABLE) transaction: + # exactly where the dialect's retry helper is needed (ydb-dbapi + # does not retry an interactive transaction on its own). + object_id = random.randint(1, max(1, max_id)) + + def do_write(oid=object_id): + with engine.connect().execution_options(isolation_level="SERIALIZABLE") as conn: + with conn.begin(): + current = conn.execute( + sa.select(table.c.payload_double).where(table.c.object_id == oid) + ).scalar() + conn.execute( + ydb_upsert(table).values(object_id=oid, payload_double=(current or 0.0) + 1.0) + ) + + # An increment is not idempotent, so only the pre-commit transient + # errors (the SDK's default, idempotent=False) are retried. + _measure(metrics, OP_TYPE_WRITE, do_write, retry=True, idempotent=False) + + else: # core + params = row_generator.get().as_params() + + def do_write(p=params): + with engine.begin() as conn: + conn.execute(ydb_upsert(table).values(**p)) + + _measure(metrics, OP_TYPE_WRITE, do_write) + + @staticmethod + def _metrics_loop(metrics, end_time, report_period_ms): + limiter = SyncRateLimiter(max(1, int(report_period_ms)) / 1000.0) + while time.monotonic() < end_time: + with limiter: + metrics.push() + + +def run_from_args(args): + workload = Workload(args) + command = args.command + if command == "create": + workload.create() + elif command == "run": + workload.run() + elif command == "cleanup": + workload.cleanup() + else: + raise ValueError(f"Unknown command: {command}") diff --git a/tox.ini b/tox.ini index 111242a..3b6b2bc 100644 --- a/tox.ini +++ b/tox.ini @@ -25,7 +25,7 @@ ignore_errors = True commands = docker-compose up -d python {toxinidir}/wait_container_ready.py - pytest -v test --dbdriver ydb --dbdriver ydb_async + pytest -v tests/integration --dbdriver ydb --dbdriver ydb_async pytest -v ydb_sqlalchemy docker-compose down -v @@ -33,7 +33,7 @@ commands = commands = docker-compose up -d python {toxinidir}/wait_container_ready.py - pytest -v test --dbdriver ydb --dbdriver ydb_async + pytest -v tests/integration --dbdriver ydb --dbdriver ydb_async docker-compose down -v [testenv:test-unit] @@ -51,22 +51,22 @@ commands = [testenv:black] skip_install = true commands = - black --diff --check ydb_sqlalchemy examples/basic_example test + black --diff --check ydb_sqlalchemy examples/basic_example tests/integration [testenv:black-format] skip_install = true commands = - black ydb_sqlalchemy examples/basic_example test + black ydb_sqlalchemy examples/basic_example tests/integration [testenv:isort] skip_install = true commands = - isort ydb_sqlalchemy examples/basic_example test + isort ydb_sqlalchemy examples/basic_example tests/integration [testenv:style] ignore_errors = True commands = - flake8 ydb_sqlalchemy examples/basic_example test + flake8 ydb_sqlalchemy examples/basic_example tests/integration [flake8] show-source = true diff --git a/ydb_sqlalchemy/__init__.py b/ydb_sqlalchemy/__init__.py index 55ade24..6e2bc08 100644 --- a/ydb_sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/__init__.py @@ -1,4 +1,5 @@ from ._version import VERSION # noqa: F401 from ydb_dbapi import IsolationLevel # noqa: F401 from .sqlalchemy import Upsert, types, upsert # noqa: F401 +from .sqlalchemy import retry_ydb, retry_ydb_operation, retry_ydb_operation_async # noqa: F401 import ydb_dbapi as dbapi diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index affed36..73198cd 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -21,6 +21,11 @@ import ydb_dbapi from ydb_sqlalchemy.sqlalchemy.dbapi_adapter import AdaptedAsyncConnection from ydb_sqlalchemy.sqlalchemy.dml import Upsert +from ydb_sqlalchemy.sqlalchemy.retry import ( # noqa: F401 + retry_ydb, + retry_ydb_operation, + retry_ydb_operation_async, +) from ydb_sqlalchemy.sqlalchemy.compiler import YqlCompiler, YqlDDLCompiler, YqlIdentifierPreparer, YqlTypeCompiler diff --git a/ydb_sqlalchemy/sqlalchemy/retry.py b/ydb_sqlalchemy/sqlalchemy/retry.py new file mode 100644 index 0000000..c812092 --- /dev/null +++ b/ydb_sqlalchemy/sqlalchemy/retry.py @@ -0,0 +1,120 @@ +"""Retry SQLAlchemy operations that fail with retryable YDB errors. + +SQLAlchemy wraps the underlying ydb-dbapi error in ``sqlalchemy.exc.DBAPIError``, +while YDB's retry logic only recognises ``ydb.Error``. These helpers unwrap the +original YDB error and re-raise it so the SDK can apply its retryable-error +classification (Unavailable, Overloaded, Aborted, BadSession, ...) and back-off. +Only YDB errors are retried; any other error propagates unchanged. + +The retried callable is run from scratch on every attempt, so it must be safe to +re-run (open its own connection/session, not depend on earlier state). + +Call form (sync and async):: + + rows = retry_ydb_operation(read, idempotent=True) + rows = await retry_ydb_operation_async(read, idempotent=True) + +Decorator form (works on both sync and async functions):: + + @retry_ydb(idempotent=True) + def read(): + with engine.connect() as conn: + return conn.execute(stmt).fetchall() +""" + +import functools +import inspect +from typing import Awaitable, Callable, Optional, TypeVar + +import sqlalchemy.exc +import ydb + +T = TypeVar("T") + +DEFAULT_MAX_RETRIES = 10 + + +def _retry_settings(max_retries: int, idempotent: bool) -> ydb.RetrySettings: + return ydb.RetrySettings(max_retries=max_retries, idempotent=idempotent) + + +def _unwrap_ydb_error(exc: BaseException) -> Optional[ydb.Error]: + """Return the ``ydb.Error`` wrapped inside a SQLAlchemy / ydb-dbapi error, if any.""" + original = getattr(getattr(exc, "orig", None), "original_error", None) + return original if isinstance(original, ydb.Error) else None + + +def retry_ydb_operation( + callee: Callable[[], T], + *, + max_retries: int = DEFAULT_MAX_RETRIES, + idempotent: bool = False, +) -> T: + """Run ``callee``, retrying transient YDB errors up to ``max_retries`` times. + + Set ``idempotent=True`` only when re-running ``callee`` is safe, so that + ambiguous errors are retried too. + """ + + def attempt() -> T: + try: + return callee() + except sqlalchemy.exc.DBAPIError as exc: + ydb_error = _unwrap_ydb_error(exc) + if ydb_error is not None: + raise ydb_error from exc + raise + + return ydb.retry_operation_sync(attempt, _retry_settings(max_retries, idempotent)) + + +async def retry_ydb_operation_async( + callee: Callable[[], Awaitable[T]], + *, + max_retries: int = DEFAULT_MAX_RETRIES, + idempotent: bool = False, +) -> T: + """Async counterpart of :func:`retry_ydb_operation`.""" + + async def attempt() -> T: + try: + return await callee() + except sqlalchemy.exc.DBAPIError as exc: + ydb_error = _unwrap_ydb_error(exc) + if ydb_error is not None: + raise ydb_error from exc + raise + + return await ydb.retry_operation_async(attempt, _retry_settings(max_retries, idempotent)) + + +def retry_ydb(*, max_retries: int = DEFAULT_MAX_RETRIES, idempotent: bool = False): + """Decorator that retries transient YDB errors raised by the wrapped function. + + Works on both sync and async functions:: + + @retry_ydb(idempotent=True) + def read(): ... + + @retry_ydb(idempotent=True) + async def read(): ... + """ + + def decorator(func): + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + return await retry_ydb_operation_async( + lambda: func(*args, **kwargs), max_retries=max_retries, idempotent=idempotent + ) + + return async_wrapper + + @functools.wraps(func) + def sync_wrapper(*args, **kwargs): + return retry_ydb_operation(lambda: func(*args, **kwargs), max_retries=max_retries, idempotent=idempotent) + + return sync_wrapper + + return decorator diff --git a/ydb_sqlalchemy/sqlalchemy/test_retry.py b/ydb_sqlalchemy/sqlalchemy/test_retry.py new file mode 100644 index 0000000..f77d5a1 --- /dev/null +++ b/ydb_sqlalchemy/sqlalchemy/test_retry.py @@ -0,0 +1,98 @@ +import asyncio + +import pytest +import sqlalchemy.exc +import ydb +import ydb_dbapi + +from .retry import _unwrap_ydb_error, retry_ydb, retry_ydb_operation, retry_ydb_operation_async + + +def _sa_dbapi_error(ydb_issue): + """Wrap a ydb issue the way the real stack does: ydb.Error -> ydb-dbapi -> SQLAlchemy.""" + dbapi_error = ydb_dbapi.OperationalError(str(ydb_issue), original_error=ydb_issue) + return sqlalchemy.exc.OperationalError("SELECT 1", {}, dbapi_error) + + +def test_unwrap_ydb_error_returns_original(): + issue = ydb.issues.Unavailable("node is down") + assert _unwrap_ydb_error(_sa_dbapi_error(issue)) is issue + + +def test_unwrap_ydb_error_none_for_plain_exception(): + assert _unwrap_ydb_error(ValueError("boom")) is None + + +def test_retry_ydb_operation_passes_through_success(): + assert retry_ydb_operation(lambda: 42) == 42 + + +def test_retry_ydb_operation_retries_transient_then_succeeds(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + def callee(): + calls.append(1) + if len(calls) < 3: + raise _sa_dbapi_error(issue) + return "ok" + + assert retry_ydb_operation(callee, max_retries=5) == "ok" + assert len(calls) == 3 # two transient failures, then success + + +def test_retry_ydb_operation_does_not_retry_non_retryable(): + issue = ydb.issues.BadRequest("bad query") + calls = [] + + def callee(): + calls.append(1) + raise _sa_dbapi_error(issue) + + with pytest.raises(ydb.issues.BadRequest): + retry_ydb_operation(callee, max_retries=5) + assert len(calls) == 1 # not retried + + +def test_retry_ydb_decorator_sync(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + @retry_ydb(max_retries=5) + def op(): + calls.append(1) + if len(calls) < 3: + raise _sa_dbapi_error(issue) + return "ok" + + assert op() == "ok" + assert len(calls) == 3 + + +def test_retry_ydb_decorator_async(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + @retry_ydb(max_retries=5) + async def op(): + calls.append(1) + if len(calls) < 3: + raise _sa_dbapi_error(issue) + return "ok" + + assert asyncio.run(op()) == "ok" + assert len(calls) == 3 + + +def test_retry_ydb_operation_async_call_form(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + async def callee(): + calls.append(1) + if len(calls) < 2: + raise _sa_dbapi_error(issue) + return "ok" + + assert asyncio.run(retry_ydb_operation_async(callee, max_retries=5)) == "ok" + assert len(calls) == 2