sqlmesh-dag-generator 0.9.7__tar.gz → 0.9.9__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.
- {sqlmesh_dag_generator-0.9.7/sqlmesh_dag_generator.egg-info → sqlmesh_dag_generator-0.9.9}/PKG-INFO +1 -1
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/setup.py +1 -1
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/__init__.py +1 -1
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/config.py +1 -1
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/generator.py +52 -8
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9/sqlmesh_dag_generator.egg-info}/PKG-INFO +1 -1
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_generator.py +76 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/LICENSE +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/README.md +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/pyproject.toml +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/setup.cfg +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/airflow_utils.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/cli.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/dag_builder.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/models.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/security.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/utils.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/validation.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator.egg-info/SOURCES.txt +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator.egg-info/dependency_links.txt +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator.egg-info/entry_points.txt +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator.egg-info/requires.txt +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator.egg-info/top_level.txt +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_airflow_utils.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_auto_schedule.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_config.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_kubernetes.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_models.py +0 -0
- {sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/tests/test_utils.py +0 -0
|
@@ -20,7 +20,7 @@ if requirements_file.exists():
|
|
|
20
20
|
|
|
21
21
|
setup(
|
|
22
22
|
name="sqlmesh-dag-generator",
|
|
23
|
-
version="0.9.
|
|
23
|
+
version="0.9.9",
|
|
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",
|
|
@@ -138,7 +138,7 @@ class GenerationConfig:
|
|
|
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
140
|
auto_replan_on_change: bool = True # Automatically plan+apply before running models
|
|
141
|
-
replan_timeout_hours: int = 6 #
|
|
141
|
+
replan_timeout_hours: Optional[int] = 6 # None disables the replan task timeout.
|
|
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
|
|
144
144
|
# Tag-based filtering - only include models with any of these tags
|
{sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/generator.py
RENAMED
|
@@ -2,6 +2,7 @@
|
|
|
2
2
|
Core DAG generator module
|
|
3
3
|
"""
|
|
4
4
|
from datetime import datetime, timedelta, timezone
|
|
5
|
+
import inspect
|
|
5
6
|
import logging
|
|
6
7
|
from pathlib import Path
|
|
7
8
|
from typing import Dict, Optional, Union, Any, List
|
|
@@ -183,7 +184,7 @@ class SQLMeshDAGGenerator:
|
|
|
183
184
|
include_source_tables=kwargs.get("include_source_tables", True), # Default: enabled
|
|
184
185
|
return_value=kwargs.get("return_value", True),
|
|
185
186
|
auto_replan_on_change=kwargs.get("auto_replan_on_change", True),
|
|
186
|
-
replan_timeout_hours=kwargs.get("replan_timeout_hours", 6), #
|
|
187
|
+
replan_timeout_hours=kwargs.get("replan_timeout_hours", 6), # None disables the replan timeout.
|
|
187
188
|
skip_audits=kwargs.get("skip_audits", False),
|
|
188
189
|
enable_health_check=kwargs.get("enable_health_check", False),
|
|
189
190
|
include_tags=kwargs.get("include_tags"),
|
|
@@ -991,14 +992,25 @@ class SQLMeshDAGGenerator:
|
|
|
991
992
|
}
|
|
992
993
|
|
|
993
994
|
from datetime import timedelta as td
|
|
995
|
+
|
|
996
|
+
replan_timeout_hours = self.config.generation.replan_timeout_hours
|
|
994
997
|
replan_task = PythonOperator(
|
|
995
998
|
task_id="sqlmesh_plan_apply",
|
|
996
999
|
python_callable=run_replan,
|
|
997
1000
|
dag=dag,
|
|
998
|
-
|
|
999
|
-
|
|
1001
|
+
execution_timeout=(
|
|
1002
|
+
td(hours=replan_timeout_hours)
|
|
1003
|
+
if replan_timeout_hours is not None
|
|
1004
|
+
else None
|
|
1005
|
+
),
|
|
1000
1006
|
)
|
|
1001
|
-
|
|
1007
|
+
if replan_timeout_hours is None:
|
|
1008
|
+
logger.info("Created auto-replan task: sqlmesh_plan_apply (timeout: disabled)")
|
|
1009
|
+
else:
|
|
1010
|
+
logger.info(
|
|
1011
|
+
"Created auto-replan task: sqlmesh_plan_apply (timeout: %sh)",
|
|
1012
|
+
replan_timeout_hours,
|
|
1013
|
+
)
|
|
1002
1014
|
|
|
1003
1015
|
# Link health check to replan if both exist
|
|
1004
1016
|
if health_check_task:
|
|
@@ -1184,7 +1196,7 @@ class SQLMeshDAGGenerator:
|
|
|
1184
1196
|
task_id = model_info.get_task_id()
|
|
1185
1197
|
|
|
1186
1198
|
# Create the execution function
|
|
1187
|
-
def make_callable(m_name, m_fqn):
|
|
1199
|
+
def make_callable(m_name, m_fqn, m_interval_minutes=None):
|
|
1188
1200
|
def execute_model(**context):
|
|
1189
1201
|
from sqlmesh import Context
|
|
1190
1202
|
|
|
@@ -1212,16 +1224,42 @@ class SQLMeshDAGGenerator:
|
|
|
1212
1224
|
start = context.get('data_interval_start') or context.get('execution_date')
|
|
1213
1225
|
end = context.get('data_interval_end') or context.get('execution_date')
|
|
1214
1226
|
|
|
1227
|
+
# When this model's interval is COARSER than the DAG tick,
|
|
1228
|
+
# the tick's data_interval window (e.g. 5 minutes on a mixed
|
|
1229
|
+
# sub-hourly + hourly project) spans no full model interval.
|
|
1230
|
+
# SQLMesh then returns NOTHING_TO_DO and the model's state
|
|
1231
|
+
# silently freezes while the task still reports success.
|
|
1232
|
+
# In that case omit start/end so SQLMesh selects the correct
|
|
1233
|
+
# interval(s) from the model's own cron instead.
|
|
1234
|
+
if (
|
|
1235
|
+
m_interval_minutes
|
|
1236
|
+
and expected_interval_minutes
|
|
1237
|
+
and m_interval_minutes > expected_interval_minutes
|
|
1238
|
+
):
|
|
1239
|
+
logger.info(
|
|
1240
|
+
"Model %s interval (%s min) is coarser than the DAG tick "
|
|
1241
|
+
"(%s min); running without an explicit window so SQLMesh "
|
|
1242
|
+
"picks the due interval(s) from the model cron.",
|
|
1243
|
+
m_fqn, m_interval_minutes, expected_interval_minutes,
|
|
1244
|
+
)
|
|
1245
|
+
start = None
|
|
1246
|
+
end = None
|
|
1247
|
+
|
|
1215
1248
|
# Run the model with proper time range
|
|
1216
1249
|
try:
|
|
1217
1250
|
# Build run kwargs - only include skip_audits if enabled
|
|
1218
1251
|
# (some SQLMesh versions may not support this parameter)
|
|
1219
1252
|
run_kwargs = {
|
|
1220
1253
|
"environment": self.config.sqlmesh.environment,
|
|
1221
|
-
"start": start,
|
|
1222
|
-
"end": end,
|
|
1223
1254
|
"select_models": [m_fqn],
|
|
1224
1255
|
}
|
|
1256
|
+
# Only pass an explicit window when we have one. For
|
|
1257
|
+
# coarser-than-tick models start/end were cleared above so
|
|
1258
|
+
# SQLMesh selects the due interval(s) from the model cron.
|
|
1259
|
+
if start is not None:
|
|
1260
|
+
run_kwargs["start"] = start
|
|
1261
|
+
if end is not None:
|
|
1262
|
+
run_kwargs["end"] = end
|
|
1225
1263
|
|
|
1226
1264
|
# Check if skip_audits is supported by inspecting the method signature
|
|
1227
1265
|
import inspect
|
|
@@ -1299,7 +1337,13 @@ class SQLMeshDAGGenerator:
|
|
|
1299
1337
|
# Create PythonOperator
|
|
1300
1338
|
task = PythonOperator(
|
|
1301
1339
|
task_id=task_id,
|
|
1302
|
-
python_callable=make_callable(
|
|
1340
|
+
python_callable=make_callable(
|
|
1341
|
+
model_name,
|
|
1342
|
+
model_info.name,
|
|
1343
|
+
get_interval_frequency_minutes(model_info.interval_unit)
|
|
1344
|
+
if model_info.interval_unit is not None
|
|
1345
|
+
else None,
|
|
1346
|
+
),
|
|
1303
1347
|
dag=dag,
|
|
1304
1348
|
)
|
|
1305
1349
|
|
|
@@ -204,6 +204,82 @@ class TestSQLMeshDAGGenerator:
|
|
|
204
204
|
assert hasattr(task, 'task_id')
|
|
205
205
|
assert hasattr(task, 'python_callable')
|
|
206
206
|
|
|
207
|
+
@patch("sqlmesh.Context")
|
|
208
|
+
@patch("sqlmesh_dag_generator.generator.Context")
|
|
209
|
+
def test_execute_model_omits_window_for_coarser_than_tick_models(self, mock_mod_context, mock_sqlmesh_context):
|
|
210
|
+
"""A model whose interval is coarser than the DAG tick must not be run
|
|
211
|
+
with the tick's narrow window.
|
|
212
|
+
|
|
213
|
+
Root cause regression: when a project mixes sub-hourly and hourly models,
|
|
214
|
+
the DAG schedule is the global minimum (5 min). execute_model passed the
|
|
215
|
+
tick's data_interval (e.g. a 5-minute window) as start/end to EVERY model,
|
|
216
|
+
including hourly ones. SQLMesh returns NOTHING_TO_DO for an hourly model
|
|
217
|
+
given a sub-hour window that spans no full hourly boundary, so the hourly
|
|
218
|
+
model's state froze while its task still reported success. The fix: for a
|
|
219
|
+
model whose interval is coarser than the scheduler tick, omit start/end so
|
|
220
|
+
SQLMesh selects the correct interval(s) from the model cron.
|
|
221
|
+
"""
|
|
222
|
+
from airflow import DAG
|
|
223
|
+
|
|
224
|
+
run_ctx = MagicMock()
|
|
225
|
+
run_result = MagicMock()
|
|
226
|
+
run_result.name = "SUCCESS"
|
|
227
|
+
run_ctx.run.return_value = run_result
|
|
228
|
+
# execute_model does `from sqlmesh import Context` locally, so the
|
|
229
|
+
# sqlmesh.Context patch is the one that matters here.
|
|
230
|
+
mock_sqlmesh_context.return_value = run_ctx
|
|
231
|
+
mock_mod_context.return_value = run_ctx
|
|
232
|
+
|
|
233
|
+
generator = SQLMeshDAGGenerator(
|
|
234
|
+
sqlmesh_project_path="/tmp/project",
|
|
235
|
+
dag_id="test",
|
|
236
|
+
)
|
|
237
|
+
# Mixed project: 5-min model sets the DAG tick to 5 minutes; the hourly
|
|
238
|
+
# model is coarser than the tick.
|
|
239
|
+
generator.models = {
|
|
240
|
+
"dwh.raw_5m": SQLMeshModelInfo(
|
|
241
|
+
name="dwh.raw_5m",
|
|
242
|
+
dependencies=set(),
|
|
243
|
+
interval_unit="FIVE_MINUTE",
|
|
244
|
+
kind="INCREMENTAL_BY_TIME_RANGE",
|
|
245
|
+
),
|
|
246
|
+
"dwh.f_hourly": SQLMeshModelInfo(
|
|
247
|
+
name="dwh.f_hourly",
|
|
248
|
+
dependencies=set(),
|
|
249
|
+
interval_unit="HOUR",
|
|
250
|
+
kind="INCREMENTAL_BY_TIME_RANGE",
|
|
251
|
+
),
|
|
252
|
+
}
|
|
253
|
+
|
|
254
|
+
with DAG("test", start_date=datetime(2024, 1, 1)) as dag:
|
|
255
|
+
tasks = generator.create_tasks_in_dag(dag)
|
|
256
|
+
|
|
257
|
+
# 5-minute tick window (what the scheduler hands every task).
|
|
258
|
+
tick_start = datetime(2024, 1, 6, 12, 25)
|
|
259
|
+
tick_end = datetime(2024, 1, 6, 12, 30)
|
|
260
|
+
|
|
261
|
+
# Hourly model (coarser than tick): must be run WITHOUT start/end.
|
|
262
|
+
run_ctx.run.reset_mock()
|
|
263
|
+
tasks["dwh.f_hourly"].python_callable(
|
|
264
|
+
data_interval_start=tick_start, data_interval_end=tick_end,
|
|
265
|
+
)
|
|
266
|
+
run_ctx.run.assert_called_once()
|
|
267
|
+
_, hourly_kwargs = run_ctx.run.call_args
|
|
268
|
+
assert hourly_kwargs.get("start") is None
|
|
269
|
+
assert hourly_kwargs.get("end") is None
|
|
270
|
+
assert hourly_kwargs["select_models"] == ["dwh.f_hourly"]
|
|
271
|
+
|
|
272
|
+
# 5-min model (matches tick): keeps the explicit window (unchanged).
|
|
273
|
+
run_ctx.run.reset_mock()
|
|
274
|
+
tasks["dwh.raw_5m"].python_callable(
|
|
275
|
+
data_interval_start=tick_start, data_interval_end=tick_end,
|
|
276
|
+
)
|
|
277
|
+
run_ctx.run.assert_called_once()
|
|
278
|
+
_, raw_kwargs = run_ctx.run.call_args
|
|
279
|
+
assert raw_kwargs["start"] == tick_start
|
|
280
|
+
assert raw_kwargs["end"] == tick_end
|
|
281
|
+
assert raw_kwargs["select_models"] == ["dwh.raw_5m"]
|
|
282
|
+
|
|
207
283
|
def test_static_dag_generation(self, demo_sqlmesh_project):
|
|
208
284
|
"""Test static DAG generation (alternative mode)"""
|
|
209
285
|
config = DAGGeneratorConfig.from_dict({
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/airflow_utils.py
RENAMED
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/dag_builder.py
RENAMED
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/security.py
RENAMED
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.9.7 → sqlmesh_dag_generator-0.9.9}/sqlmesh_dag_generator/validation.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|