onestep-postgres 0.1.0__py3-none-any.whl

onestep_postgres/__init__.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version as _package_version
+
+from .connector import (
+    IncrementalDelivery,
+    PostgresConnector,
+    PostgresIncrementalSource,
+    PostgresTableQueueDelivery,
+    PostgresTableQueueSource,
+    PostgresTableSink,
+)
+from .resources import register_resources
+from .resilience import classify_sqlalchemy_error
+from .state_sqlalchemy import SQLAlchemyCursorStore, SQLAlchemyStateStore
+
+try:
+    __version__ = _package_version("onestep-postgres")
+except PackageNotFoundError:  # pragma: no cover - local source tree before install
+    __version__ = "dev"
+
+register = register_resources
+
+__all__ = [
+    "IncrementalDelivery",
+    "PostgresConnector",
+    "PostgresIncrementalSource",
+    "PostgresTableQueueDelivery",
+    "PostgresTableQueueSource",
+    "PostgresTableSink",
+    "SQLAlchemyCursorStore",
+    "SQLAlchemyStateStore",
+    "__version__",
+    "classify_sqlalchemy_error",
+    "register",
+    "register_resources",
+]

onestep_postgres/connector.py

@@ -0,0 +1,465 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from collections import deque
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from onestep.connectors.base import Delivery, Sink, Source
+from onestep.envelope import Envelope
+from onestep.resilience import ConnectorOperation
+from onestep.state import CursorStore, InMemoryCursorStore
+
+from .resilience import as_postgres_connector_operation_error
+from .state_sqlalchemy import SQLAlchemyCursorStore, SQLAlchemyStateStore
+
+try:
+    import sqlalchemy as sa
+    from sqlalchemy import create_engine
+except ImportError:  # pragma: no cover - exercised when optional deps are missing
+    sa = None
+    create_engine = None
+
+
+class PostgresConnector:
+    def __init__(self, dsn: str, **engine_options: Any) -> None:
+        if create_engine is None:
+            raise RuntimeError("PostgresConnector requires SQLAlchemy. Install onestep-postgres.")
+        self.dsn = dsn
+        self.engine = create_engine(dsn, future=True, pool_pre_ping=True, **engine_options)
+        self._tables: dict[str, Any] = {}
+
+    async def close(self) -> None:
+        await asyncio.to_thread(self.engine.dispose)
+
+    def state_store(
+        self,
+        *,
+        table: str = "onestep_state",
+        key_column: str = "state_key",
+        value_column: str = "state_value",
+        updated_at_column: str = "updated_at",
+        auto_create: bool = True,
+    ) -> SQLAlchemyStateStore:
+        return SQLAlchemyStateStore(
+            engine=self.engine,
+            table=table,
+            key_column=key_column,
+            value_column=value_column,
+            updated_at_column=updated_at_column,
+            auto_create=auto_create,
+        )
+
+    def cursor_store(
+        self,
+        *,
+        table: str = "onestep_cursor",
+        key_column: str = "cursor_key",
+        value_column: str = "cursor_value",
+        updated_at_column: str = "updated_at",
+        auto_create: bool = True,
+    ) -> SQLAlchemyCursorStore:
+        return SQLAlchemyCursorStore(
+            engine=self.engine,
+            table=table,
+            key_column=key_column,
+            value_column=value_column,
+            updated_at_column=updated_at_column,
+            auto_create=auto_create,
+        )
+
+    def table_queue(
+        self,
+        *,
+        table: str,
+        key: str,
+        where: str,
+        claim: Mapping[str, Any],
+        ack: Mapping[str, Any],
+        nack: Mapping[str, Any] | None = None,
+        batch_size: int = 100,
+        poll_interval_s: float = 1.0,
+    ) -> "PostgresTableQueueSource":
+        return PostgresTableQueueSource(
+            connector=self,
+            table=table,
+            key=key,
+            where=where,
+            claim=dict(claim),
+            ack=dict(ack),
+            nack=dict(nack or {}),
+            batch_size=batch_size,
+            poll_interval_s=poll_interval_s,
+        )
+
+    def incremental(
+        self,
+        *,
+        table: str,
+        key: str,
+        cursor: Sequence[str],
+        where: str | None = None,
+        batch_size: int = 1000,
+        poll_interval_s: float = 1.0,
+        state: CursorStore | None = None,
+        state_key: str | None = None,
+    ) -> "PostgresIncrementalSource":
+        if len(cursor) < 1:
+            raise ValueError("cursor must contain at least one column")
+        effective_cursor = tuple(cursor) if key in cursor else (*tuple(cursor), key)
+        return PostgresIncrementalSource(
+            connector=self,
+            table=table,
+            key=key,
+            cursor=effective_cursor,
+            where=where,
+            batch_size=batch_size,
+            poll_interval_s=poll_interval_s,
+            state=state or InMemoryCursorStore(),
+            state_key=state_key or _default_incremental_state_key(
+                table=table,
+                cursor=effective_cursor,
+                key=key,
+                where=where,
+            ),
+        )
+
+    def table_sink(
+        self,
+        *,
+        table: str,
+        mode: str = "insert",
+        keys: Sequence[str] = (),
+    ) -> "PostgresTableSink":
+        return PostgresTableSink(connector=self, table=table, mode=mode, keys=tuple(keys))
+
+    def _table(self, table_name: str):
+        table = self._tables.get(table_name)
+        if table is None:
+            metadata = sa.MetaData()
+            table = sa.Table(table_name, metadata, autoload_with=self.engine)
+            self._tables[table_name] = table
+        return table
+
+
+def _default_incremental_state_key(
+    *,
+    table: str,
+    cursor: Sequence[str],
+    key: str,
+    where: str | None,
+) -> str:
+    normalized_where = " ".join((where or "").split())
+    if normalized_where:
+        where_fragment = normalized_where
+        if len(where_fragment) > 64:
+            where_fragment = f"sha1:{hashlib.sha1(where_fragment.encode('utf-8')).hexdigest()}"
+    else:
+        where_fragment = "-"
+    return f"{table}:{','.join(cursor)}:key={key}:where={where_fragment}"
+
+
+@dataclass
+class _TableRowRef:
+    table: str
+    key: str
+    key_value: Any
+
+
+class PostgresTableQueueDelivery(Delivery):
+    def __init__(self, source: "PostgresTableQueueSource", envelope: Envelope, row_ref: _TableRowRef) -> None:
+        super().__init__(envelope)
+        self._source = source
+        self._row_ref = row_ref
+
+    async def update_current_row(self, values: Mapping[str, Any]) -> None:
+        payload = dict(values)
+        await self._source.update_row(self._row_ref, payload)
+        if isinstance(self.envelope.body, dict):
+            self.envelope.body.update(payload)
+
+    async def ack(self) -> None:
+        await self._source.ack_row(self._row_ref)
+
+    async def retry(self, *, delay_s: float | None = None) -> None:
+        await self._source.retry_row(self._row_ref, delay_s=delay_s)
+
+    async def fail(self, exc: Exception | None = None) -> None:
+        await self._source.fail_row(self._row_ref, exc=exc)
+
+    async def release_unstarted(self) -> None:
+        await self._source.release_row(self._row_ref)
+
+
+class PostgresTableQueueSource(Source):
+    fetch_is_cancel_safe = False
+
+    def __init__(
+        self,
+        *,
+        connector: PostgresConnector,
+        table: str,
+        key: str,
+        where: str,
+        claim: dict[str, Any],
+        ack: dict[str, Any],
+        nack: dict[str, Any],
+        batch_size: int,
+        poll_interval_s: float,
+    ) -> None:
+        super().__init__(f"postgres.table_queue:{table}")
+        self.connector = connector
+        self.table_name = table
+        self.key = key
+        self.where = where
+        self.claim = claim
+        self.ack = ack
+        self.nack = nack
+        self.batch_size = batch_size
+        self.poll_interval_s = poll_interval_s
+
+    async def fetch(self, limit: int) -> list[Delivery]:
+        try:
+            rows = await asyncio.to_thread(self._fetch_sync, max(1, min(limit, self.batch_size)))
+        except Exception as exc:
+            connector_error = as_postgres_connector_operation_error(
+                operation=ConnectorOperation.FETCH,
+                exc=exc,
+                source_name=self.name,
+                retry_delay_s=self.poll_interval_s,
+            )
+            if connector_error is None:
+                raise
+            raise connector_error from exc
+        deliveries: list[Delivery] = []
+        for row in rows:
+            key_value = row[self.key]
+            envelope = Envelope(body=row, meta={"table": self.table_name})
+            row_ref = _TableRowRef(self.table_name, self.key, key_value)
+            deliveries.append(PostgresTableQueueDelivery(self, envelope, row_ref))
+        return deliveries
+
+    def _fetch_sync(self, limit: int) -> list[dict[str, Any]]:
+        table = self.connector._table(self.table_name)
+        with self.connector.engine.begin() as conn:
+            stmt = sa.select(table).where(sa.text(self.where)).order_by(table.c[self.key]).limit(limit)
+            try:
+                stmt = stmt.with_for_update(skip_locked=True)
+            except TypeError:
+                stmt = stmt.with_for_update()
+            rows = [dict(row) for row in conn.execute(stmt).mappings().all()]
+            if not rows:
+                return []
+            ids = [row[self.key] for row in rows]
+            conn.execute(sa.update(table).where(table.c[self.key].in_(ids)).values(**self.claim))
+            refreshed = conn.execute(
+                sa.select(table).where(table.c[self.key].in_(ids)).order_by(table.c[self.key])
+            )
+            return [dict(row) for row in refreshed.mappings().all()]
+
+    async def ack_row(self, row_ref: _TableRowRef) -> None:
+        await self.update_row(row_ref, self.ack)
+
+    async def retry_row(self, row_ref: _TableRowRef, *, delay_s: float | None = None) -> None:
+        if delay_s:
+            await asyncio.sleep(delay_s)
+        await self.update_row(row_ref, self.nack)
+
+    async def fail_row(self, row_ref: _TableRowRef, exc: Exception | None = None) -> None:
+        await self.update_row(row_ref, self.nack)
+
+    async def release_row(self, row_ref: _TableRowRef) -> None:
+        await self.update_row(row_ref, self.nack)
+
+    async def update_row(self, row_ref: _TableRowRef, values: Mapping[str, Any]) -> None:
+        await asyncio.to_thread(self._update_row_sync, row_ref, dict(values))
+
+    def _update_row_sync(self, row_ref: _TableRowRef, values: Mapping[str, Any]) -> None:
+        if not values:
+            return
+        table = self.connector._table(row_ref.table)
+        with self.connector.engine.begin() as conn:
+            conn.execute(
+                sa.update(table)
+                .where(table.c[row_ref.key] == row_ref.key_value)
+                .values(**dict(values))
+            )
+
+
+@dataclass
+class _CursorToken:
+    value: tuple[Any, ...]
+
+
+class IncrementalDelivery(Delivery):
+    def __init__(self, source: "PostgresIncrementalSource", envelope: Envelope, token: _CursorToken) -> None:
+        super().__init__(envelope)
+        self._source = source
+        self._token = token
+
+    async def ack(self) -> None:
+        await self._source.ack_token(self._token)
+
+    async def retry(self, *, delay_s: float | None = None) -> None:
+        if delay_s:
+            await asyncio.sleep(delay_s)
+
+    async def fail(self, exc: Exception | None = None) -> None:
+        return None
+
+
+class PostgresIncrementalSource(Source):
+    def __init__(
+        self,
+        *,
+        connector: PostgresConnector,
+        table: str,
+        key: str,
+        cursor: tuple[str, ...],
+        where: str | None,
+        batch_size: int,
+        poll_interval_s: float,
+        state: CursorStore,
+        state_key: str,
+    ) -> None:
+        super().__init__(f"postgres.incremental:{table}")
+        self.connector = connector
+        self.table_name = table
+        self.key = key
+        self.configured_cursor = cursor
+        self.cursor = cursor if key in cursor else (*cursor, key)
+        self.where = where
+        self.batch_size = batch_size
+        self.poll_interval_s = poll_interval_s
+        self.state = state
+        self.state_key = state_key
+        self._pending: deque[tuple[Any, ...]] = deque()
+        self._acked: set[tuple[Any, ...]] = set()
+        self._commit_lock: asyncio.Lock | None = None
+        self._commit_loop: asyncio.AbstractEventLoop | None = None
+        self._loaded = False
+        self._committed_cursor: tuple[Any, ...] | None = None
+        self._fetched_cursor: tuple[Any, ...] | None = None
+
+    async def open(self) -> None:
+        if not self._loaded:
+            loaded = await self.state.load(self.state_key)
+            if loaded is not None and len(loaded) == len(self.cursor):
+                self._committed_cursor = tuple(loaded)
+                self._fetched_cursor = self._committed_cursor
+            self._loaded = True
+
+    async def fetch(self, limit: int) -> list[Delivery]:
+        await self.open()
+        try:
+            rows = await asyncio.to_thread(self._fetch_sync, max(1, min(limit, self.batch_size)))
+        except Exception as exc:
+            connector_error = as_postgres_connector_operation_error(
+                operation=ConnectorOperation.FETCH,
+                exc=exc,
+                source_name=self.name,
+                retry_delay_s=self.poll_interval_s,
+            )
+            if connector_error is None:
+                raise
+            raise connector_error from exc
+        deliveries: list[Delivery] = []
+        for row in rows:
+            token = _CursorToken(tuple(row[column] for column in self.cursor))
+            self._pending.append(token.value)
+            self._fetched_cursor = token.value
+            envelope = Envelope(body=row, meta={"table": self.table_name})
+            deliveries.append(IncrementalDelivery(self, envelope, token))
+        return deliveries
+
+    def _fetch_sync(self, limit: int) -> list[dict[str, Any]]:
+        table = self.connector._table(self.table_name)
+        stmt = sa.select(table)
+        predicates = []
+        if self.where:
+            predicates.append(sa.text(self.where))
+        read_cursor = self._fetched_cursor or self._committed_cursor
+        if read_cursor is not None:
+            cursor_columns = [table.c[name] for name in self.cursor]
+            predicates.append(sa.tuple_(*cursor_columns) > tuple(read_cursor))
+        if predicates:
+            stmt = stmt.where(*predicates)
+        order_columns = [table.c[name] for name in self.cursor]
+        stmt = stmt.order_by(*order_columns).limit(limit)
+        with self.connector.engine.begin() as conn:
+            rows = conn.execute(stmt).mappings().all()
+        return [dict(row) for row in rows]
+
+    async def ack_token(self, token: _CursorToken) -> None:
+        lock = self._runtime_commit_lock()
+        async with lock:
+            self._acked.add(token.value)
+            advanced: tuple[Any, ...] | None = None
+            while self._pending and self._pending[0] in self._acked:
+                advanced = self._pending.popleft()
+                self._acked.remove(advanced)
+            if advanced is not None:
+                self._committed_cursor = advanced
+                if not self._pending:
+                    self._fetched_cursor = advanced
+                await self.state.save(self.state_key, list(advanced))
+
+    def _runtime_commit_lock(self) -> asyncio.Lock:
+        current_loop = asyncio.get_running_loop()
+        if self._commit_lock is None or self._commit_loop is not current_loop:
+            self._commit_lock = asyncio.Lock()
+            self._commit_loop = current_loop
+        return self._commit_lock
+
+
+class PostgresTableSink(Sink):
+    def __init__(self, *, connector: PostgresConnector, table: str, mode: str, keys: tuple[str, ...]) -> None:
+        super().__init__(f"postgres.table_sink:{table}")
+        if mode not in {"insert", "upsert"}:
+            raise ValueError("mode must be either 'insert' or 'upsert'")
+        self.connector = connector
+        self.table_name = table
+        self.mode = mode
+        self.keys = keys
+
+    async def send(self, envelope: Envelope) -> None:
+        if not isinstance(envelope.body, Mapping):
+            raise TypeError("PostgresTableSink only accepts mapping payloads")
+        try:
+            await asyncio.to_thread(self._send_sync, dict(envelope.body))
+        except Exception as exc:
+            connector_error = as_postgres_connector_operation_error(
+                operation=ConnectorOperation.SEND,
+                exc=exc,
+                source_name=self.name,
+                retry_delay_s=1.0,
+            )
+            if connector_error is None:
+                raise
+            raise connector_error from exc
+
+    def _send_sync(self, payload: dict[str, Any]) -> None:
+        table = self.connector._table(self.table_name)
+        dialect = self.connector.engine.dialect.name
+        with self.connector.engine.begin() as conn:
+            if self.mode == "insert":
+                conn.execute(sa.insert(table).values(**payload))
+                return
+            if not self.keys:
+                raise ValueError("upsert mode requires keys")
+            update_payload = {key: value for key, value in payload.items() if key not in self.keys}
+            if dialect == "postgresql":
+                from sqlalchemy.dialects.postgresql import insert as postgres_insert
+
+                stmt = postgres_insert(table).values(**payload)
+                conn.execute(stmt.on_conflict_do_update(index_elements=list(self.keys), set_=update_payload))
+                return
+            if dialect == "sqlite":
+                from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+                stmt = sqlite_insert(table).values(**payload)
+                conn.execute(stmt.on_conflict_do_update(index_elements=list(self.keys), set_=update_payload))
+                return
+            conn.execute(sa.insert(table).values(**payload))

