sqlmesh-dag-generator 0.9.5__tar.gz → 0.9.7__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. {sqlmesh_dag_generator-0.9.5/sqlmesh_dag_generator.egg-info → sqlmesh_dag_generator-0.9.7}/PKG-INFO +38 -3
  2. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/README.md +37 -2
  3. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/setup.py +1 -1
  4. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/__init__.py +1 -1
  5. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/config.py +2 -2
  6. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/generator.py +158 -2
  7. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7/sqlmesh_dag_generator.egg-info}/PKG-INFO +38 -3
  8. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_config.py +10 -0
  9. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_generator.py +86 -0
  10. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/LICENSE +0 -0
  11. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/pyproject.toml +0 -0
  12. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/setup.cfg +0 -0
  13. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/airflow_utils.py +0 -0
  14. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/cli.py +0 -0
  15. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/dag_builder.py +0 -0
  16. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/models.py +0 -0
  17. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/security.py +0 -0
  18. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/utils.py +0 -0
  19. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator/validation.py +0 -0
  20. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator.egg-info/SOURCES.txt +0 -0
  21. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator.egg-info/dependency_links.txt +0 -0
  22. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator.egg-info/entry_points.txt +0 -0
  23. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator.egg-info/requires.txt +0 -0
  24. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/sqlmesh_dag_generator.egg-info/top_level.txt +0 -0
  25. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_airflow_utils.py +0 -0
  26. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_auto_schedule.py +0 -0
  27. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_kubernetes.py +0 -0
  28. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_models.py +0 -0
  29. {sqlmesh_dag_generator-0.9.5 → sqlmesh_dag_generator-0.9.7}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlmesh-dag-generator
3
- Version: 0.9.5
3
+ Version: 0.9.7
4
4
  Summary: Open-source Airflow DAG generator for SQLMesh projects
5
5
  Home-page: https://github.com/kubolko/sqlmesh-dag-generator
6
6
  Author: Jakub Sumionka
@@ -128,9 +128,9 @@ If you run sub-hourly incremental models with `catchup=False`, outages can leave
128
128
 
129
129
  The package now supports an explicit recovery policy:
130
130
 
131
- - `recovery_mode="disabled"` (default): no runtime recovery tasks are added.
131
+ - `recovery_mode="disabled"`: no runtime recovery tasks are added.
132
132
  - `recovery_mode="warn"`: add an integrity guard task that detects missing intervals and logs them.
133
- - `recovery_mode="bounded_auto"`: add the same guard task plus a bounded recovery task that replays missing intervals when the gap is within `recovery_max_intervals`.
133
+ - `recovery_mode="bounded_auto"` (default): add the same guard task plus a bounded recovery task that replays missing intervals when the gap is within `recovery_max_intervals`.
134
134
 
135
135
  Example:
136
136
 
@@ -150,6 +150,41 @@ When `recovery_mode` is enabled, the package adds stable helper tasks to the DAG
150
150
 
151
151
  This keeps recovery explicit and observable in Airflow while preserving the default "no surprise backfills" behavior.
152
152
 
153
+ ### Manual Backfill Task
154
+
155
+ For larger historical gaps, the package also exposes a first-class manual backfill task helper.
156
+ Use it in a separate unscheduled DAG and trigger it only when you need to replay a broader window:
157
+
158
+ ```python
159
+ from datetime import datetime, timedelta
160
+ from airflow import DAG
161
+ from sqlmesh_dag_generator import SQLMeshDAGGenerator
162
+
163
+ generator = SQLMeshDAGGenerator(
164
+ sqlmesh_project_path="/path/to/project",
165
+ dag_id="my_pipeline",
166
+ gateway="prod",
167
+ )
168
+
169
+ with DAG(
170
+ dag_id="my_pipeline_manual_backfill",
171
+ schedule=None,
172
+ start_date=datetime(2024, 1, 1),
173
+ catchup=False,
174
+ ) as dag:
175
+ generator.create_manual_backfill_task(
176
+ dag,
177
+ default_start="2024-01-01T00:00:00",
178
+ execution_timeout=timedelta(hours=12),
179
+ )
180
+ ```
181
+
182
+ The task reads optional `dag_run.conf` overrides:
183
+
184
+ - `start`: ISO-8601 start boundary
185
+ - `end`: ISO-8601 end boundary. If omitted, the task backfills up to the current UTC time.
186
+ - `models`: one model name or a list of model names. If omitted, the task backfills all models.
187
+
153
188
  ## 💡 What You Get
