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

import os

import re

import time

import json

import unittest

from typing import List, Optional, cast

from datetime import datetime

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

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

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"))

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:

self._client = FelderaClient(

connection_timeout=10,

api_key=_get_effective_api_key(),

)

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)

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.

"""

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

return f"{ci_tag}_{base_name}"

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 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})")

)

if extra_rows:

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

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

if missing_rows:

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

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

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,

) -> Pipeline:

sql = generate_program(tables, views)

pipeline = PipelineBuilder(

TEST_CLIENT,

pipeline_name,

sql=sql,

compilation_profile=CompilationProfile.OPTIMIZED,

runtime_config=RuntimeConfig(

provisioning_timeout_secs=60,

resources=resources,

workers=FELDERA_TEST_NUM_WORKERS,

hosts=FELDERA_TEST_NUM_HOSTS,

),

).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()

while not check_end_of_input(pipeline):

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

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

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 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 run_workload(

pipeline_name: str, tables: dict, views: List[ViewSpec], transaction: bool = True

):

"""

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)

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)

pipeline.stop(force=True)