onestep_postgres/resilience.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from onestep.resilience import ConnectorErrorKind, ConnectorOperation, ConnectorOperationError
+
+try:  # pragma: no cover - optional dependency
+    import sqlalchemy as sa
+except ImportError:  # pragma: no cover - optional dependency
+    sa = None
+
+
+def classify_sqlalchemy_error(exc: BaseException) -> ConnectorErrorKind | None:
+    if sa is None:
+        return None
+    sql_exc = sa.exc
+    if isinstance(exc, getattr(sql_exc, "TimeoutError", ())):
+        return ConnectorErrorKind.TRANSIENT
+    if isinstance(exc, getattr(sql_exc, "InterfaceError", ())):
+        return ConnectorErrorKind.DISCONNECTED
+    if isinstance(exc, getattr(sql_exc, "ProgrammingError", ())):
+        return ConnectorErrorKind.PERMANENT
+    if isinstance(exc, getattr(sql_exc, "DBAPIError", ())):
+        if getattr(exc, "connection_invalidated", False):
+            return ConnectorErrorKind.DISCONNECTED
+        message = " ".join(
+            str(part).lower()
+            for part in (
+                getattr(exc, "orig", None),
+                exc,
+            )
+            if part is not None
+        )
+        if any(token in message for token in ("server closed the connection", "connection refused", "connection reset")):
+            return ConnectorErrorKind.DISCONNECTED
+        if any(token in message for token in ("deadlock detected", "lock timeout", "could not serialize access")):
+            return ConnectorErrorKind.TRANSIENT
+        if any(token in message for token in ("password authentication failed", "permission denied")):
+            return ConnectorErrorKind.MISCONFIGURED
+        if any(token in message for token in ("database", "role")) and "does not exist" in message:
+            return ConnectorErrorKind.MISCONFIGURED
+        if any(token in message for token in ("undefined table", "undefined column", "syntax error", "does not exist")):
+            return ConnectorErrorKind.PERMANENT
+        if isinstance(exc, getattr(sql_exc, "OperationalError", ())):
+            return ConnectorErrorKind.TRANSIENT
+    return None
+
+
+def as_postgres_connector_operation_error(
+    *,
+    operation: ConnectorOperation,
+    exc: BaseException,
+    source_name: str | None = None,
+    retry_delay_s: float | None = None,
+) -> ConnectorOperationError | None:
+    kind = classify_sqlalchemy_error(exc)
+    if kind is None:
+        return None
+    return ConnectorOperationError(
+        backend="postgres",
+        operation=operation,
+        kind=kind,
+        source_name=source_name,
+        retry_delay_s=retry_delay_s,
+        cause=exc,
+    )

