sqlmesh-dag-generator 0.7.2__tar.gz → 0.7.3__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.7.2/sqlmesh_dag_generator.egg-info → sqlmesh_dag_generator-0.7.3}/PKG-INFO +9 -1
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/README.md +8 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/setup.py +1 -1
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/__init__.py +1 -1
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/airflow_utils.py +2 -2
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/config.py +14 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/dag_builder.py +71 -5
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/generator.py +93 -4
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3/sqlmesh_dag_generator.egg-info}/PKG-INFO +9 -1
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/SOURCES.txt +1 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_airflow_utils.py +9 -5
- sqlmesh_dag_generator-0.7.3/tests/test_enterprise_features.py +240 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/LICENSE +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/pyproject.toml +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/setup.cfg +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/cli.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/models.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/security.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/utils.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/validation.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/dependency_links.txt +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/entry_points.txt +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/requires.txt +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/top_level.txt +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_auto_schedule.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_bug_fix_user_scenario.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_cache.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_config.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_environment_default.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_environment_handling.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_generator.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_kubernetes.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_models.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_runtime_connections.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_source_tables.py +0 -0
- {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_utils.py +0 -0
{sqlmesh_dag_generator-0.7.2/sqlmesh_dag_generator.egg-info → sqlmesh_dag_generator-0.7.3}/PKG-INFO
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: sqlmesh-dag-generator
|
|
3
|
-
Version: 0.7.
|
|
3
|
+
Version: 0.7.3
|
|
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
|
|
@@ -63,6 +63,14 @@ Transform your SQLMesh models into production-ready Airflow DAGs with **full dat
|
|
|
63
63
|
- 🛠️ **Dual Mode**: Dynamic (auto-discovery, default) or Static (full control)
|
|
64
64
|
- 🚫 **No Vendor Lock-in**: Open source, no cloud dependencies
|
|
65
65
|
|
|
66
|
+
### 🏢 Enterprise Features (NEW in v0.8.0)
|
|
67
|
+
|
|
68
|
+
- 🔔 **Callbacks**: `on_failure_callback`, `on_success_callback`, `sla_miss_callback` for alerting
|
|
69
|
+
- 🏷️ **Tag-Based Filtering**: `include_tags`, `exclude_tags` for Data Mesh team-specific DAGs
|
|
70
|
+
- 🎱 **Pool Configuration**: `pool`, `pool_slots` for resource management
|
|
71
|
+
- ⏩ **Trigger Downstream DAGs**: `trigger_dag_id` for ML pipeline integration
|
|
72
|
+
- 🎯 **Pattern Filtering**: `model_pattern` for regex-based model selection
|
|
73
|
+
|
|
66
74
|
## ⚠️ Important: Gateway vs Environment
|
|
67
75
|
|
|
68
76
|
**SQLMesh uses "gateways" to switch between environments, NOT an "environment" parameter.**
|
|
@@ -16,6 +16,14 @@ Transform your SQLMesh models into production-ready Airflow DAGs with **full dat
|
|
|
16
16
|
- 🛠️ **Dual Mode**: Dynamic (auto-discovery, default) or Static (full control)
|
|
17
17
|
- 🚫 **No Vendor Lock-in**: Open source, no cloud dependencies
|
|
18
18
|
|
|
19
|
+
### 🏢 Enterprise Features (NEW in v0.8.0)
|
|
20
|
+
|
|
21
|
+
- 🔔 **Callbacks**: `on_failure_callback`, `on_success_callback`, `sla_miss_callback` for alerting
|
|
22
|
+
- 🏷️ **Tag-Based Filtering**: `include_tags`, `exclude_tags` for Data Mesh team-specific DAGs
|
|
23
|
+
- 🎱 **Pool Configuration**: `pool`, `pool_slots` for resource management
|
|
24
|
+
- ⏩ **Trigger Downstream DAGs**: `trigger_dag_id` for ML pipeline integration
|
|
25
|
+
- 🎯 **Pattern Filtering**: `model_pattern` for regex-based model selection
|
|
26
|
+
|
|
19
27
|
## ⚠️ Important: Gateway vs Environment
|
|
20
28
|
|
|
21
29
|
**SQLMesh uses "gateways" to switch between environments, NOT an "environment" parameter.**
|
|
@@ -20,7 +20,7 @@ if requirements_file.exists():
|
|
|
20
20
|
|
|
21
21
|
setup(
|
|
22
22
|
name="sqlmesh-dag-generator",
|
|
23
|
-
version="0.7.
|
|
23
|
+
version="0.7.3",
|
|
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",
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/airflow_utils.py
RENAMED
|
@@ -268,7 +268,8 @@ def _build_config_from_connection(conn: Any, conn_type: str) -> Dict[str, Any]:
|
|
|
268
268
|
# Add connection details based on type
|
|
269
269
|
if conn_type == "redshift":
|
|
270
270
|
# Redshift uses 2-part naming (schema.table), not 3-part (catalog.schema.table)
|
|
271
|
-
#
|
|
271
|
+
# Note: default_catalog is NOT a valid Redshift connection config field
|
|
272
|
+
# It should be set at the SQLMesh config level, not connection level
|
|
272
273
|
database = conn.schema or "dev"
|
|
273
274
|
config.update({
|
|
274
275
|
"host": conn.host,
|
|
@@ -276,7 +277,6 @@ def _build_config_from_connection(conn: Any, conn_type: str) -> Dict[str, Any]:
|
|
|
276
277
|
"user": conn.login,
|
|
277
278
|
"password": conn.password,
|
|
278
279
|
"database": database,
|
|
279
|
-
"default_catalog": database, # Tells SQLMesh to omit this from SQL generation
|
|
280
280
|
})
|
|
281
281
|
elif conn_type in ["postgres", "postgresql"]:
|
|
282
282
|
config.update({
|
|
@@ -94,6 +94,11 @@ class AirflowConfig:
|
|
|
94
94
|
max_active_runs: int = 1
|
|
95
95
|
description: Optional[str] = None
|
|
96
96
|
env_vars: Dict[str, str] = field(default_factory=dict) # Environment variables for tasks
|
|
97
|
+
# Callback configuration - pass callable names (will be imported in generated DAG)
|
|
98
|
+
on_failure_callback: Optional[str] = None # e.g., "my_module.slack_alert"
|
|
99
|
+
on_success_callback: Optional[str] = None # e.g., "my_module.log_success"
|
|
100
|
+
sla_miss_callback: Optional[str] = None # e.g., "my_module.sla_alert"
|
|
101
|
+
sla: Optional[int] = None # SLA in seconds for all tasks
|
|
97
102
|
|
|
98
103
|
|
|
99
104
|
@dataclass
|
|
@@ -116,6 +121,15 @@ class GenerationConfig:
|
|
|
116
121
|
auto_replan_on_change: bool = False # Automatically replan if models change
|
|
117
122
|
skip_audits: bool = False # Skip audit checks during execution
|
|
118
123
|
enable_health_check: bool = False # Add a pre-flight health check task
|
|
124
|
+
# Tag-based filtering - only include models with any of these tags
|
|
125
|
+
include_tags: Optional[List[str]] = None # e.g., ["finance", "core"]
|
|
126
|
+
exclude_tags: Optional[List[str]] = None # e.g., ["deprecated", "test"]
|
|
127
|
+
# Resource management
|
|
128
|
+
pool: Optional[str] = None # Airflow pool for all tasks
|
|
129
|
+
pool_slots: int = 1 # Number of pool slots per task
|
|
130
|
+
# Trigger downstream DAG after completion
|
|
131
|
+
trigger_dag_id: Optional[str] = None # DAG to trigger on success
|
|
132
|
+
trigger_dag_conf: Optional[Dict[str, Any]] = None # Conf to pass to triggered DAG
|
|
119
133
|
|
|
120
134
|
|
|
121
135
|
@dataclass
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/dag_builder.py
RENAMED
|
@@ -92,6 +92,14 @@ class AirflowDAGBuilder:
|
|
|
92
92
|
"from kubernetes.client.models import V1EnvVar",
|
|
93
93
|
])
|
|
94
94
|
|
|
95
|
+
# Add TriggerDagRunOperator if trigger_dag_id is configured
|
|
96
|
+
if self.config.generation.trigger_dag_id:
|
|
97
|
+
imports.append("from airflow.operators.trigger_dagrun import TriggerDagRunOperator")
|
|
98
|
+
|
|
99
|
+
# Add EmptyOperator for source tables
|
|
100
|
+
if self.config.generation.include_source_tables:
|
|
101
|
+
imports.append("from airflow.operators.empty import EmptyOperator")
|
|
102
|
+
|
|
95
103
|
imports.extend([
|
|
96
104
|
"",
|
|
97
105
|
"from sqlmesh import Context",
|
|
@@ -99,6 +107,20 @@ class AirflowDAGBuilder:
|
|
|
99
107
|
"logger = logging.getLogger(__name__)",
|
|
100
108
|
])
|
|
101
109
|
|
|
110
|
+
# Add callback imports if configured
|
|
111
|
+
if self.config.airflow.on_failure_callback:
|
|
112
|
+
module_path = self.config.airflow.on_failure_callback.rsplit('.', 1)
|
|
113
|
+
if len(module_path) == 2:
|
|
114
|
+
imports.insert(4, f"from {module_path[0]} import {module_path[1]}")
|
|
115
|
+
if self.config.airflow.on_success_callback:
|
|
116
|
+
module_path = self.config.airflow.on_success_callback.rsplit('.', 1)
|
|
117
|
+
if len(module_path) == 2:
|
|
118
|
+
imports.insert(4, f"from {module_path[0]} import {module_path[1]}")
|
|
119
|
+
if self.config.airflow.sla_miss_callback:
|
|
120
|
+
module_path = self.config.airflow.sla_miss_callback.rsplit('.', 1)
|
|
121
|
+
if len(module_path) == 2:
|
|
122
|
+
imports.insert(4, f"from {module_path[0]} import {module_path[1]}")
|
|
123
|
+
|
|
102
124
|
return "\n".join(imports)
|
|
103
125
|
|
|
104
126
|
def _build_config(self) -> str:
|
|
@@ -128,6 +150,11 @@ class AirflowDAGBuilder:
|
|
|
128
150
|
else:
|
|
129
151
|
start_date_str = "datetime.now() - timedelta(days=1)" # Default to yesterday
|
|
130
152
|
|
|
153
|
+
# Build optional DAG arguments
|
|
154
|
+
optional_args = ""
|
|
155
|
+
if self.config.airflow.sla_miss_callback:
|
|
156
|
+
optional_args += f"\n sla_miss_callback={self.config.airflow.sla_miss_callback},"
|
|
157
|
+
|
|
131
158
|
return f"""# DAG Definition
|
|
132
159
|
dag = DAG(
|
|
133
160
|
dag_id="{self.config.airflow.dag_id}",
|
|
@@ -137,7 +164,7 @@ dag = DAG(
|
|
|
137
164
|
start_date={start_date_str},
|
|
138
165
|
catchup={self.config.airflow.catchup},
|
|
139
166
|
max_active_runs={self.config.airflow.max_active_runs},
|
|
140
|
-
tags={tags},
|
|
167
|
+
tags={tags},{optional_args}
|
|
141
168
|
)"""
|
|
142
169
|
|
|
143
170
|
def _format_default_args(self) -> str:
|
|
@@ -237,11 +264,30 @@ dag = DAG(
|
|
|
237
264
|
logger.info(f"Model {model_name_escaped} completed successfully")
|
|
238
265
|
return result'''
|
|
239
266
|
|
|
240
|
-
# Build the operator
|
|
267
|
+
# Build the operator with additional configuration
|
|
268
|
+
operator_args = [
|
|
269
|
+
f'task_id="{task_id}"',
|
|
270
|
+
f'python_callable={function_name}',
|
|
271
|
+
'dag=dag',
|
|
272
|
+
]
|
|
273
|
+
|
|
274
|
+
# Add pool if configured
|
|
275
|
+
if self.config.generation.pool:
|
|
276
|
+
operator_args.append(f'pool="{self.config.generation.pool}"')
|
|
277
|
+
operator_args.append(f'pool_slots={self.config.generation.pool_slots}')
|
|
278
|
+
|
|
279
|
+
# Add SLA if configured
|
|
280
|
+
if self.config.airflow.sla:
|
|
281
|
+
operator_args.append(f'sla=timedelta(seconds={self.config.airflow.sla})')
|
|
282
|
+
|
|
283
|
+
# Add callbacks if configured
|
|
284
|
+
if self.config.airflow.on_failure_callback:
|
|
285
|
+
operator_args.append(f'on_failure_callback={self.config.airflow.on_failure_callback}')
|
|
286
|
+
if self.config.airflow.on_success_callback:
|
|
287
|
+
operator_args.append(f'on_success_callback={self.config.airflow.on_success_callback}')
|
|
288
|
+
|
|
241
289
|
operator_def = f'''{task_id} = PythonOperator(
|
|
242
|
-
|
|
243
|
-
python_callable={function_name},
|
|
244
|
-
dag=dag,
|
|
290
|
+
{(","+chr(10)+" ").join(operator_args)},
|
|
245
291
|
)'''
|
|
246
292
|
|
|
247
293
|
return f"{func_def}\n\n{operator_def}"
|
|
@@ -327,6 +373,26 @@ dag = DAG(
|
|
|
327
373
|
if len(dep_lines) == 1:
|
|
328
374
|
dep_lines.append("# No dependencies")
|
|
329
375
|
|
|
376
|
+
# Add trigger downstream DAG if configured
|
|
377
|
+
if self.config.generation.trigger_dag_id:
|
|
378
|
+
trigger_conf = repr(self.config.generation.trigger_dag_conf or {})
|
|
379
|
+
dep_lines.append("")
|
|
380
|
+
dep_lines.append("# Trigger downstream DAG on completion")
|
|
381
|
+
dep_lines.append(f'''trigger_downstream = TriggerDagRunOperator(
|
|
382
|
+
task_id="trigger_{self.config.generation.trigger_dag_id}",
|
|
383
|
+
trigger_dag_id="{self.config.generation.trigger_dag_id}",
|
|
384
|
+
conf={trigger_conf},
|
|
385
|
+
dag=dag,
|
|
386
|
+
)''')
|
|
387
|
+
# Find leaf tasks and set them as upstream
|
|
388
|
+
leaf_tasks = [
|
|
389
|
+
info.get_task_id()
|
|
390
|
+
for name, info in self.dag_structure.models.items()
|
|
391
|
+
if name not in {dep for m in self.dag_structure.models.values() for dep in m.dependencies}
|
|
392
|
+
]
|
|
393
|
+
if leaf_tasks:
|
|
394
|
+
dep_lines.append(f"[{', '.join(leaf_tasks)}] >> trigger_downstream")
|
|
395
|
+
|
|
330
396
|
return "\n".join(dep_lines)
|
|
331
397
|
|
|
332
398
|
# ========================================================================
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/generator.py
RENAMED
|
@@ -158,6 +158,10 @@ class SQLMeshDAGGenerator:
|
|
|
158
158
|
tags=kwargs.get("tags", ["sqlmesh"]),
|
|
159
159
|
catchup=kwargs.get("catchup", False),
|
|
160
160
|
max_active_runs=kwargs.get("max_active_runs", 1),
|
|
161
|
+
on_failure_callback=kwargs.get("on_failure_callback"),
|
|
162
|
+
on_success_callback=kwargs.get("on_success_callback"),
|
|
163
|
+
sla_miss_callback=kwargs.get("sla_miss_callback"),
|
|
164
|
+
sla=kwargs.get("sla"),
|
|
161
165
|
)
|
|
162
166
|
|
|
163
167
|
generation_config = GenerationConfig(
|
|
@@ -167,11 +171,18 @@ class SQLMeshDAGGenerator:
|
|
|
167
171
|
parallel_tasks=kwargs.get("parallel_tasks", True),
|
|
168
172
|
include_models=kwargs.get("include_models"),
|
|
169
173
|
exclude_models=kwargs.get("exclude_models"),
|
|
174
|
+
model_pattern=kwargs.get("model_pattern"),
|
|
170
175
|
include_source_tables=kwargs.get("include_source_tables", True), # Default: enabled
|
|
171
176
|
return_value=kwargs.get("return_value", True),
|
|
172
177
|
auto_replan_on_change=kwargs.get("auto_replan_on_change", False),
|
|
173
178
|
skip_audits=kwargs.get("skip_audits", False),
|
|
174
179
|
enable_health_check=kwargs.get("enable_health_check", False),
|
|
180
|
+
include_tags=kwargs.get("include_tags"),
|
|
181
|
+
exclude_tags=kwargs.get("exclude_tags"),
|
|
182
|
+
pool=kwargs.get("pool"),
|
|
183
|
+
pool_slots=kwargs.get("pool_slots", 1),
|
|
184
|
+
trigger_dag_id=kwargs.get("trigger_dag_id"),
|
|
185
|
+
trigger_dag_conf=kwargs.get("trigger_dag_conf"),
|
|
175
186
|
)
|
|
176
187
|
|
|
177
188
|
self.config = DAGGeneratorConfig(
|
|
@@ -186,6 +197,33 @@ class SQLMeshDAGGenerator:
|
|
|
186
197
|
self.merged_config = None # Store merged config for runtime task execution
|
|
187
198
|
self.runtime_gateway = None # Store gateway name for runtime task execution
|
|
188
199
|
|
|
200
|
+
# Validate configuration and show helpful warnings
|
|
201
|
+
self._validate_config()
|
|
202
|
+
|
|
203
|
+
def _validate_config(self):
|
|
204
|
+
"""Validate configuration and show helpful warnings."""
|
|
205
|
+
# Warn if using Redshift without default_catalog
|
|
206
|
+
if self.config.sqlmesh.connection_config:
|
|
207
|
+
conn_type = self.config.sqlmesh.connection_config.get("type", "").lower()
|
|
208
|
+
if conn_type == "redshift" and not self.config.sqlmesh.default_catalog:
|
|
209
|
+
import warnings
|
|
210
|
+
warnings.warn(
|
|
211
|
+
f"\n{'='*80}\n"
|
|
212
|
+
f"⚠️ Redshift detected without default_catalog setting!\n\n"
|
|
213
|
+
f"Redshift uses 2-part naming (schema.table), not 3-part (catalog.schema.table).\n"
|
|
214
|
+
f"Without default_catalog, SQLMesh may generate invalid SQL like:\n"
|
|
215
|
+
f" FROM \"your_database\".\"schema\".\"table\" -- ERROR!\n\n"
|
|
216
|
+
f"🔧 SOLUTION: Add default_catalog to your SQLMeshDAGGenerator:\n\n"
|
|
217
|
+
f" generator = SQLMeshDAGGenerator(\n"
|
|
218
|
+
f" ...\n"
|
|
219
|
+
f" default_catalog='your_database_name', # Add this!\n"
|
|
220
|
+
f" )\n\n"
|
|
221
|
+
f"📚 See docs/REDSHIFT_3PART_NAMING_FIX.md for more information.\n"
|
|
222
|
+
f"{'='*80}\n",
|
|
223
|
+
UserWarning,
|
|
224
|
+
stacklevel=3
|
|
225
|
+
)
|
|
226
|
+
|
|
189
227
|
def load_sqlmesh_context(self) -> Context:
|
|
190
228
|
"""
|
|
191
229
|
Load the SQLMesh context from the project path.
|
|
@@ -361,8 +399,8 @@ class SQLMeshDAGGenerator:
|
|
|
361
399
|
logger.warning("Could not find models in context")
|
|
362
400
|
|
|
363
401
|
for model_name, model in sqlmesh_models.items():
|
|
364
|
-
# Filter models based on include/exclude patterns
|
|
365
|
-
if not self._should_include_model(model_name):
|
|
402
|
+
# Filter models based on include/exclude patterns and tags
|
|
403
|
+
if not self._should_include_model(model_name, model):
|
|
366
404
|
continue
|
|
367
405
|
|
|
368
406
|
model_info = self._extract_model_info(model_name, model)
|
|
@@ -386,8 +424,10 @@ class SQLMeshDAGGenerator:
|
|
|
386
424
|
|
|
387
425
|
return models
|
|
388
426
|
|
|
389
|
-
def _should_include_model(self, model_name: str) -> bool:
|
|
427
|
+
def _should_include_model(self, model_name: str, model: Model = None) -> bool:
|
|
390
428
|
"""Check if a model should be included based on filters"""
|
|
429
|
+
import re
|
|
430
|
+
|
|
391
431
|
# Check include patterns
|
|
392
432
|
if self.config.generation.include_models:
|
|
393
433
|
if model_name not in self.config.generation.include_models:
|
|
@@ -398,6 +438,31 @@ class SQLMeshDAGGenerator:
|
|
|
398
438
|
if model_name in self.config.generation.exclude_models:
|
|
399
439
|
return False
|
|
400
440
|
|
|
441
|
+
# Check model_pattern (regex)
|
|
442
|
+
if self.config.generation.model_pattern:
|
|
443
|
+
pattern = self.config.generation.model_pattern
|
|
444
|
+
if not re.match(pattern, model_name):
|
|
445
|
+
logger.debug(f"Model {model_name} excluded by pattern: {pattern}")
|
|
446
|
+
return False
|
|
447
|
+
|
|
448
|
+
# Tag-based filtering (requires model object)
|
|
449
|
+
if model is not None:
|
|
450
|
+
model_tags = set(getattr(model, 'tags', []) or [])
|
|
451
|
+
|
|
452
|
+
# Check include_tags - model must have at least one of these tags
|
|
453
|
+
if self.config.generation.include_tags:
|
|
454
|
+
include_tags = set(self.config.generation.include_tags)
|
|
455
|
+
if not model_tags.intersection(include_tags):
|
|
456
|
+
logger.debug(f"Model {model_name} excluded: no matching include_tags")
|
|
457
|
+
return False
|
|
458
|
+
|
|
459
|
+
# Check exclude_tags - model must not have any of these tags
|
|
460
|
+
if self.config.generation.exclude_tags:
|
|
461
|
+
exclude_tags = set(self.config.generation.exclude_tags)
|
|
462
|
+
if model_tags.intersection(exclude_tags):
|
|
463
|
+
logger.debug(f"Model {model_name} excluded: has exclude_tags")
|
|
464
|
+
return False
|
|
465
|
+
|
|
401
466
|
return True
|
|
402
467
|
|
|
403
468
|
def _extract_model_info(self, model_name: str, model: Model) -> SQLMeshModelInfo:
|
|
@@ -865,8 +930,10 @@ class SQLMeshDAGGenerator:
|
|
|
865
930
|
"model": m_fqn,
|
|
866
931
|
}
|
|
867
932
|
except Exception as e:
|
|
933
|
+
error_msg = str(e)
|
|
934
|
+
|
|
868
935
|
# Check for "Environment not found" error
|
|
869
|
-
if "Environment" in
|
|
936
|
+
if "Environment" in error_msg and "was not found" in error_msg:
|
|
870
937
|
env_name = self.config.sqlmesh.environment
|
|
871
938
|
raise RuntimeError(
|
|
872
939
|
f"SQLMesh environment '{env_name}' was not found.\n\n"
|
|
@@ -886,6 +953,28 @@ class SQLMeshDAGGenerator:
|
|
|
886
953
|
f" See docs/SQLMESH_ENVIRONMENTS.md for complete explanation.\n\n"
|
|
887
954
|
f"Original error: {e}"
|
|
888
955
|
) from e
|
|
956
|
+
|
|
957
|
+
# Check for Redshift catalog/3-part naming error
|
|
958
|
+
elif "does not exist" in error_msg and (
|
|
959
|
+
"redshift" in error_msg.lower() or
|
|
960
|
+
"3D000" in error_msg or # Redshift error code for invalid catalog
|
|
961
|
+
("." in error_msg and error_msg.count('"') >= 6) # 3-part naming pattern
|
|
962
|
+
):
|
|
963
|
+
raise RuntimeError(
|
|
964
|
+
f"SQLMesh catalog error (likely 3-part naming issue):\n"
|
|
965
|
+
f"{error_msg}\n\n"
|
|
966
|
+
f"🔧 SOLUTION: For Redshift (2-part naming), set default_catalog:\n\n"
|
|
967
|
+
f" generator = SQLMeshDAGGenerator(\n"
|
|
968
|
+
f" sqlmesh_project_path='/path/to/project',\n"
|
|
969
|
+
f" connection='REDSHIFT',\n"
|
|
970
|
+
f" default_catalog='your_database_name', # ✅ Add this!\n"
|
|
971
|
+
f" )\n\n"
|
|
972
|
+
f"📚 Why? Redshift uses schema.table (2-part), not catalog.schema.table (3-part).\n"
|
|
973
|
+
f" Setting default_catalog tells SQLMesh to omit the catalog prefix.\n\n"
|
|
974
|
+
f" See docs/REDSHIFT_3PART_NAMING_FIX.md for more information.\n\n"
|
|
975
|
+
f"Original error: {e}"
|
|
976
|
+
) from e
|
|
977
|
+
|
|
889
978
|
else:
|
|
890
979
|
# Re-raise other errors as-is
|
|
891
980
|
raise
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3/sqlmesh_dag_generator.egg-info}/PKG-INFO
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: sqlmesh-dag-generator
|
|
3
|
-
Version: 0.7.
|
|
3
|
+
Version: 0.7.3
|
|
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
|
|
@@ -63,6 +63,14 @@ Transform your SQLMesh models into production-ready Airflow DAGs with **full dat
|
|
|
63
63
|
- 🛠️ **Dual Mode**: Dynamic (auto-discovery, default) or Static (full control)
|
|
64
64
|
- 🚫 **No Vendor Lock-in**: Open source, no cloud dependencies
|
|
65
65
|
|
|
66
|
+
### 🏢 Enterprise Features (NEW in v0.8.0)
|
|
67
|
+
|
|
68
|
+
- 🔔 **Callbacks**: `on_failure_callback`, `on_success_callback`, `sla_miss_callback` for alerting
|
|
69
|
+
- 🏷️ **Tag-Based Filtering**: `include_tags`, `exclude_tags` for Data Mesh team-specific DAGs
|
|
70
|
+
- 🎱 **Pool Configuration**: `pool`, `pool_slots` for resource management
|
|
71
|
+
- ⏩ **Trigger Downstream DAGs**: `trigger_dag_id` for ML pipeline integration
|
|
72
|
+
- 🎯 **Pattern Filtering**: `model_pattern` for regex-based model selection
|
|
73
|
+
|
|
66
74
|
## ⚠️ Important: Gateway vs Environment
|
|
67
75
|
|
|
68
76
|
**SQLMesh uses "gateways" to switch between environments, NOT an "environment" parameter.**
|
|
@@ -21,8 +21,12 @@ class TestConnectionTypeMapping:
|
|
|
21
21
|
def test_redshift_mapping(self):
|
|
22
22
|
assert _map_airflow_conn_type_to_sqlmesh("redshift") == "redshift"
|
|
23
23
|
|
|
24
|
-
def
|
|
25
|
-
"""Test that Redshift connections include default_catalog
|
|
24
|
+
def test_redshift_no_default_catalog_in_connection(self):
|
|
25
|
+
"""Test that Redshift connections do NOT include default_catalog.
|
|
26
|
+
|
|
27
|
+
default_catalog is a SQLMesh config-level setting, NOT a connection config field.
|
|
28
|
+
For Redshift 2-part naming, set default_catalog on SQLMeshDAGGenerator, not here.
|
|
29
|
+
"""
|
|
26
30
|
from sqlmesh_dag_generator.airflow_utils import _build_config_from_connection
|
|
27
31
|
|
|
28
32
|
mock_conn = Mock()
|
|
@@ -36,11 +40,11 @@ class TestConnectionTypeMapping:
|
|
|
36
40
|
|
|
37
41
|
result = _build_config_from_connection(mock_conn, "redshift")
|
|
38
42
|
|
|
39
|
-
# Redshift should have default_catalog
|
|
40
|
-
#
|
|
43
|
+
# Redshift connection should NOT have default_catalog
|
|
44
|
+
# default_catalog is set at SQLMesh config level, not in connection config
|
|
41
45
|
assert result["type"] == "redshift"
|
|
42
46
|
assert result["database"] == "stg_carrier_quality_redshift_db"
|
|
43
|
-
assert
|
|
47
|
+
assert "default_catalog" not in result # Must NOT be in connection config
|
|
44
48
|
assert result["port"] == 5439
|
|
45
49
|
|
|
46
50
|
def test_mysql_mapping(self):
|
|
@@ -0,0 +1,240 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Tests for enterprise features:
|
|
3
|
+
- Callback support (on_failure, on_success, sla_miss)
|
|
4
|
+
- Tag-based filtering
|
|
5
|
+
- Pool configuration
|
|
6
|
+
- Pattern-based model selection
|
|
7
|
+
- Trigger downstream DAG
|
|
8
|
+
"""
|
|
9
|
+
import pytest
|
|
10
|
+
from sqlmesh_dag_generator import SQLMeshDAGGenerator, DAGGeneratorConfig
|
|
11
|
+
from sqlmesh_dag_generator.config import SQLMeshConfig, AirflowConfig, GenerationConfig
|
|
12
|
+
from sqlmesh_dag_generator.models import SQLMeshModelInfo
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class TestCallbackSupport:
|
|
16
|
+
"""Test callback configuration"""
|
|
17
|
+
|
|
18
|
+
def test_on_failure_callback_in_config(self):
|
|
19
|
+
"""Test that on_failure_callback can be set"""
|
|
20
|
+
config = AirflowConfig(
|
|
21
|
+
dag_id="test_dag",
|
|
22
|
+
on_failure_callback="my_module.slack_alert",
|
|
23
|
+
)
|
|
24
|
+
assert config.on_failure_callback == "my_module.slack_alert"
|
|
25
|
+
|
|
26
|
+
def test_on_success_callback_in_config(self):
|
|
27
|
+
"""Test that on_success_callback can be set"""
|
|
28
|
+
config = AirflowConfig(
|
|
29
|
+
dag_id="test_dag",
|
|
30
|
+
on_success_callback="my_module.log_success",
|
|
31
|
+
)
|
|
32
|
+
assert config.on_success_callback == "my_module.log_success"
|
|
33
|
+
|
|
34
|
+
def test_sla_miss_callback_in_config(self):
|
|
35
|
+
"""Test that sla_miss_callback can be set"""
|
|
36
|
+
config = AirflowConfig(
|
|
37
|
+
dag_id="test_dag",
|
|
38
|
+
sla_miss_callback="my_module.sla_alert",
|
|
39
|
+
)
|
|
40
|
+
assert config.sla_miss_callback == "my_module.sla_alert"
|
|
41
|
+
|
|
42
|
+
def test_sla_seconds_in_config(self):
|
|
43
|
+
"""Test that SLA can be set in seconds"""
|
|
44
|
+
config = AirflowConfig(
|
|
45
|
+
dag_id="test_dag",
|
|
46
|
+
sla=3600, # 1 hour
|
|
47
|
+
)
|
|
48
|
+
assert config.sla == 3600
|
|
49
|
+
|
|
50
|
+
def test_generator_with_callbacks(self, tmp_path):
|
|
51
|
+
"""Test generator accepts callback parameters"""
|
|
52
|
+
project_path = tmp_path / "project"
|
|
53
|
+
project_path.mkdir()
|
|
54
|
+
(project_path / "config.yaml").write_text("gateways:\n default:\n connection:\n type: duckdb")
|
|
55
|
+
|
|
56
|
+
generator = SQLMeshDAGGenerator(
|
|
57
|
+
sqlmesh_project_path=str(project_path),
|
|
58
|
+
dag_id="test_callback_dag",
|
|
59
|
+
on_failure_callback="alerts.slack_notify",
|
|
60
|
+
on_success_callback="alerts.log_success",
|
|
61
|
+
sla=7200,
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
assert generator.config.airflow.on_failure_callback == "alerts.slack_notify"
|
|
65
|
+
assert generator.config.airflow.on_success_callback == "alerts.log_success"
|
|
66
|
+
assert generator.config.airflow.sla == 7200
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
class TestTagBasedFiltering:
|
|
70
|
+
"""Test tag-based model filtering"""
|
|
71
|
+
|
|
72
|
+
def test_include_tags_in_config(self):
|
|
73
|
+
"""Test that include_tags can be set"""
|
|
74
|
+
config = GenerationConfig(
|
|
75
|
+
include_tags=["finance", "core"],
|
|
76
|
+
)
|
|
77
|
+
assert config.include_tags == ["finance", "core"]
|
|
78
|
+
|
|
79
|
+
def test_exclude_tags_in_config(self):
|
|
80
|
+
"""Test that exclude_tags can be set"""
|
|
81
|
+
config = GenerationConfig(
|
|
82
|
+
exclude_tags=["deprecated", "test"],
|
|
83
|
+
)
|
|
84
|
+
assert config.exclude_tags == ["deprecated", "test"]
|
|
85
|
+
|
|
86
|
+
def test_generator_with_tags(self, tmp_path):
|
|
87
|
+
"""Test generator accepts tag filter parameters"""
|
|
88
|
+
project_path = tmp_path / "project"
|
|
89
|
+
project_path.mkdir()
|
|
90
|
+
(project_path / "config.yaml").write_text("gateways:\n default:\n connection:\n type: duckdb")
|
|
91
|
+
|
|
92
|
+
generator = SQLMeshDAGGenerator(
|
|
93
|
+
sqlmesh_project_path=str(project_path),
|
|
94
|
+
dag_id="test_tag_dag",
|
|
95
|
+
include_tags=["production", "tier1"],
|
|
96
|
+
exclude_tags=["experimental"],
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
assert generator.config.generation.include_tags == ["production", "tier1"]
|
|
100
|
+
assert generator.config.generation.exclude_tags == ["experimental"]
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
class TestPoolConfiguration:
|
|
104
|
+
"""Test Airflow pool configuration"""
|
|
105
|
+
|
|
106
|
+
def test_pool_in_config(self):
|
|
107
|
+
"""Test that pool can be set"""
|
|
108
|
+
config = GenerationConfig(
|
|
109
|
+
pool="sqlmesh_pool",
|
|
110
|
+
pool_slots=2,
|
|
111
|
+
)
|
|
112
|
+
assert config.pool == "sqlmesh_pool"
|
|
113
|
+
assert config.pool_slots == 2
|
|
114
|
+
|
|
115
|
+
def test_pool_slots_default(self):
|
|
116
|
+
"""Test that pool_slots defaults to 1"""
|
|
117
|
+
config = GenerationConfig()
|
|
118
|
+
assert config.pool_slots == 1
|
|
119
|
+
|
|
120
|
+
def test_generator_with_pool(self, tmp_path):
|
|
121
|
+
"""Test generator accepts pool parameters"""
|
|
122
|
+
project_path = tmp_path / "project"
|
|
123
|
+
project_path.mkdir()
|
|
124
|
+
(project_path / "config.yaml").write_text("gateways:\n default:\n connection:\n type: duckdb")
|
|
125
|
+
|
|
126
|
+
generator = SQLMeshDAGGenerator(
|
|
127
|
+
sqlmesh_project_path=str(project_path),
|
|
128
|
+
dag_id="test_pool_dag",
|
|
129
|
+
pool="db_pool",
|
|
130
|
+
pool_slots=3,
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
assert generator.config.generation.pool == "db_pool"
|
|
134
|
+
assert generator.config.generation.pool_slots == 3
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
class TestTriggerDownstreamDAG:
|
|
138
|
+
"""Test trigger downstream DAG configuration"""
|
|
139
|
+
|
|
140
|
+
def test_trigger_dag_id_in_config(self):
|
|
141
|
+
"""Test that trigger_dag_id can be set"""
|
|
142
|
+
config = GenerationConfig(
|
|
143
|
+
trigger_dag_id="ml_training_dag",
|
|
144
|
+
)
|
|
145
|
+
assert config.trigger_dag_id == "ml_training_dag"
|
|
146
|
+
|
|
147
|
+
def test_trigger_dag_conf_in_config(self):
|
|
148
|
+
"""Test that trigger_dag_conf can be set"""
|
|
149
|
+
config = GenerationConfig(
|
|
150
|
+
trigger_dag_id="ml_training_dag",
|
|
151
|
+
trigger_dag_conf={"model_version": "v2", "retrain": True},
|
|
152
|
+
)
|
|
153
|
+
assert config.trigger_dag_conf == {"model_version": "v2", "retrain": True}
|
|
154
|
+
|
|
155
|
+
def test_generator_with_trigger_dag(self, tmp_path):
|
|
156
|
+
"""Test generator accepts trigger DAG parameters"""
|
|
157
|
+
project_path = tmp_path / "project"
|
|
158
|
+
project_path.mkdir()
|
|
159
|
+
(project_path / "config.yaml").write_text("gateways:\n default:\n connection:\n type: duckdb")
|
|
160
|
+
|
|
161
|
+
generator = SQLMeshDAGGenerator(
|
|
162
|
+
sqlmesh_project_path=str(project_path),
|
|
163
|
+
dag_id="test_trigger_dag",
|
|
164
|
+
trigger_dag_id="downstream_dag",
|
|
165
|
+
trigger_dag_conf={"key": "value"},
|
|
166
|
+
)
|
|
167
|
+
|
|
168
|
+
assert generator.config.generation.trigger_dag_id == "downstream_dag"
|
|
169
|
+
assert generator.config.generation.trigger_dag_conf == {"key": "value"}
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
class TestModelPatternFiltering:
|
|
173
|
+
"""Test model pattern (regex) filtering"""
|
|
174
|
+
|
|
175
|
+
def test_model_pattern_in_config(self):
|
|
176
|
+
"""Test that model_pattern can be set"""
|
|
177
|
+
config = GenerationConfig(
|
|
178
|
+
model_pattern=r"^analytics\..*",
|
|
179
|
+
)
|
|
180
|
+
assert config.model_pattern == r"^analytics\..*"
|
|
181
|
+
|
|
182
|
+
def test_generator_with_pattern(self, tmp_path):
|
|
183
|
+
"""Test generator accepts model_pattern parameter"""
|
|
184
|
+
project_path = tmp_path / "project"
|
|
185
|
+
project_path.mkdir()
|
|
186
|
+
(project_path / "config.yaml").write_text("gateways:\n default:\n connection:\n type: duckdb")
|
|
187
|
+
|
|
188
|
+
generator = SQLMeshDAGGenerator(
|
|
189
|
+
sqlmesh_project_path=str(project_path),
|
|
190
|
+
dag_id="test_pattern_dag",
|
|
191
|
+
model_pattern=r".*_fact$",
|
|
192
|
+
)
|
|
193
|
+
|
|
194
|
+
assert generator.config.generation.model_pattern == r".*_fact$"
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
class TestCombinedEnterpriseFeatures:
|
|
198
|
+
"""Test multiple enterprise features together"""
|
|
199
|
+
|
|
200
|
+
def test_full_enterprise_config(self, tmp_path):
|
|
201
|
+
"""Test full enterprise configuration with all new features"""
|
|
202
|
+
project_path = tmp_path / "project"
|
|
203
|
+
project_path.mkdir()
|
|
204
|
+
(project_path / "config.yaml").write_text("gateways:\n default:\n connection:\n type: duckdb")
|
|
205
|
+
|
|
206
|
+
generator = SQLMeshDAGGenerator(
|
|
207
|
+
sqlmesh_project_path=str(project_path),
|
|
208
|
+
dag_id="enterprise_dag",
|
|
209
|
+
# Callbacks
|
|
210
|
+
on_failure_callback="alerts.pagerduty_alert",
|
|
211
|
+
on_success_callback="metrics.record_success",
|
|
212
|
+
sla_miss_callback="alerts.sla_breach",
|
|
213
|
+
sla=1800, # 30 minutes
|
|
214
|
+
# Tag filtering
|
|
215
|
+
include_tags=["production"],
|
|
216
|
+
exclude_tags=["deprecated"],
|
|
217
|
+
# Resource management
|
|
218
|
+
pool="prod_pool",
|
|
219
|
+
pool_slots=2,
|
|
220
|
+
# Pattern filtering
|
|
221
|
+
model_pattern=r"^prod\..*",
|
|
222
|
+
# Trigger downstream
|
|
223
|
+
trigger_dag_id="feature_store_update",
|
|
224
|
+
trigger_dag_conf={"source": "sqlmesh"},
|
|
225
|
+
)
|
|
226
|
+
|
|
227
|
+
# Verify all config values
|
|
228
|
+
assert generator.config.airflow.on_failure_callback == "alerts.pagerduty_alert"
|
|
229
|
+
assert generator.config.airflow.on_success_callback == "metrics.record_success"
|
|
230
|
+
assert generator.config.airflow.sla_miss_callback == "alerts.sla_breach"
|
|
231
|
+
assert generator.config.airflow.sla == 1800
|
|
232
|
+
assert generator.config.generation.include_tags == ["production"]
|
|
233
|
+
assert generator.config.generation.exclude_tags == ["deprecated"]
|
|
234
|
+
assert generator.config.generation.pool == "prod_pool"
|
|
235
|
+
assert generator.config.generation.pool_slots == 2
|
|
236
|
+
assert generator.config.generation.model_pattern == r"^prod\..*"
|
|
237
|
+
assert generator.config.generation.trigger_dag_id == "feature_store_update"
|
|
238
|
+
assert generator.config.generation.trigger_dag_conf == {"source": "sqlmesh"}
|
|
239
|
+
|
|
240
|
+
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/security.py
RENAMED
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/validation.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_bug_fix_user_scenario.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_environment_default.py
RENAMED
|
File without changes
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_environment_handling.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_runtime_connections.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|