"Utility functions for writing tests against a Feldera instance."

import base64

import logging

import math

import os

import urllib.error

import urllib.parse

import urllib.request

import platform

import re

import time

import json

import unittest

from typing import List, Optional, cast

from datetime import datetime

from feldera._long_operation_warning import LongOperationWarning

from feldera.enums import CompilationProfile

from feldera.pipeline import Pipeline

from feldera.pipeline_builder import PipelineBuilder

from feldera.runtime_config import Resources, RuntimeConfig

from feldera.rest import FelderaClient

from feldera.rest._helpers import requests_verify_from_env

logger = logging.getLogger(__name__)

API_KEY = os.environ.get("FELDERA_API_KEY")

# OIDC authentication support

def _get_oidc_token():

"""Get OIDC token if environment is configured, otherwise return None"""

try:

from feldera.testutils_oidc import get_oidc_test_helper

oidc_helper = get_oidc_test_helper()

if oidc_helper is not None:

return oidc_helper.obtain_access_token()

except ImportError:

pass

return None

def _get_effective_api_key():

"""Get effective API key - OIDC token takes precedence over static API key"""

oidc_token = _get_oidc_token()

return oidc_token if oidc_token else API_KEY

# Audience -> (token, seconds since the epoch after which it is re-minted).

_oidc_token_cache: dict[str, tuple[str, float]] = {}

# Re-mint a token this many seconds before its own expiry, so a request issued

# just before the check cannot arrive after it.

_OIDC_REFRESH_MARGIN_SECONDS = 120.0

# GitHub's own token endpoint occasionally 503s or drops the connection; these

# are as transient as the pipeline-side errors FelderaClient already retries,

# so retry the same way rather than letting one blip fail the whole run.

_OIDC_MINT_RETRYABLE_STATUS_CODES = frozenset({429, 502, 503, 504})

_OIDC_MINT_MAX_RETRIES = 3

_OIDC_MINT_INITIAL_BACKOFF_SECONDS = 1.0

_OIDC_MINT_BACKOFF_MULTIPLIER = 2.0

def _mint_github_oidc_token(request: urllib.request.Request) -> str:

"""Issue the token-mint request, retrying transient failures."""

backoff = _OIDC_MINT_INITIAL_BACKOFF_SECONDS

for attempt in range(_OIDC_MINT_MAX_RETRIES + 1):

try:

with urllib.request.urlopen(request, timeout=30) as response:

return json.load(response)["value"]

except urllib.error.HTTPError as e:

if (

e.code not in _OIDC_MINT_RETRYABLE_STATUS_CODES

or attempt == _OIDC_MINT_MAX_RETRIES

):

raise

except urllib.error.URLError:

if attempt == _OIDC_MINT_MAX_RETRIES:

raise

time.sleep(backoff)

backoff *= _OIDC_MINT_BACKOFF_MULTIPLIER

raise AssertionError("unreachable") # loop always returns or raises

def _github_oidc_token() -> str:

"""A GitHub Actions ID token, re-minted shortly before it expires.

The SDK resolves this before every request, so the implementation has to be

cheap. Minting per request adds a round trip to GitHub each time, and a

suite that polls in loops across parallel workers sends enough of them to be

throttled, which results in a connection timeout rather than an error.

"""

audience = os.environ.get("FELDERA_OIDC_AUDIENCE", "")

cached = _oidc_token_cache.get(audience)

if cached is not None and time.time() < cached[1]:

return cached[0]

request_url = os.environ["ACTIONS_ID_TOKEN_REQUEST_URL"]

if audience:

request_url += "&audience=" + urllib.parse.quote(audience, safe="")

request = urllib.request.Request(request_url)

request.add_header(

"Authorization", f"bearer {os.environ['ACTIONS_ID_TOKEN_REQUEST_TOKEN']}"

)

token = _mint_github_oidc_token(request)

_oidc_token_cache[audience] = (

token,

_token_expiry(token) - _OIDC_REFRESH_MARGIN_SECONDS,

)

return token

def _token_expiry(token: str) -> float:

"""`exp` from a JWT payload, in seconds since the epoch.

Returns the current time if the payload cannot be read, which re-mints on

every call rather than serving a token past its expiry.

"""

try:

payload = token.split(".")[1]

payload += "=" * (-len(payload) % 4)

return float(json.loads(base64.urlsafe_b64decode(payload))["exp"])

except Exception:

return time.time()

def feldera_bearer_token() -> Optional[str]:

