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.
Files changed (36) hide show
  1. {sqlmesh_dag_generator-0.7.2/sqlmesh_dag_generator.egg-info → sqlmesh_dag_generator-0.7.3}/PKG-INFO +9 -1
  2. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/README.md +8 -0
  3. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/setup.py +1 -1
  4. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/__init__.py +1 -1
  5. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/airflow_utils.py +2 -2
  6. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/config.py +14 -0
  7. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/dag_builder.py +71 -5
  8. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/generator.py +93 -4
  9. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3/sqlmesh_dag_generator.egg-info}/PKG-INFO +9 -1
  10. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/SOURCES.txt +1 -0
  11. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_airflow_utils.py +9 -5
  12. sqlmesh_dag_generator-0.7.3/tests/test_enterprise_features.py +240 -0
  13. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/LICENSE +0 -0
  14. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/pyproject.toml +0 -0
  15. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/setup.cfg +0 -0
  16. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/cli.py +0 -0
  17. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/models.py +0 -0
  18. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/security.py +0 -0
  19. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/utils.py +0 -0
  20. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator/validation.py +0 -0
  21. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/dependency_links.txt +0 -0
  22. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/entry_points.txt +0 -0
  23. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/requires.txt +0 -0
  24. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/sqlmesh_dag_generator.egg-info/top_level.txt +0 -0
  25. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_auto_schedule.py +0 -0
  26. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_bug_fix_user_scenario.py +0 -0
  27. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_cache.py +0 -0
  28. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_config.py +0 -0
  29. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_environment_default.py +0 -0
  30. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_environment_handling.py +0 -0
  31. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_generator.py +0 -0
  32. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_kubernetes.py +0 -0
  33. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_models.py +0 -0
  34. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_runtime_connections.py +0 -0
  35. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/tests/test_source_tables.py +0 -0
  36. {sqlmesh_dag_generator-0.7.2 → sqlmesh_dag_generator-0.7.3}/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.7.2
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.2",
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",
@@ -2,7 +2,7 @@
2
2
  SQLMesh DAG Generator - Open Source Airflow Integration for SQLMesh
3
3
  """
4
4
 
5
- __version__ = "0.7.2"
5
+ __version__ = "0.8.0"
6
6
 
7
7
 
8
8
  from sqlmesh_dag_generator.generator import SQLMeshDAGGenerator
@@ -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
- # Set default_catalog to the database name so SQLMesh omits it from generated SQL
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
@@ -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
- task_id="{task_id}",
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
  # ========================================================================
@@ -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 str(e) and "was not found" in str(e):
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
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlmesh-dag-generator
3
- Version: 0.7.2
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.**
@@ -23,6 +23,7 @@ tests/test_auto_schedule.py
23
23
  tests/test_bug_fix_user_scenario.py
24
24
  tests/test_cache.py
25
25
  tests/test_config.py
26
+ tests/test_enterprise_features.py
26
27
  tests/test_environment_default.py
27
28
  tests/test_environment_handling.py
28
29
  tests/test_generator.py
@@ -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 test_redshift_default_catalog(self):
25
- """Test that Redshift connections include default_catalog for 2-part naming"""
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 set to database name
40
- # This tells SQLMesh to use 2-part naming (schema.table) not 3-part (catalog.schema.table)
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 result["default_catalog"] == "stg_carrier_quality_redshift_db"
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
+