onestep_postgres/resources.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from onestep.resource_registry import ResourceBuildContext, ResourceRegistry, ResourceSpecHandler
+
+from .connector import PostgresConnector
+
+_POSTGRES_FIELDS = frozenset({"type", "dsn", "engine_options"})
+_POSTGRES_STATE_STORE_FIELDS = frozenset(
+    {"type", "connector", "table", "key_column", "value_column", "updated_at_column", "auto_create"}
+)
+_POSTGRES_CURSOR_STORE_FIELDS = frozenset(
+    {"type", "connector", "table", "key_column", "value_column", "updated_at_column", "auto_create"}
+)
+_POSTGRES_TABLE_QUEUE_FIELDS = frozenset(
+    {"type", "connector", "table", "key", "where", "claim", "ack", "nack", "batch_size", "poll_interval_s"}
+)
+_POSTGRES_INCREMENTAL_FIELDS = frozenset(
+    {"type", "connector", "table", "key", "cursor", "where", "batch_size", "poll_interval_s", "state", "state_key"}
+)
+_POSTGRES_TABLE_SINK_FIELDS = frozenset({"type", "connector", "table", "mode", "keys"})
+
+
+def register_resources(registry: ResourceRegistry) -> None:
+    registry.register_resource_type(
+        ResourceSpecHandler(
+            type="postgres",
+            allowed_fields=_POSTGRES_FIELDS,
+            build=_build_postgres,
+        )
+    )
+    registry.register_resource_type(
+        ResourceSpecHandler(
+            type="postgres_state_store",
+            allowed_fields=_POSTGRES_STATE_STORE_FIELDS,
+            build=_build_postgres_state_store,
+        )
+    )
+    registry.register_resource_type(
+        ResourceSpecHandler(
+            type="postgres_cursor_store",
+            allowed_fields=_POSTGRES_CURSOR_STORE_FIELDS,
+            build=_build_postgres_cursor_store,
+        )
+    )
+    registry.register_resource_type(
+        ResourceSpecHandler(
+            type="postgres_table_queue",
+            allowed_fields=_POSTGRES_TABLE_QUEUE_FIELDS,
+            build=_build_postgres_table_queue,
+        )
+    )
+    registry.register_resource_type(
+        ResourceSpecHandler(
+            type="postgres_incremental",
+            allowed_fields=_POSTGRES_INCREMENTAL_FIELDS,
+            build=_build_postgres_incremental,
+        )
+    )
+    registry.register_resource_type(
+        ResourceSpecHandler(
+            type="postgres_table_sink",
+            allowed_fields=_POSTGRES_TABLE_SINK_FIELDS,
+            build=_build_postgres_table_sink,
+        )
+    )
+
+
+def _build_postgres(ctx: ResourceBuildContext, spec: Mapping[str, Any]) -> PostgresConnector:
+    return PostgresConnector(
+        ctx.require_string(spec, "dsn"),
+        **ctx.mapping_value(spec.get("engine_options"), field=f"{ctx.field}.engine_options"),
+    )
+
+
+def _build_postgres_state_store(ctx: ResourceBuildContext, spec: Mapping[str, Any]) -> Any:
+    connector = ctx.resolve_dependency(spec, "connector")
+    if not hasattr(connector, "state_store"):
+        raise TypeError(f"resource {spec['connector']!r} cannot build postgres_state_store")
+    return connector.state_store(
+        table=spec.get("table", "onestep_state"),
+        key_column=spec.get("key_column", "state_key"),
+        value_column=spec.get("value_column", "state_value"),
+        updated_at_column=spec.get("updated_at_column", "updated_at"),
+        auto_create=spec.get("auto_create", True),
+    )
+
+
+def _build_postgres_cursor_store(ctx: ResourceBuildContext, spec: Mapping[str, Any]) -> Any:
+    connector = ctx.resolve_dependency(spec, "connector")
+    if not hasattr(connector, "cursor_store"):
+        raise TypeError(f"resource {spec['connector']!r} cannot build postgres_cursor_store")
+    return connector.cursor_store(
+        table=spec.get("table", "onestep_cursor"),
+        key_column=spec.get("key_column", "cursor_key"),
+        value_column=spec.get("value_column", "cursor_value"),
+        updated_at_column=spec.get("updated_at_column", "updated_at"),
+        auto_create=spec.get("auto_create", True),
+    )
+
+
+def _build_postgres_table_queue(ctx: ResourceBuildContext, spec: Mapping[str, Any]) -> Any:
+    connector = ctx.resolve_dependency(spec, "connector")
+    if not hasattr(connector, "table_queue"):
+        raise TypeError(f"resource {spec['connector']!r} cannot build postgres_table_queue")
+    return connector.table_queue(
+        table=ctx.require_string(spec, "table"),
+        key=ctx.require_string(spec, "key"),
+        where=ctx.require_string(spec, "where"),
+        claim=ctx.require_mapping(spec, "claim"),
+        ack=ctx.require_mapping(spec, "ack"),
+        nack=ctx.optional_mapping(spec.get("nack"), field=f"{ctx.field}.nack") or None,
+        batch_size=spec.get("batch_size", 100),
+        poll_interval_s=spec.get("poll_interval_s", 1.0),
+    )
+
+
+def _build_postgres_incremental(ctx: ResourceBuildContext, spec: Mapping[str, Any]) -> Any:
+    connector = ctx.resolve_dependency(spec, "connector")
+    if not hasattr(connector, "incremental"):
+        raise TypeError(f"resource {spec['connector']!r} cannot build postgres_incremental")
+    raw_state_name = spec.get("state")
+    state = None
+    if raw_state_name is not None:
+        state_name = ctx.string_value(raw_state_name, field=f"{ctx.field}.state")
+        state = ctx.resolve(state_name)
+        if not ctx.is_cursor_store(state):
+            raise TypeError(f"resource {state_name!r} cannot be used as incremental state")
+    return connector.incremental(
+        table=ctx.require_string(spec, "table"),
+        key=ctx.require_string(spec, "key"),
+        cursor=tuple(ctx.string_list(spec.get("cursor"), field=f"{ctx.field}.cursor")),
+        where=spec.get("where"),
+        batch_size=spec.get("batch_size", 1000),
+        poll_interval_s=spec.get("poll_interval_s", 1.0),
+        state=state,
+        state_key=spec.get("state_key"),
+    )
+
+
+def _build_postgres_table_sink(ctx: ResourceBuildContext, spec: Mapping[str, Any]) -> Any:
+    connector = ctx.resolve_dependency(spec, "connector")
+    if not hasattr(connector, "table_sink"):
+        raise TypeError(f"resource {spec['connector']!r} cannot build postgres_table_sink")
+    keys = spec.get("keys")
+    return connector.table_sink(
+        table=ctx.require_string(spec, "table"),
+        mode=spec.get("mode", "insert"),
+        keys=tuple(ctx.string_list(keys, field=f"{ctx.field}.keys")) if keys is not None else (),
+    )