"""The bearer token to send with a request issued now.

A configured OIDC login flow wins, then a GitHub Actions ID token, then the

static API key. None when nothing is configured, which is how a local

instance without authentication runs. Callers that build their own requests

resolve this per request: under Actions the token expires well inside a test

run, and one read at import time starts returning 401 partway through.

"""

if os.environ.get("OIDC_TEST_ISSUER"):

return _get_effective_api_key()

if os.environ.get("ACTIONS_ID_TOKEN_REQUEST_URL"):

return _github_oidc_token()

return API_KEY

BASE_URL = os.environ.get("FELDERA_HOST") or "http://localhost:8080"

FELDERA_REQUESTS_VERIFY = requests_verify_from_env()

FELDERA_TEST_NUM_WORKERS = int(os.environ.get("FELDERA_TEST_NUM_WORKERS", "8"))

FELDERA_TEST_NUM_HOSTS = int(os.environ.get("FELDERA_TEST_NUM_HOSTS", "1"))

# Ad-hoc queries reserve 64 MiB per worker for sort spill before sorting a

# single row (SORT_SPILL_RESERVATION_BYTES in

# adapterlib/src/utils/datafusion.rs), so a `datafusion_memory_mb` pool

# smaller than `workers * 64 MiB` fails any ORDER BY/EXCEPT/hash-join ad-hoc

# query with "Resources exhausted", independent of data size. This holds

# regardless of `hosts`: the reservation is sized from the *total* configured

# `workers`, whether they all run on one host or are split across several

# (the coordinator's own ad-hoc engine also uses the total worker count as

# its partition count -- see crates/coord/src/adhoc.rs).

_ADHOC_SORT_RESERVATION_MB_PER_WORKER = (64 * 1024 * 1024) / 1_000_000

def min_datafusion_memory_mb(

workers: int = FELDERA_TEST_NUM_WORKERS,

hosts: int = FELDERA_TEST_NUM_HOSTS,

headroom_mb: int = 512,

) -> int:

"""Minimum `datafusion_memory_mb` for a sort-heavy ad-hoc query to avoid

"Resources exhausted" with the given number of workers, plus

`headroom_mb` for the query's actual data. `hosts` doesn't change the

result (see module comment above); it's accepted so callers can pass

both test parameters without worrying about which one matters.

"""

return math.ceil(workers * _ADHOC_SORT_RESERVATION_MB_PER_WORKER) + headroom_mb

class _LazyClient:

"Construct the FelderaClient only when accessed as opposed to when imported."

__slots__ = ("_client",)

def __init__(self):

self._client = None

def _ensure(self):

if self._client is None:

# Under Actions the token expires inside a run, so the SDK gets the

# resolver itself: it re-resolves per request and retries once on

# 401. Elsewhere the credential is fixed for the process.

self._client = FelderaClient(

connection_timeout=10,

api_key=(

feldera_bearer_token

if os.environ.get("ACTIONS_ID_TOKEN_REQUEST_URL")

else feldera_bearer_token()

),

)

return self._client

def __getattr__(self, name):

return getattr(self._ensure(), name)

def __call__(self, *a, **kw) -> FelderaClient:

return self._ensure()

TEST_CLIENT = cast(FelderaClient, _LazyClient())

# SQL index definition.

class IndexSpec:

def __init__(self, name: str, columns: List[str]):

self.name = name

self.columns = columns

def __repr__(self):

return f"IndexSpec(name={self.name!r},columns={self.columns!r})"

class ViewSpec:

"""

SQL view definition consisting of a query that can run in Feldera or

datafusion, optional connector spec and aux SQL statements, e.g., indexes

and lateness clauses following view definition.

"""

def __init__(

self,

name: str,

query: str,

indexes: List[IndexSpec] = [],

connectors: Optional[str] = None,

aux: Optional[str] = None,

expected_hash: Optional[str] = None,

):

if not isinstance(query, str):

raise TypeError("query must be a string")

self.name = name

self.query = query

self.connectors = connectors

self.indexes = indexes

self.aux = aux

self.expected_hash = expected_hash

def __repr__(self):

return f"ViewSpec(name={self.name!r}, query={self.query!r}, indexes={self.indexes!r}, connectors={self.connectors!r}, aux={self.aux!r}, expected_hash={self.expected_hash!r})"

def clone(self):

return ViewSpec(

self.name,

self.query,

self.indexes,

self.connectors,

self.aux,

self.expected_hash,

)

def clone_with_name(self, name: str):

return ViewSpec(name, self.query, self.indexes, self.connectors, self.aux)

