servomexlib.sinks¶

async def close(self) -> None:
    """Flush and close the CSV file. Idempotent."""
    if self._file is None:
        return
    try:
        self._file.flush()
    finally:
        self._file.close()
        self._file = None
        self._writer = None

async def open(self) -> None:
    """Open the CSV file for writing (truncating). Idempotent."""
    if self._file is not None:
        return
    self._file = self._path.open("w", encoding="utf-8", newline="")

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Append ``samples`` as CSV rows, locking the header on first call."""
    if self._file is None:
        raise RuntimeError("CsvSink: write_many called before open()")
    if not samples:
        return

    rows = [sample_to_row(s) for s in samples]
    if self._writer is None:
        self._columns = tuple(rows[0].keys())
        self._writer = csv.DictWriter(self._file, fieldnames=list(self._columns))
        self._writer.writeheader()

    columns = self._columns
    assert columns is not None  # noqa: S101 — narrow for type checker

    for row in rows:
        for key in row.keys() - set(columns):
            if key not in self._unknown_columns_warned:
                self._unknown_columns_warned.add(key)
                _logger.warning(
                    "sinks.csv.unknown_column path=%s column=%s action=drop",
                    str(self._path),
                    key,
                )
        self._writer.writerow({k: row.get(k) for k in columns})
    self._file.flush()

async def close(self) -> None:
    """Flip the closed flag — no I/O, buffer preserved for inspection."""
    self._closed = True

async def open(self) -> None:
    """No backing resource — just flips the open flag."""
    self._open = True
    self._closed = False

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Append every sample to the internal buffer."""
    if not self.is_open:
        raise RuntimeError("InMemorySink: write_many called before open()")
    self._samples.extend(samples)

async def close(self) -> None:
    """Flush and close the JSONL file. Idempotent."""
    if self._file is None:
        return
    try:
        self._file.flush()
    finally:
        self._file.close()
        self._file = None

async def open(self) -> None:
    """Open the JSONL file for writing in append mode. Idempotent."""
    if self._file is not None:
        return
    self._file = self._path.open("a", encoding="utf-8", newline="")

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Serialise each sample as one JSON object per line."""
    if self._file is None:
        raise RuntimeError("JsonlSink: write_many called before open()")
    if not samples:
        return
    for sample in samples:
        self._file.write(json.dumps(sample_to_row(sample), ensure_ascii=False))
        self._file.write("\n")
    self._file.flush()

async def close(self) -> None:
    """Flush the footer and close the writer. Idempotent."""
    if self._writer is not None:
        try:
            self._writer.close()
        finally:
            self._writer = None
    self._pa = None
    self._pq = None
    _logger.info(
        "sinks.parquet.close path=%s rows_written=%s", str(self._path), self._rows_written
    )

async def open(self) -> None:
    """Load pyarrow and create the parent directory. Idempotent.

    The ``ParquetWriter`` itself opens lazily on the first
    :meth:`write_many`, when the concrete schema is known.
    """
    if self._pa is not None:
        return
    self._pa, self._pq = _load_pyarrow()
    self._path.parent.mkdir(parents=True, exist_ok=True)
    _logger.info(
        "sinks.parquet.open path=%s compression=%s", str(self._path), self._compression
    )

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Append ``samples`` as one Parquet row group."""
    if self._pa is None:
        raise RuntimeError("ParquetSink: write_many called before open()")
    if not samples:
        return

    rows = [sample_to_row(s) for s in samples]

    if not self._schema.is_locked:
        self._schema.lock(rows)
        self._arrow_schema = self._build_arrow_schema()
        self._writer = self._open_writer()

    assert self._writer is not None  # noqa: S101 — narrow for type checker
    assert self._arrow_schema is not None  # noqa: S101

    projected = [self._schema.project(r) for r in rows]
    columns = self._schema.columns
    assert columns is not None  # noqa: S101

    arrays = {spec.name: [row[spec.name] for row in projected] for spec in columns}
    try:
        table = self._pa.Table.from_pydict(arrays, schema=self._arrow_schema)
        self._writer.write_table(table, row_group_size=self._row_group_size)
    except Exception as exc:
        raise ServomexSinkWriteError(
            f"ParquetSink: write failed for {self._path}: {exc}"
        ) from exc
    self._rows_written += len(projected)

def target(self) -> str:
    """Return a log-safe ``host:port/db.schema.table`` (no password)."""
    if self.dsn is not None:
        parsed = urlparse(self.dsn)
        host = parsed.hostname or "?"
        port = parsed.port or self.port
        db = (parsed.path or "/?").lstrip("/") or "?"
    else:
        host = self.host or "?"
        port = self.port
        db = self.database or "?"
    return f"{host}:{port}/{db}.{self.schema}.{self.table}"

async def close(self) -> None:
    """Close the pool. Idempotent."""
    if self._pool is None:
        return
    pool = self._pool
    self._pool = None
    try:
        await pool.close()
    finally:
        _logger.info(
            "sinks.postgres.close target=%s rows_written=%s",
            self._config.target(),
            self._rows_written,
        )
    self._asyncpg = None

