import logging

import os

import pathlib

import shutil

import subprocess

import tempfile

import time

import uuid

from collections.abc import Callable, Sequence

from dataclasses import dataclass

from urllib.parse import urlparse

MINIO_BUCKET = os.environ.get("CI_MINIO_BUCKET", "ci-tests")

MINIO_ENDPOINT = os.environ.get(

"CI_MINIO_ENDPOINT", "http://minio.minio.svc.cluster.local:9000"

)

MINIO_REGION = os.environ.get("CI_MINIO_REGION", "us-east-1")

KAFKA_BOOTSTRAP = os.environ.get(

"KAFKA_BOOTSTRAP_SERVERS", "ci-kafka-bootstrap.kafka:9092"

)

def env_truthy(name: str) -> bool:

"""Return True when an environment variable is set to a truthy value."""

value = os.environ.get(name)

return value is not None and value.lower() not in {"", "0", "false", "no"}

def required_env(name: str) -> str:

"""Return a required environment variable or raise a descriptive error."""

value = os.environ.get(name)

if value:

return value

raise RuntimeError(f"required environment variable '{name}' is not set")

def runs_in_ci() -> bool:

"""Return True when the test suite is running under CI."""

return env_truthy("CI")

@dataclass

class DeltaTestLocation:

"""Describe where the Delta sink writes test data and how to read it back."""

uri: str

connector_config: dict[str, object]

root_path: str

local_dir: pathlib.Path | None = None

# True when ``stable_subpath`` was used at construction time. ``cleanup()``

# honors this by leaving the directory in place so the next run reuses the

# cached fixture instead of paying to rebuild it.

stable: bool = False

@classmethod

def create(

cls,

pipeline_name: str,

*,

mode: str = "truncate",

stable_subpath: str | None = None,

) -> "DeltaTestLocation":

"""Use the local filesystem for local runs and MinIO-backed S3 in CI.

:param mode: Value of the connector's ``mode`` field. Output

connectors use ``"truncate"`` (the default); input connectors

should pass ``"snapshot"``.

:param stable_subpath: When set, locates the table at a fixed,

shared path (``_fixtures/<stable_subpath>``) that is *not*

namespaced by the pipeline name or commit SHA, so a fixture

built once is reused by every later run and every commit.

When unset, a fresh random subpath under ``pipeline_name`` is

used (the original behavior). Use a stable path only for

fixtures whose contents are deterministic across runs.

"""

if runs_in_ci():

access_key = required_env("CI_K8S_MINIO_ACCESS_KEY_ID")

secret_key = required_env("CI_K8S_MINIO_SECRET_ACCESS_KEY")

if stable_subpath is not None:

prefix = f"_fixtures/{stable_subpath}"

else:

prefix = f"{pipeline_name}/{uuid.uuid4().hex}"

root_path = f"{MINIO_BUCKET}/{prefix}"

minio_endpoint = MINIO_ENDPOINT.rstrip("/")

parsed_endpoint = urlparse(minio_endpoint)

if (

parsed_endpoint.scheme not in {"http", "https"}

or not parsed_endpoint.netloc

):

raise ValueError(

"CI_MINIO_ENDPOINT must be a full URL, e.g. "

"'http://minio.minio.svc.cluster.local:9000'"

)

return cls(

uri=f"s3://{root_path}",

connector_config={

"uri": f"s3://{root_path}",

"mode": mode,

"aws_access_key_id": access_key,

"aws_secret_access_key": secret_key,

"aws_region": MINIO_REGION,

"aws_endpoint": minio_endpoint,

"aws_allow_http": str(parsed_endpoint.scheme == "http").lower(),

},

root_path=root_path,

stable=stable_subpath is not None,

)

if stable_subpath is not None:

local_dir = pathlib.Path("/tmp/feldera_fixtures") / stable_subpath

local_dir.mkdir(parents=True, exist_ok=True)

else:

local_dir = pathlib.Path(

tempfile.mkdtemp(prefix=f"{pipeline_name}_delta_", dir="/tmp")

)

return cls(

uri=f"file://{local_dir}",

connector_config={

"uri": f"file://{local_dir}",

"mode": mode,

},

root_path=str(local_dir),

local_dir=local_dir,

stable=stable_subpath is not None,

)

def delta_storage_options(self) -> dict[str, str]:

"""Return `deltalake` storage_options derived from the connector config."""

return {

k: str(v)

for k, v in self.connector_config.items()

if k not in ("uri", "mode")

}

def _s3_filesystem(self):

"""Build a pyarrow ``S3FileSystem`` from the connector config.

Pyarrow imports are deferred to keep module-level test collection

cheap on hosts that don't read Delta tables.

"""