onestep_postgres/state_sqlalchemy.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import threading
+from datetime import datetime, timezone
+from typing import Any
+
+try:
+    import sqlalchemy as sa
+    from sqlalchemy import create_engine
+except ImportError:  # pragma: no cover - exercised when optional deps are missing
+    sa = None
+    create_engine = None
+
+
+class SQLAlchemyStateStore:
+    def __init__(
+        self,
+        *,
+        dsn: str | None = None,
+        engine: Any | None = None,
+        table: str = "onestep_state",
+        key_column: str = "state_key",
+        value_column: str = "state_value",
+        updated_at_column: str = "updated_at",
+        auto_create: bool = True,
+        **engine_options: Any,
+    ) -> None:
+        if create_engine is None or sa is None:
+            raise RuntimeError("SQLAlchemyStateStore requires SQLAlchemy. Install onestep-postgres.")
+        if engine is None and dsn is None:
+            raise ValueError("dsn or engine is required")
+        if engine is not None and dsn is not None:
+            raise ValueError("pass either dsn or engine, not both")
+        self.engine = engine or create_engine(dsn, future=True, pool_pre_ping=True, **engine_options)
+        self._owns_engine = engine is None
+        self.table_name = table
+        self.key_column_name = key_column
+        self.value_column_name = value_column
+        self.updated_at_column_name = updated_at_column
+        self.auto_create = auto_create
+        self._metadata = sa.MetaData()
+        self._table = sa.Table(
+            table,
+            self._metadata,
+            sa.Column(key_column, sa.String(255), primary_key=True),
+            sa.Column(value_column, sa.Text(), nullable=False),
+            sa.Column(updated_at_column, sa.DateTime(timezone=True), nullable=False),
+        )
+        self._ready = False
+        self._ready_lock = threading.Lock()
+
+    async def load(self, key: str) -> Any | None:
+        return await asyncio.to_thread(self._load_sync, key)
+
+    async def save(self, key: str, value: Any) -> None:
+        await asyncio.to_thread(self._save_sync, key, value)
+
+    async def delete(self, key: str) -> None:
+        await asyncio.to_thread(self._delete_sync, key)
+
+    async def close(self) -> None:
+        if self._owns_engine:
+            await asyncio.to_thread(self.engine.dispose)
+
+    def _load_sync(self, key: str) -> Any | None:
+        self._ensure_ready_sync()
+        key_column = self._table.c[self.key_column_name]
+        value_column = self._table.c[self.value_column_name]
+        with self.engine.begin() as conn:
+            row = conn.execute(sa.select(value_column).where(key_column == key)).scalar_one_or_none()
+        if row is None:
+            return None
+        return json.loads(row)
+
+    def _save_sync(self, key: str, value: Any) -> None:
+        self._ensure_ready_sync()
+        key_column = self._table.c[self.key_column_name]
+        payload = {
+            self.key_column_name: key,
+            self.value_column_name: json.dumps(value, ensure_ascii=False),
+            self.updated_at_column_name: datetime.now(timezone.utc),
+        }
+        with self.engine.begin() as conn:
+            exists = conn.execute(sa.select(key_column).where(key_column == key)).scalar_one_or_none()
+            if exists is None:
+                conn.execute(sa.insert(self._table).values(**payload))
+                return
+            conn.execute(
+                sa.update(self._table)
+                .where(key_column == key)
+                .values(
+                    **{
+                        self.value_column_name: payload[self.value_column_name],
+                        self.updated_at_column_name: payload[self.updated_at_column_name],
+                    }
+                )
+            )
+
+    def _delete_sync(self, key: str) -> None:
+        self._ensure_ready_sync()
+        key_column = self._table.c[self.key_column_name]
+        with self.engine.begin() as conn:
+            conn.execute(sa.delete(self._table).where(key_column == key))
+
+    def _ensure_ready_sync(self) -> None:
+        if self._ready or not self.auto_create:
+            return
+        with self._ready_lock:
+            if self._ready:
+                return
+            self._metadata.create_all(self.engine, tables=[self._table], checkfirst=True)
+            self._ready = True
+
+
+class SQLAlchemyCursorStore(SQLAlchemyStateStore):
+    pass