def sql(self) -> str:

sql = ""

if self.connectors:

with_clause = f"\nwith('connectors' = '{self.connectors}')\n"

else:

with_clause = ""

sql += (

f"create materialized view {self.name}{with_clause} as\n{self.query};\n\n"

)

for index in self.indexes:

columns = ",".join(index.columns)

sql += f"create index {index.name} on {self.name}({columns});\n"

if self.aux:

sql += f"{self.aux}\n"

sql += "\n"

return sql

def log(*args, **kwargs):

"""Print like built-in print(), but prefix each line with current time."""

prefix = datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")

print(prefix, *args, **kwargs)

# Cap on rows logged when a view validation fails. A failing TPC-H or test_now run

# can otherwise emit tens of megabytes on a single line, which blows up CI log capture.

_ROW_DIFF_LOG_LIMIT = 20

def _log_row_diff(label: str, rows: list) -> None:

if not rows:

return

log(

f"{label} ({len(rows)} rows; showing first {min(len(rows), _ROW_DIFF_LOG_LIMIT)}):"

)

for row in rows[:_ROW_DIFF_LOG_LIMIT]:

log(json.dumps(row, default=str))

def unique_pipeline_name(base_name: str) -> str:

"""

In CI, multiple tests of different runs can run against the same Feldera instance, we

make sure the pipeline names they use are unique by appending the first 5 characters

of the commit SHA or 'local' if not in CI. FELDERA_TEST_TAG_SUFFIX distinguishes

test suites of the same commit running concurrently against one instance.

"""

ci_tag = os.getenv("GITHUB_SHA", "local")[:5] + os.getenv(

"FELDERA_TEST_TAG_SUFFIX", ""

)

name = f"{ci_tag}_{base_name}"

# The pipeline name becomes a Kubernetes label value (max 63 chars). Fail here

# with a clear message rather than letting provisioning hit a cryptic 422.

assert len(name) <= 62, (

f"Generated pipeline name '{name}' is {len(name)} chars, exceeding the 62-char "

f"limit; shorten the test name '{base_name}' to at most {62 - len(ci_tag) - 1} chars."

)

return name

# Teardown must not block on one wedged pipeline: the SDK otherwise polls for

# `Stopped` forever, and when the runner finally kills the job every pipeline

# the run started keeps consuming compute on the shared instance.

RECLAIM_TIMEOUT_SECONDS = 60.0

def reclaim_pipeline(name: str, delete: bool = False) -> List[str]:

"""Force-stop `name`, clear its storage, and optionally delete it.

Every step runs even when the one before it raised: a pipeline that never

reports `Stopped` can still release its storage, and abandoning the rest on

the first error leaves them running on a shared instance. Returns one

message per failed step, empty when the pipeline is fully reclaimed, so the

caller decides whether a cleanup failure is worth failing a test over.

"""

failures = []

try:

TEST_CLIENT.stop_pipeline(name, force=True, timeout_s=RECLAIM_TIMEOUT_SECONDS)

except Exception as error:

failures.append(f"{name}: stop: {error}")

try:

TEST_CLIENT.clear_storage(name, timeout_s=RECLAIM_TIMEOUT_SECONDS)

except Exception as error:

failures.append(f"{name}: clear storage: {error}")

if delete:

try:

TEST_CLIENT.delete_pipeline(name)

except Exception as error:

failures.append(f"{name}: delete: {error}")

return failures

def enterprise_only(fn):

fn._enterprise_only = True

return unittest.skipUnless(

TEST_CLIENT.get_config().edition.is_enterprise(),

f"{fn.__name__} is enterprise only, skipping",

)(fn)

def single_host_only(fn):

fn._single_host_only = True

return unittest.skipUnless(

FELDERA_TEST_NUM_HOSTS == 1,

f"multihost not yet supported for {fn.__name__}, skipping",

)(fn)

def skip_on_arm64(fn):

"""Skip the test on Linux aarch64.

Some native Python wheels (e.g. `deltalake`) bundle a jemalloc built

for 4 KB pages and abort on import on Ubuntu's 64 KB-page aarch64

kernel used by some CI runners. Tests that depend on such wheels can

use this decorator to opt out of running there.

"""

fn._skip_on_arm64 = True

return unittest.skipIf(

platform.machine() == "aarch64",

f"{fn.__name__} skipped on aarch64 (incompatible native wheel)",

)(fn)

def datafusionize(query: str) -> str:

sort_array_pattern = re.compile(re.escape("SORT_ARRAY"), re.IGNORECASE)