import pyarrow.fs as pafs

cfg = self.connector_config

endpoint = str(cfg["aws_endpoint"]).rstrip("/")

parsed_endpoint = urlparse(endpoint)

return pafs.S3FileSystem(

access_key=str(cfg["aws_access_key_id"]),

secret_key=str(cfg["aws_secret_access_key"]),

region=str(cfg["aws_region"]),

scheme=parsed_endpoint.scheme,

endpoint_override=parsed_endpoint.netloc,

)

def log_json_paths(self) -> list[str]:

"""List Delta transaction log JSON files in version order."""

if self.local_dir is not None:

return [

str(path)

for path in sorted((self.local_dir / "_delta_log").glob("*.json"))

]

import pyarrow.fs as pafs

fs = self._s3_filesystem()

infos = fs.get_file_info(

pafs.FileSelector(f"{self.root_path}/_delta_log", recursive=False)

)

return sorted(

info.path

for info in infos

if info.type == pafs.FileType.File and info.path.endswith(".json")

)

def _read_text(self, path: str) -> str:

"""Read a text file from the configured backend."""

if self.local_dir is not None:

return pathlib.Path(path).read_text(encoding="utf-8")

with self._s3_filesystem().open_input_file(path) as handle:

return handle.readall().decode("utf-8")

def _read_parquet(self, relative_path: str):

"""Read a Delta data file (relative to ``root_path``) as a pyarrow.Table."""

import pyarrow.parquet as pq

if self.local_dir is not None:

return pq.read_table(self.local_dir / relative_path)

return pq.read_table(

f"{self.root_path}/{relative_path}", filesystem=self._s3_filesystem()

)

def read_rows(self) -> list[dict]:

"""Read the active rows of the Delta table by replaying its log.

Walks ``_delta_log/*.json``, follows ``add``/``remove`` actions to

derive the current set of parquet files, reads them with pyarrow

and drops Feldera-internal ``__feldera_*`` columns. Works against

both local-filesystem and S3-backed (MinIO) tables; avoids the

``deltalake`` Python package, whose wheel aborts on aarch64 hosts.

"""

import json as _json

import pyarrow as pa

active: dict[str, None] = {}

for log_path in self.log_json_paths():

for line in self._read_text(log_path).splitlines():

action = _json.loads(line)

if (add := action.get("add")) is not None:

active[add["path"]] = None

if (remove := action.get("remove")) is not None:

active.pop(remove["path"], None)

if not active:

return []

tables = [self._read_parquet(rel) for rel in sorted(active)]

return [

{

key: value

for key, value in row.items()

if not key.startswith("__feldera_")

}

for row in pa.concat_tables(tables).to_pylist()

]

def row_count(self, missing_ok: bool = False) -> int:

"""Return the row count of the current Delta snapshot.

Uses the per-file `numRecords` stats recorded in the delta log, so

we never need to scan parquet (and never need pyarrow). The

`deltalake` import is deferred to keep module-level test collection

cheap on hosts that don't read delta tables.

:param missing_ok: When True, return ``-1`` if the table does not

exist instead of raising. Useful when probing a cache.

"""

from deltalake import DeltaTable

from deltalake.exceptions import TableNotFoundError

try:

dt = DeltaTable(self.uri, storage_options=self.delta_storage_options())

except TableNotFoundError:

if missing_ok:

return -1

raise

return dt.count()

def delta_log_exists(self) -> bool:

"""Return True when a Delta log is present at this location.

Used to decide whether a cached fixture (see ``stable_subpath``) can

be reused instead of rebuilt.

"""

try:

return len(self.log_json_paths()) > 0

except FileNotFoundError:

return False

def _place_tree(self, staging: pathlib.Path) -> None:

"""Copy a locally-built Delta table tree to where this location stores

its data: the local directory, or the S3/MinIO bucket, depending on

how this location was created.

Some fixtures can only be produced on the local filesystem (e.g. a

PySpark-written table). For a local target, any existing content at

``self.local_dir`` is deleted before the copy. S3/MinIO targets get

the data files first and ``_delta_log`` last, so a reader observing

the upload mid-flight never sees a log referencing a missing parquet.

"""

if self.local_dir is not None:

if self.local_dir.exists():

shutil.rmtree(self.local_dir)

shutil.copytree(staging, self.local_dir)

return

fs = self._s3_filesystem()

files = [path for path in sorted(staging.rglob("*")) if path.is_file()]

for path in sorted(files, key=lambda p: ("_delta_log" in p.parts, p.name)):

rel = path.relative_to(staging).as_posix()