154
189
 
155
190
  ### Your SQLMesh Project:
@@ -81,9 +81,9 @@ If you run sub-hourly incremental models with `catchup=False`, outages can leave
81
81
 
82
82
  The package now supports an explicit recovery policy:
83
83
 
84
- - `recovery_mode="disabled"` (default): no runtime recovery tasks are added.
84
+ - `recovery_mode="disabled"`: no runtime recovery tasks are added.
85
85
  - `recovery_mode="warn"`: add an integrity guard task that detects missing intervals and logs them.
86
- - `recovery_mode="bounded_auto"`: add the same guard task plus a bounded recovery task that replays missing intervals when the gap is within `recovery_max_intervals`.
86
+ - `recovery_mode="bounded_auto"` (default): add the same guard task plus a bounded recovery task that replays missing intervals when the gap is within `recovery_max_intervals`.
87
87
 
88
88
  Example:
89
89
 
@@ -103,6 +103,41 @@ When `recovery_mode` is enabled, the package adds stable helper tasks to the DAG
103
103
 
104
104
  This keeps recovery explicit and observable in Airflow while preserving the default "no surprise backfills" behavior.
105
105
 
106
+ ### Manual Backfill Task
107
+
108
+ For larger historical gaps, the package also exposes a first-class manual backfill task helper.
109
+ Use it in a separate unscheduled DAG and trigger it only when you need to replay a broader window:
110
+
111
+ ```python
112
+ from datetime import datetime, timedelta
113
+ from airflow import DAG
114
+ from sqlmesh_dag_generator import SQLMeshDAGGenerator
115
+
116
+ generator = SQLMeshDAGGenerator(
117
+ sqlmesh_project_path="/path/to/project",
118
+ dag_id="my_pipeline",
119
+ gateway="prod",
120
+ )
121
+
122
+ with DAG(
123
+ dag_id="my_pipeline_manual_backfill",
124
+ schedule=None,
125
+ start_date=datetime(2024, 1, 1),
126
+ catchup=False,
127
+ ) as dag:
128
+ generator.create_manual_backfill_task(
129
+ dag,
130
+ default_start="2024-01-01T00:00:00",
131
+ execution_timeout=timedelta(hours=12),
132
+ )
133
+ ```
134
+
135
+ The task reads optional `dag_run.conf` overrides:
136
+
137
+ - `start`: ISO-8601 start boundary
138
+ - `end`: ISO-8601 end boundary. If omitted, the task backfills up to the current UTC time.
139
+ - `models`: one model name or a list of model names. If omitted, the task backfills all models.
140
+
106
141
  ## 💡 What You Get
107
142
 
108
143
  ### Your SQLMesh Project:
@@ -20,7 +20,7 @@ if requirements_file.exists():
20
20
 
