import time

import unittest

from feldera.enums import PipelineStatus

from feldera.pipeline import Pipeline

from feldera.runtime_config import Resources

from feldera.testutils import (

ViewSpec,

build_pipeline,

log,

validate_outputs,

unique_pipeline_name,

)

# Add now() value to a table in a transaction.

INPUT_RECORDS = 2000000

# After restart, the implicit `now` clock connector uses this resolution (see

# `now_endpoint_config` in the adapters crate). Stats only expose `stream` for

# the endpoint; the effective resolution is the pipeline runtime setting.

CLOCK_RESOLUTION_USECS_AFTER_RESTART = 3_000_000

def _clock_input_endpoint(pipeline: Pipeline):

"""Return stats for the built-in real-time clock input (`endpoint_name` ``now``)."""

for inp in pipeline.stats().inputs:

if inp.endpoint_name == "now":

return inp

names = [i.endpoint_name for i in pipeline.stats().inputs]

raise AssertionError(f"clock input endpoint 'now' not found; have {names!r}")

def _wait_clock_record_ticks(

pipeline: Pipeline, count: int, *, timeout_s: float

) -> None:

"""Wait until the clock connector's ingested record count increases by ``count``."""

baseline = _clock_input_endpoint(pipeline).metrics.total_records

assert baseline is not None

deadline = time.monotonic() + timeout_s

while time.monotonic() < deadline:

cur = _clock_input_endpoint(pipeline).metrics.total_records

if cur is not None and cur >= baseline + count:

return

time.sleep(0.25)

raise TimeoutError(

f"clock total_records did not increase by {count} within {timeout_s}s "

f"(baseline was {baseline})"

)

class TestNow(unittest.TestCase):

def test_now(self):

tables = {

"t": f"""

create table t(

id bigint not null primary key,

company_id bigint,

name string

) with (

'materialized' = 'true',

'connectors' = '[{{

"name": "datagen",

"transport": {{

"name": "datagen",

"config": {{

"plan": [{{

"limit": {INPUT_RECORDS},

"fields": {{

"name": {{ "strategy": "sentence" }}

}}

}}]

}}

}}

}}]');

"""

}

views = [

ViewSpec(

# Have a deterministic now() for testing against datafusion.

"v_now",

"""

select

now() as now_ts

""",

),

ViewSpec(

"v",

"""

select

t.*,

v_now.now_ts as now_ts

from t, v_now

""",

),

]

pipeline = build_pipeline(

unique_pipeline_name("now-test"),

tables,

views,

# This test uses >2GB of storage in the ad hoc query.

resources=Resources(storage_mb_max=8192),

)

pipeline.start()

start_time = time.monotonic()

pipeline.start_transaction()

while (

pipeline.stats().global_metrics.total_input_records < INPUT_RECORDS

or pipeline.stats().global_metrics.buffered_input_records > 0

):

log(f"Waiting for {INPUT_RECORDS} records to be ingested...")

time.sleep(1)

elapsed = time.monotonic() - start_time

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

# Freeze the value of now().

pipeline.pause()

pipeline.wait_for_idle()

start_time = time.monotonic()

pipeline.commit_transaction(transaction_id=None, wait=True, timeout_s=600)

elapsed = time.monotonic() - start_time

log(f"Commit took {elapsed}")

# Don't validate v_now which depends on the real-time clock.

views = [view for view in views if view.name != "v_now"]

validate_outputs(pipeline, tables, views)

# Tighten clock resolution, restart, confirm runtime and clock endpoint,

# then wait for two more clock ticks before validating again.

pipeline.stop(force=False)

runtime_cfg = pipeline.runtime_config()

runtime_cfg.clock_resolution_usecs = CLOCK_RESOLUTION_USECS_AFTER_RESTART

pipeline.set_runtime_config(runtime_cfg)

pipeline.start()

pipeline.wait_for_status(PipelineStatus.RUNNING, timeout=300)

assert (

pipeline.runtime_config().clock_resolution_usecs

== CLOCK_RESOLUTION_USECS_AFTER_RESTART

)

clock_status = _clock_input_endpoint(pipeline)

assert clock_status.config is not None

assert clock_status.config.get("stream") == "now"

# Two ticks at 3s resolution need at least ~6s; allow slack for scheduling.

_wait_clock_record_ticks(pipeline, 2, timeout_s=30.0)

validate_outputs(pipeline, tables, views)

assert next(pipeline.query("select count(*) as cnt from v_now"))["cnt"] == 1

pipeline.stop(force=True)

if __name__ == "__main__":

unittest.main()