with fs.open_output_stream(f"{self.root_path}/{rel}") as out:

out.write(path.read_bytes())

def cleanup(self) -> None:

"""Remove the local temp directory, if any.

No-op when ``stable_subpath`` was used at construction time:

the whole point of a stable path is to cache contents across

runs, so deleting it would defeat the cache. Also a no-op on

the CI/MinIO path: there is no local directory to remove, so any

objects this run wrote to the bucket are simply left in place.

The shared MinIO bucket is long-lived and these tests never delete

from it, so those objects accumulate; non-stable runs each use a

unique uuid prefix so they never collide, and the leftover volume

for this internal CI bucket is accepted rather than swept here.

"""

if self.local_dir is not None and not self.stable:

shutil.rmtree(self.local_dir, ignore_errors=True)

self.local_dir = None

def ensure_delta_spark_fixture(

loc: DeltaTestLocation,

builder_script: str | os.PathLike[str],

builder_args: Sequence[object] = (),

*,

delta_spark_spec: str = "delta-spark>=4.2,<5",

is_present: Callable[[DeltaTestLocation], bool] | None = None,

) -> None:

"""Ensure a PySpark-authored Delta fixture exists at ``loc`` (cached).

Some Delta features (deletion vectors, column-mapping schema evolution)

can only be written by Delta Spark, not by delta-rs or the ``deltalake``

wheel. This builds such a fixture once and reuses it:

* If the fixture is already present (``is_present``, defaulting to

:meth:`DeltaTestLocation.delta_log_exists`), do nothing — so a

``stable_subpath`` cache is reused across runs.

* Otherwise run ``builder_script`` in an isolated environment

(``uv run --no-project --with <delta_spark_spec> python <builder_script>

<staging> *builder_args``), staging into a temp dir so a half-finished

build can never leak into the upload, then place the tree onto ``loc``'s

backend. ``--no-project`` keeps the builder hermetic: it depends only on

``delta_spark_spec``, never on building the enclosing project.

The heavy PySpark + JVM stack is pulled only on this rare rebuild path.

:param builder_script: Path to a standalone script that writes a Delta

table to the directory given as its first argument.

:param builder_args: Extra positional arguments passed after the staging

directory. Each is stringified verbatim with ``str()`` — pass

primitives; ``None`` would become the literal string ``"None"``.

:param is_present: Predicate deciding whether the fixture already exists;

also re-checked after upload to catch partial uploads.

"""

present = (

is_present if is_present is not None else DeltaTestLocation.delta_log_exists

)

if present(loc):

return

if shutil.which("uv") is None:

raise RuntimeError(

"`uv` is required on PATH to build the PySpark Delta fixture "

f"(builder runs via `uv run --with {delta_spark_spec}`)."

)

staging = pathlib.Path(tempfile.mkdtemp(prefix="feldera_delta_fixture_"))

try:

subprocess.run(

[

"uv",

"run",

"--no-project",

"--with",

delta_spark_spec,

"python",

str(builder_script),

str(staging),

*(str(arg) for arg in builder_args),

],

check=True,

)

loc._place_tree(staging)

finally:

shutil.rmtree(staging, ignore_errors=True)

if not present(loc):

raise RuntimeError(

f"Delta fixture at {loc.uri} is still absent after the builder "

"ran and the tree was uploaded."

)

def wait_for_condition(

description: str,

predicate_func,

timeout_s: float | None,

poll_interval_s: float,

) -> None:

"""Poll ``predicate_func`` until it returns truthy or the timeout elapses.

:param description: Human-readable description used in timeout/errors.

:param predicate_func: Callable returning ``True`` when condition is met.

:param timeout_s: Maximum wait time in seconds. ``None`` means wait forever.

:param poll_interval_s: Poll interval in seconds.

:raises TimeoutError: If the condition is not met within ``timeout_s``.

"""

if timeout_s is not None and poll_interval_s > timeout_s:

raise ValueError(

f"poll interval ({poll_interval_s}s) cannot be larger than"

f" timeout ({timeout_s}s)"

)

timestamp_start_s = time.monotonic()

timestamp_deadline_s = (

timestamp_start_s + timeout_s if timeout_s is not None else float("inf")

)

attempt = 0

while True:

if time.monotonic() > timestamp_deadline_s:

raise TimeoutError(

f"timeout ({timeout_s:.1f}s) waiting for condition '{description}'"

)

attempt += 1

if predicate_func():

logging.debug(

f"condition '{description}' met after"

f" {time.monotonic() - timestamp_start_s:.1f}s"

)

return

time.sleep(poll_interval_s)