onestep_postgres-0.1.0.dist-info/METADATA

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: onestep-postgres
+Version: 0.1.0
+Summary: PostgreSQL connector plugin for onestep.
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: onestep>=1.4.2
+Requires-Dist: psycopg[binary]>=3.2.0
+Requires-Dist: sqlalchemy>=2.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == 'dev'
+Provides-Extra: test
+Requires-Dist: pytest>=8.0.0; extra == 'test'
+Description-Content-Type: text/markdown
+
+# onestep-postgres
+
+PostgreSQL connector plugin for onestep.
+
+Install it with:
+
+```bash
+pip install onestep-postgres
+```
+
+YAML resources are available after the plugin is installed:
+
+```yaml
+resources:
+  pg:
+    type: postgres
+    dsn: "${POSTGRES_DSN}"
+
+  cursor:
+    type: postgres_cursor_store
+    connector: pg
+
+  users:
+    type: postgres_incremental
+    connector: pg
+    table: users
+    key: id
+    cursor: [updated_at, id]
+    state: cursor
+
+  processed:
+    type: postgres_table_sink
+    connector: pg
+    table: processed_users
+    mode: upsert
+    keys: [id]
+```
+
+The first version supports table queues, incremental polling, table sinks, and
+SQLAlchemy-backed state/cursor stores. It does not support PostgreSQL logical
+replication or CDC.