truncate_pattern = re.compile(re.escape("TRUNCATE"), re.IGNORECASE)

timestamp_trunc_pattern = re.compile(

r"TIMESTAMP_TRUNC\s*\(\s*MAKE_TIMESTAMP\s*\(\s*([^)]+)\s*\)\s*,\s*([A-Z]+)\s*\)",

re.IGNORECASE,

)

result = sort_array_pattern.sub("array_sort", query)

result = truncate_pattern.sub("trunc", result)

result = timestamp_trunc_pattern.sub(r"DATE_TRUNC('\2', TO_TIMESTAMP(\1))", result)

return result

def validate_view(pipeline: Pipeline, view: ViewSpec):

log(f"Validating view '{view.name}'")

# We have two modes to verify the view, either we run the same SQL as the view against datafusion

# by `datafusionizing` the query, or a weaker form where we pass a hash of what the result

# should look like and check that the hash hasn't changed

if view.expected_hash:

view_query = f"select * from {view.name}"

computed_hash = pipeline.query_hash(view_query)

if computed_hash != view.expected_hash:

raise AssertionError(

f"View {view.name} hash {computed_hash} was but expected hash {view.expected_hash}"

)

else:

# TODO: count records

view_query = datafusionize(view.query)

try:

extra_rows = list(

pipeline.query(f"(select * from {view.name}) except ({view_query})")

)

missing_rows = list(

pipeline.query(f"({view_query}) except (select * from {view.name})")

)

_log_row_diff(

"Extra rows in Feldera output, but not in the ad hoc query output",

extra_rows,

)

_log_row_diff(

"Extra rows in the ad hoc query output, but not in Feldera output",

missing_rows,

)

except Exception as e:

log(f"Error querying view '{view.name}': {e}")

log(f"Ad-hoc Query: {view_query}")

raise

if extra_rows or missing_rows:

raise AssertionError(f"Validation failed for view {view.name}")

def generate_program(tables: dict, views: List[ViewSpec]) -> str:

sql = ""

for table_sql in tables.values():

sql += f"{table_sql}\n"

for view in views:

sql += view.sql()

return sql

def build_pipeline(

pipeline_name: str,

tables: dict,

views: List[ViewSpec],

resources: Optional[Resources] = None,

dev_tweaks: Optional[dict] = None,

datafusion_memory_mb: Optional[int] = None,

) -> Pipeline:

sql = generate_program(tables, views)

pipeline = PipelineBuilder(

TEST_CLIENT,

pipeline_name,

sql=sql,

compilation_profile=CompilationProfile.OPTIMIZED,

runtime_config=RuntimeConfig(

# Covers node auto-provisioning: a pipeline that needs a fresh

# node shape waits for node boot plus image pull.

provisioning_timeout_secs=180,

resources=resources,

workers=FELDERA_TEST_NUM_WORKERS,

hosts=FELDERA_TEST_NUM_HOSTS,

dev_tweaks=dev_tweaks,

datafusion_memory_mb=datafusion_memory_mb,

),

).create_or_replace()

return pipeline

def validate_outputs(pipeline: Pipeline, tables: dict, views: List[ViewSpec]):

for table in tables.keys():

row_count = list(pipeline.query(f"select count(*) from {table}"))

log(f"Table '{table}' count(*):\n{row_count}")

for view in views:

validate_view(pipeline, view)

def check_end_of_input(pipeline: Pipeline) -> bool:

return all(

input_endpoint.metrics.end_of_input

for input_endpoint in pipeline.stats().inputs

)

def wait_end_of_input(pipeline: Pipeline, timeout_s: Optional[int] = None):

start_time = time.monotonic()

# Reused by the warning message below so a stalled ingest can be told

# apart from a slow one; check_end_of_input() would fetch and discard it.

latest_stats = None

wait_for_input_end = LongOperationWarning(

logger,

lambda elapsed: (

f"still waiting for end of input on pipeline {pipeline.name}, "

f"waited {elapsed:.1f} seconds ({latest_stats.global_metrics.progress_summary()})"

),

lambda elapsed: (

f"end of input reached on pipeline {pipeline.name} "

f"after {elapsed:.1f} seconds"

),

)

while True:

latest_stats = pipeline.stats()

if all(

input_endpoint.metrics.end_of_input

for input_endpoint in latest_stats.inputs

):

wait_for_input_end.done()

return

if timeout_s is not None and time.monotonic() - start_time > timeout_s:

raise TimeoutError("Timeout waiting for end of input")