async def open(self) -> None:
    """Load asyncpg, open the pool, and (optionally) introspect the table."""
    if self._pool is not None:
        return
    self._asyncpg = _load_asyncpg()
    cfg = self._config
    server_settings = {
        "application_name": "servomexlib",
        "statement_timeout": str(int(cfg.statement_timeout_ms)),
    }
    try:
        if cfg.dsn is not None:
            self._pool = await self._asyncpg.create_pool(
                dsn=cfg.dsn,
                min_size=cfg.pool_min_size,
                max_size=cfg.pool_max_size,
                command_timeout=cfg.command_timeout_s,
                server_settings=server_settings,
            )
        else:
            self._pool = await self._asyncpg.create_pool(
                host=cfg.host,
                port=cfg.port,
                user=cfg.user,
                password=cfg.password,
                database=cfg.database,
                min_size=cfg.pool_min_size,
                max_size=cfg.pool_max_size,
                command_timeout=cfg.command_timeout_s,
                server_settings=server_settings,
            )
    except Exception as exc:
        raise ServomexSinkWriteError(
            f"PostgresSink: failed to open pool for {cfg.target()}: {exc}"
        ) from exc

    _logger.info(
        "sinks.postgres.open target=%s pool_min=%s pool_max=%s use_copy=%s create_table=%s",
        cfg.target(),
        cfg.pool_min_size,
        cfg.pool_max_size,
        cfg.use_copy,
        cfg.create_table,
    )

    if not cfg.create_table:
        await self._introspect_existing_table()

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Append ``samples`` — one COPY (or executemany) per call."""
    if self._pool is None:
        raise RuntimeError("PostgresSink: write_many called before open()")
    if not samples:
        return

    rows = [sample_to_row(s) for s in samples]

    if not self._schema.is_locked:
        assert self._config.create_table  # noqa: S101
        self._schema.lock(rows)
        await self._create_table()
        self._insert_sql = self._build_insert_sql()

    columns = self._schema.columns
    assert columns is not None  # noqa: S101
    assert self._insert_sql is not None  # noqa: S101

    projected_tuples = [
        tuple(self._schema.project(row)[spec.name] for spec in columns) for row in rows
    ]

    try:
        if self._config.use_copy:
            await self._write_copy(projected_tuples, columns)
        else:
            await self._write_executemany(projected_tuples)
    except ServomexSinkWriteError:
        raise
    except Exception as exc:
        raise ServomexSinkWriteError(
            f"PostgresSink: write failed for {self._config.target()}: {exc}"
        ) from exc
    self._rows_written += len(projected_tuples)

async def __aenter__(self) -> Self:
    """Open the sink and return ``self`` for chaining."""
    ...

async def __aexit__(
    self,
    exc_type: type[BaseException] | None,
    exc: BaseException | None,
    tb: TracebackType | None,
) -> None:
    """Close the sink on exit."""
    ...

async def close(self) -> None:
    """Flush and release the backing resource. Idempotent."""
    ...

async def open(self) -> None:
    """Allocate the sink's backing resource (file handle, DB conn, ...)."""
    ...

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Append ``samples`` to the sink (``Sequence`` so the sink sees ``len``)."""
    ...

async def close(self) -> None:
    """Close the connection. Idempotent."""
    if self._conn is None:
        return
    conn = self._conn
    self._conn = None
    try:
        await run_sync(conn.close)
    finally:
        _logger.info("sinks.sqlite.close path=%s table=%s", str(self._path), self._table)

async def open(self) -> None:
    """Open the connection, apply PRAGMAs, and introspect the target. Idempotent."""
    if self._conn is not None:
        return
    self._conn = await run_sync(self._connect_blocking)
    _logger.info(
        "sinks.sqlite.open path=%s table=%s journal_mode=%s synchronous=%s",
        str(self._path),
        self._table,
        self._journal_mode,
        self._synchronous,
    )
    if not self._create_table:
        try:
            await run_sync(self._introspect_existing_table_blocking)
        except BaseException:
            conn = self._conn
            self._conn = None
            await run_sync(conn.close)
            raise

async def write_many(self, samples: Sequence[Sample]) -> None:
    """Append ``samples`` as rows in a single transaction (parameterised)."""
    if self._conn is None:
        raise RuntimeError("SqliteSink: write_many called before open()")
    if not samples:
        return

    rows = [sample_to_row(s) for s in samples]

    if not self._schema.is_locked and self._create_table:
        self._schema.lock(rows)
        await run_sync(self._create_table_blocking)
        self._insert_sql = self._build_insert_sql()

    assert self._insert_sql is not None  # noqa: S101 — narrow for type checker
    columns = self._schema.columns
    assert columns is not None  # noqa: S101

    projected = [
        tuple(self._schema.project(row)[spec.name] for spec in columns) for row in rows
    ]
    await run_sync(self._executemany_blocking, projected)