onestep_postgres-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+onestep_postgres/__init__.py,sha256=BQfpRdzWJGyJx0Tc2jnI3BA2Mm0u61ZVFuUVue1Wrg0,1017
+onestep_postgres/connector.py,sha256=aj1A9XXRITDgkB9mYfJs-54THqf50gb1uf0U7HUyD-w,16987
+onestep_postgres/resilience.py,sha256=tr0Jhgv13xJrvVjuoaSI2_tsYhrCrOOUz2K8IMYwBSY,2586
+onestep_postgres/resources.py,sha256=qu3pWzeAv42UeMOW6MYKM7rmdcX-UZe0gwx2OaYW7is,6266
+onestep_postgres/state_sqlalchemy.py,sha256=g2yFDgygvnHFMutLzEnMaZNmDs5Pb-m9N1EReO4wERk,4400
+onestep_postgres-0.1.0.dist-info/METADATA,sha256=0i05tUfWh2_f7RPQiW-TOCbDOr87kZrgbV-t_6K2e9A,1173
+onestep_postgres-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+onestep_postgres-0.1.0.dist-info/entry_points.txt,sha256=XuAxF8oKTQStyyD4DJ2phsFC1J5N2EKD-3uXvP-hh9I,57
+onestep_postgres-0.1.0.dist-info/RECORD,,

onestep_postgres-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

onestep_postgres-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[onestep.resources]
+postgres = onestep_postgres:register