wait_for_input_end.check()

time.sleep(3)

def transaction(pipeline: Pipeline, duration_seconds: int):

"""Run a transaction for a specified duration."""

log(f"Running transaction for {duration_seconds} seconds")

pipeline.start_transaction()

time.sleep(duration_seconds)

log("Committing transaction")

commit_start = time.monotonic()

pipeline.commit_transaction()

log(f"Transaction committed in {time.monotonic() - commit_start} seconds")

def transaction_num_records(pipeline: Pipeline, num_records: int):

"""Run a transaction until it ingests a record count or reaches end of input."""

log(f"Running transaction for {num_records} records or end of input")

initial_records = number_of_processed_records(pipeline)

pipeline.start_transaction()

while not check_end_of_input(pipeline):

processed_records = number_of_processed_records(pipeline) - initial_records

if processed_records >= num_records:

break

time.sleep(3)

log("Committing transaction")

commit_start = time.monotonic()

pipeline.commit_transaction()

log(f"Transaction committed in {time.monotonic() - commit_start} seconds")

def checkpoint_pipeline(pipeline: Pipeline):

"""Create a checkpoint and wait for it to complete."""

log("Creating checkpoint")

checkpoint_start = time.monotonic()

pipeline.checkpoint(wait=True)

log(f"Checkpoint complete in {time.monotonic() - checkpoint_start} seconds")

def check_for_endpoint_errors(pipeline: Pipeline):

"""Check for errors on all input and output endpoints."""

for input_endpoint_status in pipeline.stats().inputs:

input_endpoint_status.metrics

if input_endpoint_status.metrics.num_transport_errors > 0:

raise RuntimeError(

f"Transport errors detected on input endpoint: {input_endpoint_status.endpoint_name}"

)

if input_endpoint_status.metrics.num_parse_errors > 0:

raise RuntimeError(

f"Parse errors on input endpoint: {input_endpoint_status.endpoint_name}"

)

log(f" Input endpoint {input_endpoint_status.endpoint_name} OK")

for output_endpoint_status in pipeline.stats().outputs:

output_endpoint_status.metrics

if output_endpoint_status.metrics.num_transport_errors > 0:

raise RuntimeError(

f"Transport errors detected on output endpoint: {output_endpoint_status.endpoint_name}"

)

if output_endpoint_status.metrics.num_encode_errors > 0:

raise RuntimeError(

f"Encode errors on output endpoint: {output_endpoint_status.endpoint_name}"

)

log(f" Output endpoint {output_endpoint_status.endpoint_name} OK")

def number_of_processed_records(pipeline: Pipeline) -> int:

"""Get the total_processed_records metric."""

return pipeline.stats().global_metrics.total_processed_records

def number_of_input_records(pipeline: Pipeline) -> int:

"""Get the total_input_records metric."""

return pipeline.stats().global_metrics.total_input_records

def run_workload(

pipeline_name: str,

tables: dict,

views: List[ViewSpec],

transaction: bool = True,

stop: bool = True,

resources: Optional[Resources] = None,

) -> Pipeline:

"""

Helper to run a pipeline to completion and validate the views afterwards using ad-hoc queries.

Use this for large-scale workload and standard benchmarks (like TPC-H etc.) where you plan to

ingest a lot of data and validate the results. For testing more specific functionality, see

frameworks in the `tests` directory.

"""

pipeline = build_pipeline(pipeline_name, tables, views, resources)

try:

pipeline.start()

start_time = time.monotonic()

if transaction:

try:

pipeline.start_transaction()

except Exception as e:

log(f"Error starting transaction: {e}")

if transaction:

wait_end_of_input(pipeline, timeout_s=3600)

else:

pipeline.wait_for_completion(force_stop=False, timeout_s=3600)

elapsed = time.monotonic() - start_time

log(f"Data ingested in {elapsed}")

if transaction:

start_time = time.monotonic()

try:

pipeline.commit_transaction(

transaction_id=None, wait=True, timeout_s=None

)

log(f"Commit took {time.monotonic() - start_time}")

except Exception as e:

log(f"Error committing transaction: {e}")

log("Waiting for outputs to flush")

start_time = time.monotonic()

pipeline.wait_for_completion(force_stop=False, timeout_s=3600)

log(f"Flushing outputs took {time.monotonic() - start_time}")

validate_outputs(pipeline, tables, views)

finally:

if stop:

try:

pipeline.stop(force=True)

pipeline.clear_storage()

except Exception as e:

log(f"Error during pipeline cleanup: {e}")

return pipeline