feat(airflow): добавлен пульт управления генератором
- Зачем: - нужен основной ручной интерфейс стенда для backfill/import/check без консольной матрицы переменных. - Что: - добавлен DAG generator_control с параметрами Airflow, ветвлением операций и ожиданием ETL. - вынесена общая логика запуска и предпроверок генератора для Airflow. - обновлены compose-настройки, зависимости, тесты и документация по пульту. - Проверка: - uv run --with pytest --with-requirements generator/requirements.txt pytest generator/tests -q. - docker compose config --quiet.
This commit is contained in:
@@ -1,5 +1,5 @@
|
||||
# Kafka клиент
|
||||
kafka-python==2.0.5
|
||||
kafka-python==2.0.6
|
||||
|
||||
# Prometheus метрики
|
||||
prometheus-client==0.21.1
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
"""Чистая логика пульта Airflow для стартовой истории."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Iterator
|
||||
|
||||
from clickstream_generator.config import Config
|
||||
from clickstream_generator.kafka_io import (
|
||||
KafkaStateManager,
|
||||
KafkaStartupHistoryManifest,
|
||||
ensure_topics,
|
||||
)
|
||||
from clickstream_generator.launch import build_launch_env
|
||||
from clickstream_generator.service import GeneratorService
|
||||
from clickstream_generator.startup_history_artifact import (
|
||||
KafkaRawPublisher,
|
||||
KafkaTopicInspector,
|
||||
compare_clickhouse_stats_to_manifest,
|
||||
import_startup_history_artifact,
|
||||
load_startup_history_artifact,
|
||||
validate_startup_history_artifact,
|
||||
)
|
||||
|
||||
|
||||
AIRFLOW_DATA_DIR = "/opt/airflow/data"
|
||||
KAFKA_BOOTSTRAP_SERVERS = "kafka:29092"
|
||||
GENERATOR_METRICS_URL = "http://generator:9109/metrics"
|
||||
|
||||
WORLD_OVERRIDE_KEYS = {
|
||||
"GEN_SEED",
|
||||
"GEN_MODEL_T0",
|
||||
"GEN_MODEL_TIMEZONE",
|
||||
"GEN_MODEL_TIME_SPEED",
|
||||
"GEN_TICK_SECONDS",
|
||||
"GEN_LAMBDA_BASE_PER_MIN",
|
||||
"GEN_JITTER_PCT",
|
||||
"GEN_MIN_EVENTS_PER_TICK",
|
||||
"GEN_MAX_EVENTS_PER_TICK",
|
||||
}
|
||||
|
||||
STG_EMPTY_SQL = """
|
||||
SELECT
|
||||
(SELECT count() FROM stg.browser_raw) AS browser_raw,
|
||||
(SELECT count() FROM stg.location_raw) AS location_raw,
|
||||
(SELECT count() FROM stg.device_raw) AS device_raw,
|
||||
(SELECT count() FROM stg.geo_raw) AS geo_raw
|
||||
"""
|
||||
|
||||
CLICKHOUSE_STATS_SQL = """
|
||||
WITH
|
||||
toDateTime64('{model_t0}', 6) AS t0,
|
||||
toDateTime64('{model_t_end}', 6) AS t_end
|
||||
SELECT
|
||||
count() AS events,
|
||||
uniqExact(click_id) AS visits,
|
||||
uniqExact(user_domain_id) AS users,
|
||||
toString(min(event_ts)) AS min_event_timestamp,
|
||||
toString(max(event_ts)) AS max_event_timestamp
|
||||
FROM dm.v_events_enriched
|
||||
WHERE event_ts >= t0 AND event_ts < t_end
|
||||
"""
|
||||
|
||||
|
||||
def build_control_env(
|
||||
operation: str,
|
||||
*,
|
||||
profile_name: str,
|
||||
duration: str | None = None,
|
||||
artifact_path: str | None = None,
|
||||
overrides: dict[str, str] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Готовит env для операции пульта без доступа к Docker."""
|
||||
if operation not in {"backfill", "import"}:
|
||||
raise ValueError("operation must be backfill or import")
|
||||
if operation == "import" and not artifact_path:
|
||||
raise ValueError("artifact_path is required for import")
|
||||
|
||||
selected_overrides = {
|
||||
key: value
|
||||
for key, value in (overrides or {}).items()
|
||||
if key in WORLD_OVERRIDE_KEYS and value != ""
|
||||
}
|
||||
env = build_launch_env(
|
||||
"backfill",
|
||||
profile_name=profile_name,
|
||||
duration=duration or None,
|
||||
overrides=selected_overrides,
|
||||
)
|
||||
env["KAFKA_BOOTSTRAP_SERVERS"] = KAFKA_BOOTSTRAP_SERVERS
|
||||
env["GEN_DATA_DIR"] = AIRFLOW_DATA_DIR
|
||||
env["GEN_METRICS_ENABLED"] = "false"
|
||||
if artifact_path:
|
||||
env["GEN_STARTUP_HISTORY_ARTIFACT"] = artifact_path
|
||||
return env
|
||||
|
||||
|
||||
def assert_stand_clean(kafka_bootstrap_servers: str, clickhouse_hook) -> None:
|
||||
"""Проверяет, что backfill/import не смешает миры."""
|
||||
try:
|
||||
KafkaTopicInspector(kafka_bootstrap_servers).assert_data_topics_empty()
|
||||
except RuntimeError as exc:
|
||||
raise RuntimeError(
|
||||
"Стенд не чистый: в Kafka data-топиках уже есть сообщения. "
|
||||
"Выполните make clean с консоли и повторите операцию."
|
||||
) from exc
|
||||
assert_stg_tables_empty(clickhouse_hook)
|
||||
|
||||
|
||||
def assert_stg_tables_empty(clickhouse_hook) -> None:
|
||||
"""Падает, если в STG уже есть строки."""
|
||||
result = clickhouse_hook.execute(STG_EMPTY_SQL)
|
||||
counts = result[0] if result else ()
|
||||
names = ("stg.browser_raw", "stg.location_raw", "stg.device_raw", "stg.geo_raw")
|
||||
dirty = [
|
||||
f"{name}={int(count)}"
|
||||
for name, count in zip(names, counts)
|
||||
if int(count) > 0
|
||||
]
|
||||
if dirty:
|
||||
raise RuntimeError(
|
||||
"Стенд не чистый: в STG уже есть строки ("
|
||||
+ ", ".join(dirty)
|
||||
+ "). Выполните make clean с консоли и повторите операцию."
|
||||
)
|
||||
|
||||
|
||||
def assert_live_generator_not_running(url: str = GENERATOR_METRICS_URL) -> None:
|
||||
"""Мягко предупреждает о live-сервисе по HTTP-метрикам."""
|
||||
try:
|
||||
urllib.request.urlopen(url, timeout=2).close()
|
||||
except (urllib.error.URLError, TimeoutError, OSError):
|
||||
return
|
||||
raise RuntimeError(
|
||||
"Live-генератор отвечает на metrics-порту. Остановите его с консоли "
|
||||
"перед backfill/import, чтобы не смешать миры."
|
||||
)
|
||||
|
||||
|
||||
def run_backfill(env: dict[str, str]) -> None:
|
||||
"""Запускает backfill в процессе Airflow worker."""
|
||||
with _patched_environ(env):
|
||||
config = Config()
|
||||
GeneratorService(config).start()
|
||||
|
||||
|
||||
def run_import(env: dict[str, str], artifact_path: str) -> dict:
|
||||
"""Импортирует портативный артефакт в Kafka."""
|
||||
with _patched_environ(env):
|
||||
config = Config()
|
||||
artifact = load_startup_history_artifact(artifact_path)
|
||||
ensure_topics(config.kafka_bootstrap_servers)
|
||||
publisher = KafkaRawPublisher(config.kafka_bootstrap_servers)
|
||||
state_manager = KafkaStateManager(config.kafka_bootstrap_servers)
|
||||
manifest_manager = KafkaStartupHistoryManifest(config.kafka_bootstrap_servers)
|
||||
topic_inspector = KafkaTopicInspector(config.kafka_bootstrap_servers)
|
||||
try:
|
||||
return import_startup_history_artifact(
|
||||
artifact,
|
||||
publisher=publisher,
|
||||
state_manager=state_manager,
|
||||
manifest_manager=manifest_manager,
|
||||
expected_config=config,
|
||||
topic_inspector=topic_inspector,
|
||||
)
|
||||
finally:
|
||||
publisher.close()
|
||||
state_manager.close()
|
||||
manifest_manager.close()
|
||||
|
||||
|
||||
def validate_import_artifact(env: dict[str, str], artifact_path: str) -> None:
|
||||
"""Проверяет артефакт и настройки import до записи в Kafka."""
|
||||
with _patched_environ(env):
|
||||
config = Config()
|
||||
artifact = load_startup_history_artifact(artifact_path)
|
||||
validate_startup_history_artifact(artifact, expected_config=config)
|
||||
|
||||
|
||||
def load_manifest_from_kafka(
|
||||
kafka_bootstrap_servers: str = KAFKA_BOOTSTRAP_SERVERS,
|
||||
) -> dict:
|
||||
"""Читает manifest стартовой истории из compact-топика."""
|
||||
manager = KafkaStartupHistoryManifest(kafka_bootstrap_servers)
|
||||
try:
|
||||
manifest = manager.load()
|
||||
finally:
|
||||
manager.close()
|
||||
if not manifest:
|
||||
raise RuntimeError("Manifest стартовой истории не найден в Kafka.")
|
||||
return manifest
|
||||
|
||||
|
||||
def assert_clickhouse_matches_manifest(manifest: dict, clickhouse_hook) -> None:
|
||||
"""Сверяет контрольные числа ClickHouse с manifest."""
|
||||
model_t0 = _clickhouse_datetime_literal(str(manifest["model_t0"]))
|
||||
model_t_end = _clickhouse_datetime_literal(str(manifest["model_t_end"]))
|
||||
result = clickhouse_hook.execute(
|
||||
CLICKHOUSE_STATS_SQL.format(model_t0=model_t0, model_t_end=model_t_end)
|
||||
)
|
||||
if not result:
|
||||
raise RuntimeError("ClickHouse не вернул контрольные числа.")
|
||||
row = result[0]
|
||||
stats = {
|
||||
"events": str(row[0]),
|
||||
"visits": str(row[1]),
|
||||
"users": str(row[2]),
|
||||
"min_event_timestamp": str(row[3]),
|
||||
"max_event_timestamp": str(row[4]),
|
||||
}
|
||||
mismatches = compare_clickhouse_stats_to_manifest(manifest, stats)
|
||||
if mismatches:
|
||||
raise RuntimeError(
|
||||
"ClickHouse расходится с manifest: " + ", ".join(mismatches)
|
||||
)
|
||||
|
||||
|
||||
def default_artifact_path(operation: str) -> str:
|
||||
"""Возвращает путь артефакта по умолчанию в общем томе data."""
|
||||
filename = (
|
||||
"startup-history-import.json"
|
||||
if operation == "import"
|
||||
else "startup-history.json"
|
||||
)
|
||||
return str(Path(AIRFLOW_DATA_DIR) / filename)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_environ(env: dict[str, str]) -> Iterator[None]:
|
||||
old_values = {key: os.environ.get(key) for key in env}
|
||||
os.environ.update(env)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for key, value in old_values.items():
|
||||
if value is None:
|
||||
os.environ.pop(key, None)
|
||||
else:
|
||||
os.environ[key] = value
|
||||
|
||||
|
||||
def _clickhouse_datetime_literal(value: str) -> str:
|
||||
normalized = value.replace("T", " ").removesuffix("Z")
|
||||
if len(normalized) >= 6 and normalized[-6] in "+-" and normalized[-3] == ":":
|
||||
normalized = normalized[:-6]
|
||||
return normalized
|
||||
@@ -67,6 +67,10 @@ class Config:
|
||||
metrics_port: int = field(
|
||||
default_factory=lambda: int(os.getenv("GEN_METRICS_PORT", "9109"))
|
||||
)
|
||||
metrics_enabled: bool = field(
|
||||
default_factory=lambda: os.getenv("GEN_METRICS_ENABLED", "true").lower()
|
||||
== "true"
|
||||
)
|
||||
state_enabled: bool = field(
|
||||
default_factory=lambda: os.getenv("GEN_STATE_ENABLED", "true").lower() == "true"
|
||||
)
|
||||
|
||||
@@ -62,8 +62,11 @@ class GeneratorService:
|
||||
logger.warning("Generator is disabled (GEN_ENABLED=false)")
|
||||
return
|
||||
|
||||
logger.info(f"Starting metrics server on port {self.config.metrics_port}")
|
||||
start_http_server(self.config.metrics_port)
|
||||
if self.config.metrics_enabled:
|
||||
logger.info(f"Starting metrics server on port {self.config.metrics_port}")
|
||||
start_http_server(self.config.metrics_port)
|
||||
else:
|
||||
logger.info("Metrics HTTP server is disabled")
|
||||
|
||||
logger.info("Starting generator service...")
|
||||
logger.info(
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
"""
|
||||
Тесты чистой логики пульта Airflow для генератора.
|
||||
"""
|
||||
|
||||
from dataclasses import replace
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_build_control_env_uses_profile_duration_and_airflow_data_dir():
|
||||
"""Пульт строит env для backfill из профиля и каталога Airflow."""
|
||||
from clickstream_generator.airflow_control import build_control_env
|
||||
|
||||
env = build_control_env(
|
||||
"backfill",
|
||||
profile_name="ci",
|
||||
duration="",
|
||||
artifact_path="/opt/airflow/data/startup-history.json",
|
||||
overrides={"GEN_LAMBDA_BASE_PER_MIN": "120", "GEN_SEED": "99"},
|
||||
)
|
||||
|
||||
assert env["GEN_RUN_MODE"] == "backfill"
|
||||
assert env["GEN_HISTORY_DURATION"] == "6h"
|
||||
assert env["GEN_MODEL_T_END"] == "2026-01-01T06:00:00+00:00"
|
||||
assert env["GEN_DATA_DIR"] == "/opt/airflow/data"
|
||||
assert env["GEN_METRICS_ENABLED"] == "false"
|
||||
assert env["GEN_STARTUP_HISTORY_ARTIFACT"] == "/opt/airflow/data/startup-history.json"
|
||||
assert env["GEN_LAMBDA_BASE_PER_MIN"] == "120"
|
||||
assert env["GEN_SEED"] == "99"
|
||||
|
||||
|
||||
def test_import_env_requires_artifact_path_and_uses_backfill_contract():
|
||||
"""Import валидируется как тот же мир, что backfill."""
|
||||
from clickstream_generator.airflow_control import build_control_env
|
||||
|
||||
with pytest.raises(ValueError, match="artifact_path"):
|
||||
build_control_env("import", profile_name="ci", artifact_path="")
|
||||
|
||||
env = build_control_env(
|
||||
"import",
|
||||
profile_name="daily-wave",
|
||||
artifact_path="/opt/airflow/data/history.json",
|
||||
)
|
||||
|
||||
assert env["GEN_RUN_MODE"] == "backfill"
|
||||
assert env["GEN_STATE_RESET"] == "true"
|
||||
assert env["GEN_HISTORY_DURATION"] == "2d"
|
||||
|
||||
|
||||
def test_assert_stand_clean_rejects_non_empty_stg_before_writes():
|
||||
"""Backfill/import не стартуют на непустом STG."""
|
||||
from clickstream_generator.airflow_control import assert_stg_tables_empty
|
||||
|
||||
class Hook:
|
||||
def execute(self, sql):
|
||||
self.sql = sql
|
||||
return [(0, 2, 0, 0)]
|
||||
|
||||
hook = Hook()
|
||||
|
||||
with pytest.raises(RuntimeError, match="make clean"):
|
||||
assert_stg_tables_empty(hook)
|
||||
assert "stg.browser_raw" in hook.sql
|
||||
assert "stg.location_raw" in hook.sql
|
||||
|
||||
|
||||
def test_assert_stand_clean_rejects_non_empty_kafka_with_make_clean_hint(monkeypatch):
|
||||
"""Непустые Kafka-топики дают ту же подсказку про make clean."""
|
||||
from clickstream_generator import airflow_control
|
||||
|
||||
class Inspector:
|
||||
def __init__(self, bootstrap_servers):
|
||||
self.bootstrap_servers = bootstrap_servers
|
||||
|
||||
def assert_data_topics_empty(self):
|
||||
raise RuntimeError("Kafka data topics are not empty: browser_events")
|
||||
|
||||
class Hook:
|
||||
def execute(self, sql):
|
||||
return [(0, 0, 0, 0)]
|
||||
|
||||
monkeypatch.setattr(airflow_control, "KafkaTopicInspector", Inspector)
|
||||
|
||||
with pytest.raises(RuntimeError, match="make clean"):
|
||||
airflow_control.assert_stand_clean("kafka:29092", Hook())
|
||||
|
||||
|
||||
def test_check_manifest_compares_clickhouse_stats(base_config):
|
||||
"""Check падает, когда контрольные числа ClickHouse расходятся с manifest."""
|
||||
from clickstream_generator.airflow_control import assert_clickhouse_matches_manifest
|
||||
from clickstream_generator.startup_history_artifact import (
|
||||
StartupHistoryArtifactBuilder,
|
||||
build_manifest,
|
||||
)
|
||||
from test_startup_history_artifact import _batch, _state
|
||||
|
||||
state = _state()
|
||||
builder = StartupHistoryArtifactBuilder()
|
||||
builder.add_batch(_batch())
|
||||
manifest = build_manifest(
|
||||
config=replace(
|
||||
base_config,
|
||||
model_t0=state.model_t0,
|
||||
model_t_end=state.model_timestamp,
|
||||
model_time_speed=1,
|
||||
model_timezone="UTC",
|
||||
seed=42,
|
||||
),
|
||||
counters=builder.counters,
|
||||
state=state,
|
||||
)
|
||||
|
||||
class Hook:
|
||||
def execute(self, sql):
|
||||
self.sql = sql
|
||||
return [(
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
"2026-01-01 00:00:00.000000",
|
||||
"2026-01-01 00:00:00.000000",
|
||||
)]
|
||||
|
||||
with pytest.raises(RuntimeError, match="events"):
|
||||
assert_clickhouse_matches_manifest(manifest, Hook())
|
||||
|
||||
|
||||
def test_generator_metrics_server_can_be_disabled(base_config, monkeypatch):
|
||||
"""Airflow-задача может запускать генератор без HTTP-сервера метрик."""
|
||||
from clickstream_generator.service import GeneratorService
|
||||
|
||||
config = replace(base_config, enabled=False, metrics_enabled=False)
|
||||
service = GeneratorService(config)
|
||||
called = False
|
||||
|
||||
def start_http_server(_port):
|
||||
nonlocal called
|
||||
called = True
|
||||
|
||||
monkeypatch.setattr(
|
||||
"clickstream_generator.service.start_http_server",
|
||||
start_http_server,
|
||||
)
|
||||
|
||||
service.start()
|
||||
|
||||
assert called is False
|
||||
@@ -0,0 +1,76 @@
|
||||
"""
|
||||
Контракт DAG generator_control без запуска Airflow.
|
||||
"""
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
from clickstream_generator.launch import PROFILES
|
||||
|
||||
|
||||
DAG_PATH = Path(__file__).parents[2] / "airflow" / "dags" / "generator_control_dag.py"
|
||||
REPO_ROOT = Path(__file__).parents[2]
|
||||
|
||||
|
||||
def _tree():
|
||||
return ast.parse(DAG_PATH.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def test_dag_file_exists_and_uses_dynamic_profiles():
|
||||
"""DAG берёт варианты профилей из PROFILES, а не из ручного списка."""
|
||||
text = DAG_PATH.read_text(encoding="utf-8")
|
||||
|
||||
assert "dag_id=\"generator_control\"" in text
|
||||
assert "sorted(PROFILES)" in text
|
||||
for profile in PROFILES:
|
||||
assert profile not in {"hardcoded-profile"}
|
||||
|
||||
|
||||
def test_trigger_form_has_expected_param_enums():
|
||||
"""Форма запуска ограничивает операции и профили."""
|
||||
text = DAG_PATH.read_text(encoding="utf-8")
|
||||
|
||||
assert 'enum=["backfill", "import", "check"]' in text
|
||||
assert "enum=sorted(PROFILES)" in text
|
||||
assert '"duration": Param(' in text
|
||||
assert '"artifact_path": Param(' in text
|
||||
|
||||
|
||||
def test_dag_branches_and_waits_for_etl_completion():
|
||||
"""Backfill/import запускают ETL и ждут его завершения перед check."""
|
||||
text = DAG_PATH.read_text(encoding="utf-8")
|
||||
|
||||
assert "BranchPythonOperator" in text
|
||||
assert "TriggerDagRunOperator" in text
|
||||
assert 'trigger_dag_id="etl_pipeline"' in text
|
||||
assert "wait_for_completion=True" in text
|
||||
assert 'allowed_states=["success"]' in text
|
||||
assert 'failed_states=["failed"]' in text
|
||||
|
||||
|
||||
def test_no_docker_or_continue_operation_in_dag():
|
||||
"""Пульт не управляет Docker и не содержит операцию continue."""
|
||||
text = DAG_PATH.read_text(encoding="utf-8")
|
||||
|
||||
assert "docker" not in text.lower()
|
||||
assert '"continue"' not in text
|
||||
|
||||
|
||||
def test_compose_mounts_generator_code_without_socket_and_gates_live_service():
|
||||
"""Airflow видит код генератора, но не получает Docker socket."""
|
||||
text = (REPO_ROOT / "docker-compose.yml").read_text(encoding="utf-8")
|
||||
|
||||
assert "PYTHONPATH: /opt/airflow/generator_src" in text
|
||||
assert "./generator/src:/opt/airflow/generator_src:ro" in text
|
||||
assert "/var/run/docker.sock" not in text
|
||||
assert "profiles:\n - live-generator" in text
|
||||
|
||||
|
||||
def test_airflow_and_generator_kafka_dependency_versions_match():
|
||||
"""Airflow и генератор используют одну версию kafka-python."""
|
||||
airflow_req = (REPO_ROOT / "airflow" / "requirements.txt").read_text(encoding="utf-8")
|
||||
generator_req = (REPO_ROOT / "generator" / "requirements.txt").read_text(encoding="utf-8")
|
||||
|
||||
assert "kafka-python==2.0.6" in airflow_req
|
||||
assert "kafka-python==2.0.6" in generator_req
|
||||
assert "prometheus-client==0.21.1" in airflow_req
|
||||
Reference in New Issue
Block a user