21
21
  setup(
22
22
  name="sqlmesh-dag-generator",
23
- version="0.9.5",
23
+ version="0.9.7",
24
24
  description="Open-source Airflow DAG generator for SQLMesh projects",
25
25
  long_description=long_description,
26
26
  long_description_content_type="text/markdown",
@@ -2,7 +2,7 @@
2
2
  SQLMesh DAG Generator - Open Source Airflow Integration for SQLMesh
3
3
  """
4
4
 
5
- __version__ = "0.9.5"
5
+ __version__ = "0.9.7"
6
6
 
7
7
 
8
8
  from sqlmesh_dag_generator.generator import SQLMeshDAGGenerator
@@ -84,7 +84,7 @@ class SQLMeshConfig:
84
84
  @dataclass
85
85
  class RecoveryConfig:
86
86
  """Optional runtime recovery behavior for missed Airflow intervals."""
87
- mode: str = "disabled" # disabled, warn, bounded_auto
87
+ mode: str = "bounded_auto" # disabled, warn, bounded_auto
88
88
  max_intervals: int = 6
89
89
  fail_on_excess_gap: bool = False
90
90
 
@@ -137,7 +137,7 @@ class GenerationConfig:
137
137
  dry_run: bool = False
138
138
  include_source_tables: bool = True # Include upstream source tables as dummy tasks
139
139
  return_value: bool = True # Whether to return execution result (XCom)
140
- auto_replan_on_change: bool = False # Automatically replan if models change
140
+ auto_replan_on_change: bool = True # Automatically plan+apply before running models
141
141
  replan_timeout_hours: int = 6 # Timeout for replan task (initial backfills can be long!)
142
142
  skip_audits: bool = False # Skip audit checks during execution
143
143
  enable_health_check: bool = False # Add a pre-flight health check task
@@ -1,7 +1,7 @@
1
1
  """
2
2
  Core DAG generator module
3
3
  """
4
- from datetime import datetime, timedelta
4
+ from datetime import datetime, timedelta, timezone
5
5
  import logging
6
6
  from pathlib import Path
7
7
  from typing import Dict, Optional, Union, Any, List
@@ -182,7 +182,7 @@ class SQLMeshDAGGenerator:
182
182
  model_pattern=kwargs.get("model_pattern"),
183
183
  include_source_tables=kwargs.get("include_source_tables", True), # Default: enabled
184
184
  return_value=kwargs.get("return_value", True),
185
- auto_replan_on_change=kwargs.get("auto_replan_on_change", False),
185
+ auto_replan_on_change=kwargs.get("auto_replan_on_change", True),
186
186
  replan_timeout_hours=kwargs.get("replan_timeout_hours", 6), # Default: 6 hours for backfills
187
187
  skip_audits=kwargs.get("skip_audits", False),
188
188
  enable_health_check=kwargs.get("enable_health_check", False),
@@ -1337,6 +1337,162 @@ class SQLMeshDAGGenerator:
1337
1337
  all_tasks = {**tasks, **source_table_tasks}
1338
1338
  return all_tasks
1339
1339
 
1340
+ def _build_runtime_context_kwargs(self) -> Dict[str, Any]:
1341
+ """Build SQLMesh Context kwargs using merged runtime connection config when available."""
1342
+ context_kwargs = {
1343
+ "paths": self.config.sqlmesh.project_path,
1344
+ }
1345
+
1346
+ if self.merged_config is not None:
1347
+ context_kwargs["config"] = self.merged_config
1348
+ if self.runtime_gateway is not None:
1349
+ context_kwargs["gateway"] = self.runtime_gateway
1350
+ else:
1351
+ context_kwargs["gateway"] = self.config.sqlmesh.gateway
1352
+
1353
+ return context_kwargs
1354
+
1355
+ @staticmethod
1356
+ def _parse_manual_backfill_datetime(value: Union[str, datetime], field_name: str) -> datetime:
1357
+ """Parse a manual backfill boundary into a timezone-aware UTC datetime."""
1358
+ if isinstance(value, datetime):
1359
+ parsed = value
1360
+ else:
1361
+ raw_value = str(value).strip()
1362
+ if not raw_value:
1363
+ raise ValueError(f"Manual backfill requires a non-empty '{field_name}' value.")
1364
+ parsed = datetime.fromisoformat(raw_value.replace("Z", "+00:00"))
1365
+
1366
+ if parsed.tzinfo is None:
1367
+ return parsed.replace(tzinfo=timezone.utc)
1368
+ return parsed.astimezone(timezone.utc)
1369
+
1370
+ @staticmethod
1371
+ def _normalize_manual_backfill_models(
1372
+ raw_models: Optional[Union[str, List[str]]]
1373
+ ) -> Optional[List[str]]:
1374
+ """Normalize manual backfill model selection into a list of model names."""
1375
+ if raw_models is None:
1376
+ return None
1377
+ if isinstance(raw_models, str):
1378
+ return [raw_models]
1379
+ if not isinstance(raw_models, list) or not all(isinstance(model, str) for model in raw_models):
1380
+ raise ValueError("Manual backfill 'models' must be a string or a list of model names.")
1381
+ if not raw_models:
1382
+ raise ValueError("Manual backfill 'models' cannot be an empty list.")
1383
+ return raw_models
1384
+
1385
+ def create_manual_backfill_task(
1386
+ self,
1387
+ dag,
1388
+ task_id: str = "sqlmesh_manual_backfill",
1389
+ default_start: Optional[Union[str, datetime]] = None,
1390
+ default_end: Optional[Union[str, datetime]] = None,
1391
+ default_models: Optional[Union[str, List[str]]] = None,
1392
+ execution_timeout: Optional[timedelta] = None,
1393
+ ):
1394
+ """
1395
+ Create a manual SQLMesh backfill task for ad hoc historical replay.
1396
+
1397
+ The returned task reads optional overrides from ``dag_run.conf``:
1398
+
1399
+ - ``start``: ISO-8601 start boundary
1400
+ - ``end``: ISO-8601 end boundary (defaults to current UTC time when omitted)
1401
+ - ``models``: single model name or list of model names to replay
1402
+
1403
+ Args:
1404
+ dag: Airflow DAG object.
1405
+ task_id: Task ID for the manual backfill operator.
1406
+ default_start: Default start boundary when ``dag_run.conf.start`` is omitted.
1407
+ default_end: Default end boundary when ``dag_run.conf.end`` is omitted.
1408
+ default_models: Optional default model selection.
1409
+ execution_timeout: Optional Airflow execution timeout for the task.
1410
+
1411
+ Returns:
1412
+ PythonOperator configured for manual SQLMesh backfill.
1413
+ """
1414
+ from airflow.exceptions import AirflowException
1415
+ from airflow.operators.python import PythonOperator
1416
+
1417
+ normalized_default_models = self._normalize_manual_backfill_models(default_models)
1418
+
1419
+ def run_manual_backfill(**context):
1420
+ dag_run = context.get("dag_run")
1421
+ conf = dag_run.conf if dag_run and dag_run.conf else {}
1422
+
1423
+ requested_start = conf.get("start", default_start)
1424
+ requested_end = conf.get("end", default_end)
1425
+ requested_models = conf.get("models", normalized_default_models)
1426
+
1427
+ if requested_start is None:
1428
+ raise AirflowException(
1429
+ "Manual backfill requires a start boundary. Provide default_start or dag_run.conf['start']."
1430
+ )
1431
+
1432
+ try:
1433
+ backfill_start = self._parse_manual_backfill_datetime(requested_start, "start")
1434
+ if requested_end is None:
1435
+ backfill_end = datetime.now(timezone.utc).replace(microsecond=0)
1436
+ else:
1437
+ backfill_end = self._parse_manual_backfill_datetime(requested_end, "end")
1438
+ selected_models = self._normalize_manual_backfill_models(requested_models)
1439
+ except ValueError as exc:
1440
+ raise AirflowException(str(exc)) from exc
1441
+
1442
+ if backfill_start >= backfill_end:
1443
+ raise AirflowException(
1444
+ f"Manual backfill start must be earlier than end. Got {backfill_start} >= {backfill_end}."
1445
+ )
1446
+
1447
+ logger.info(
1448
+ "Starting manual SQLMesh backfill for %s -> %s (models=%s)",
1449
+ backfill_start,
1450
+ backfill_end,
1451
+ selected_models or "ALL",
1452
+ )
1453
+
1454
+ self.load_sqlmesh_context()
1455
+ if selected_models is not None:
1456
+ if not self.models:
1457
+ self.extract_models()
1458
+ unknown_models = sorted(set(selected_models) - set(self.models))
1459
+ if unknown_models:
1460
+ raise AirflowException(
1461
+ f"Unknown SQLMesh models requested for manual backfill: {unknown_models}"
1462
+ )
1463
+
1464
+ run_ctx = Context(**self._build_runtime_context_kwargs())
1465
+
1466
+ run_kwargs = {
1467
+ "environment": self.config.sqlmesh.environment,
1468
+ "start": backfill_start,
1469
+ "end": backfill_end,
1470
+ }
1471
+ if selected_models is not None:
1472
+ run_kwargs["select_models"] = selected_models
1473
+
1474
+ run_sig = inspect.signature(run_ctx.run)
1475
+ if "skip_audits" in run_sig.parameters and self.config.generation.skip_audits:
1476
+ run_kwargs["skip_audits"] = True
1477
+
1478
+ result = run_ctx.run(**run_kwargs)
1479
+ status = result.name if hasattr(result, "name") else str(result)
1480
+ logger.info("Manual SQLMesh backfill finished with status: %s", status)
1481
+
1482
+ return {
1483
+ "status": status,
1484
+ "start": backfill_start.isoformat(),
1485
+ "end": backfill_end.isoformat(),
1486
+ "models": selected_models or "ALL",
1487
+ }
1488
+
1489
+ return PythonOperator(
1490
+ task_id=task_id,
1491
+ python_callable=run_manual_backfill,
1492
+ execution_timeout=execution_timeout,
1493
+ dag=dag,
1494
+ )
1495
+
1340
1496
  def _get_output_path(self) -> Path:
1341
1497
  """Get the output file path for the generated DAG"""
1342
1498
  output_dir = Path(self.config.generation.output_dir)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlmesh-dag-generator
3
- Version: 0.9.5
3
+ Version: 0.9.7
4
4
  Summary: Open-source Airflow DAG generator for SQLMesh projects
5
5
  Home-page: https://github.com/kubolko/sqlmesh-dag-generator
6
6
  Author: Jakub Sumionka
@@ -128,9 +128,9 @@ If you run sub-hourly incremental models with `catchup=False`, outages can leave
128
128
 
129
129
  The package now supports an explicit recovery policy:
130
130
 
131
- - `recovery_mode="disabled"` (default): no runtime recovery tasks are added.
131
+ - `recovery_mode="disabled"`: no runtime recovery tasks are added.
132
132
  - `recovery_mode="warn"`: add an integrity guard task that detects missing intervals and logs them.
133
- - `recovery_mode="bounded_auto"`: add the same guard task plus a bounded recovery task that replays missing intervals when the gap is within `recovery_max_intervals`.
133
+ - `recovery_mode="bounded_auto"` (default): add the same guard task plus a bounded recovery task that replays missing intervals when the gap is within `recovery_max_intervals`.
134
134
 
135
135
  Example:
136
136
 
@@ -150,6 +150,41 @@ When `recovery_mode` is enabled, the package adds stable helper tasks to the DAG
150
150
 
151
151
  This keeps recovery explicit and observable in Airflow while preserving the default "no surprise backfills" behavior.
152
152
 
153
+ ### Manual Backfill Task
154
+
155
+ For larger historical gaps, the package also exposes a first-class manual backfill task helper.
156
+ Use it in a separate unscheduled DAG and trigger it only when you need to replay a broader window:
157
+
158
+ ```python
159
+ from datetime import datetime, timedelta
160
+ from airflow import DAG
161
+ from sqlmesh_dag_generator import SQLMeshDAGGenerator
162
+
163
+ generator = SQLMeshDAGGenerator(
164
+ sqlmesh_project_path="/path/to/project",
165
+ dag_id="my_pipeline",
166
+ gateway="prod",
167
+ )
168
+
169
+ with DAG(
170
+ dag_id="my_pipeline_manual_backfill",
171
+ schedule=None,
172
+ start_date=datetime(2024, 1, 1),
173
+ catchup=False,
174
+ ) as dag:
175
+ generator.create_manual_backfill_task(
176
+ dag,
177
+ default_start="2024-01-01T00:00:00",
178
+ execution_timeout=timedelta(hours=12),
179
+ )
180
+ ```
181
+
182
+ The task reads optional `dag_run.conf` overrides:
183
+
184
+ - `start`: ISO-8601 start boundary
185
+ - `end`: ISO-8601 end boundary. If omitted, the task backfills up to the current UTC time.
186
+ - `models`: one model name or a list of model names. If omitted, the task backfills all models.
187
+
153
188
  ## 💡 What You Get
154
189
 
155
190
  ### Your SQLMesh Project:
@@ -10,6 +10,7 @@ from sqlmesh_dag_generator.config import (
10
10
  AirflowConfig,
11
11
  GenerationConfig,
12
12
  DAGGeneratorConfig,
13
+ RecoveryConfig,
13
14
  )
14
15
 
15
16
 
@@ -39,6 +40,15 @@ def test_airflow_config_creation():
39
40
  assert "test" in config.tags
40
41
 
41
42
 
43
+ def test_recovery_config_defaults_to_bounded_auto():
44
+ """Test default recovery mode enables bounded auto replay."""
45
+ config = RecoveryConfig()
46
+
47
+ assert config.mode == "bounded_auto"
48
+ assert config.max_intervals == 6
49
+ assert config.fail_on_excess_gap is False
50
+
51
+
42
52
  def test_generation_config_defaults():
43
53
  """Test GenerationConfig default values"""
44
54
  config = GenerationConfig()
@@ -383,6 +383,92 @@ class TestSQLMeshDAGGenerator:
383
383
  assert run_kwargs["select_models"] == ["dwh.raw_5m"]
384
384
  assert recovery_result["gap_intervals"] == 3
385
385
 
386
+ @patch("sqlmesh_dag_generator.generator.Context")
387
+ def test_create_manual_backfill_task_uses_defaults_and_conf_overrides(self, mock_context):
388
+ """Test manual backfill task supports default windows and dag_run.conf overrides."""
389
+ from airflow import DAG
390
+
391
+ run_ctx = MagicMock()
392
+ run_result = MagicMock()
393
+ run_result.name = "SUCCESS"
394
+ run_ctx.run.return_value = run_result
395
+ mock_context.return_value = run_ctx
396
+
397
+ generator = SQLMeshDAGGenerator(
398
+ sqlmesh_project_path="/tmp/project",
399
+ dag_id="test",
400
+ )
401
+ generator.load_sqlmesh_context = MagicMock(return_value=run_ctx)
402
+ generator.models = {
403
+ "dwh.raw_5m": SQLMeshModelInfo(
404
+ name="dwh.raw_5m",
405
+ dependencies=set(),
406
+ interval_unit="FIVE_MINUTE",
407
+ kind="INCREMENTAL_BY_TIME_RANGE",
408
+ ),
409
+ "dwh.report": SQLMeshModelInfo(
410
+ name="dwh.report",
411
+ dependencies={"dwh.raw_5m"},
412
+ interval_unit="HOUR",
413
+ kind="FULL",
414
+ ),
415
+ }
416
+
417
+ with DAG("test", start_date=datetime(2024, 1, 1)) as dag:
418
+ task = generator.create_manual_backfill_task(
419
+ dag,
420
+ default_start="2024-01-01T00:00:00",
421
+ )
422
+
423
+ dag_run = MagicMock()
424
+ dag_run.conf = {
425
+ "end": "2024-01-01T02:00:00",
426
+ "models": ["dwh.raw_5m"],
427
+ }
428
+ result = task.python_callable(dag_run=dag_run)
429
+
430
+ run_ctx.run.assert_called_once()
431
+ _, run_kwargs = run_ctx.run.call_args
432
+ assert run_kwargs["start"].isoformat() == "2024-01-01T00:00:00+00:00"
433
+ assert run_kwargs["end"].isoformat() == "2024-01-01T02:00:00+00:00"
434
+ assert run_kwargs["select_models"] == ["dwh.raw_5m"]
435
+ assert result["status"] == "SUCCESS"
436
+ assert result["models"] == ["dwh.raw_5m"]
437
+
438
+ def test_create_manual_backfill_task_rejects_unknown_models(self):
439
+ """Test manual backfill task validates requested model names."""
440
+ from airflow import DAG
441
+ from airflow.exceptions import AirflowException
442
+
443
+ generator = SQLMeshDAGGenerator(
444
+ sqlmesh_project_path="/tmp/project",
445
+ dag_id="test",
446
+ )
447
+ generator.load_sqlmesh_context = MagicMock()
448
+ generator.models = {
449
+ "dwh.raw_5m": SQLMeshModelInfo(
450
+ name="dwh.raw_5m",
451
+ dependencies=set(),
452
+ interval_unit="FIVE_MINUTE",
453
+ kind="INCREMENTAL_BY_TIME_RANGE",
454
+ ),
455
+ }
456
+
457
+ with DAG("test", start_date=datetime(2024, 1, 1)) as dag:
458
+ task = generator.create_manual_backfill_task(
459
+ dag,
460
+ default_start="2024-01-01T00:00:00",
461
+ )
462
+
463
+ dag_run = MagicMock()
464
+ dag_run.conf = {
465
+ "end": "2024-01-01T02:00:00",
466
+ "models": ["dwh.missing"],
467
+ }
468
+
469
+ with pytest.raises(AirflowException, match="Unknown SQLMesh models"):
470
+ task.python_callable(dag_run=dag_run)
471
+
386
472
 
387
473
  class TestDynamicFeatures:
388
474
  """Test dynamic DAG specific features"""