Complete overhaul of senior-data-engineer skill (previously Grade F: 43/100): SKILL.md (~550 lines): - Added table of contents and trigger phrases - 3 actionable workflows: Batch ETL Pipeline, Real-Time Streaming, Data Quality Framework - Architecture decision framework (Batch vs Stream, Lambda vs Kappa) - Tech stack overview with decision matrix - Troubleshooting section with common issues and solutions Reference Files (all rewritten from 81-line boilerplate): - data_pipeline_architecture.md (~700 lines): Lambda/Kappa architectures, batch processing with Spark, stream processing with Kafka/Flink, exactly-once semantics, error handling strategies, orchestration patterns - data_modeling_patterns.md (~650 lines): Dimensional modeling (Star/Snowflake/OBT), SCD Types 0-6 with SQL implementations, Data Vault (Hub/Satellite/Link), dbt best practices, partitioning and clustering strategies - dataops_best_practices.md (~750 lines): Data testing (Great Expectations, dbt), data contracts with YAML definitions, CI/CD pipelines, observability with OpenLineage, incident response runbooks, cost optimization Python Scripts (all rewritten from 101-line placeholders): - pipeline_orchestrator.py (~600 lines): Generates Airflow DAGs, Prefect flows, and Dagster jobs with configurable ETL patterns - data_quality_validator.py (~1640 lines): Schema validation, data profiling, Great Expectations suite generation, data contract validation, anomaly detection - etl_performance_optimizer.py (~1680 lines): SQL query analysis, Spark job optimization, partition strategy recommendations, cost estimation for BigQuery/Snowflake/Redshift/Databricks Resolves #53 Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Alireza Rezvani
GitHub
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,17 +1,28 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Pipeline Orchestrator
|
||||
Production-grade tool for senior data engineer
|
||||
|
||||
Generate pipeline configurations for Airflow, Prefect, and Dagster.
|
||||
Supports ETL pattern generation, dependency management, and scheduling.
|
||||
|
||||
Usage:
|
||||
python pipeline_orchestrator.py generate --type airflow --source postgres --destination snowflake
|
||||
python pipeline_orchestrator.py generate --type prefect --config pipeline.yaml
|
||||
python pipeline_orchestrator.py visualize --dag dags/my_dag.py
|
||||
python pipeline_orchestrator.py validate --dag dags/my_dag.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import yaml
|
||||
import logging
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional, Any
|
||||
from datetime import datetime, timedelta
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
@@ -19,82 +30,847 @@ logging.basicConfig(
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class PipelineOrchestrator:
|
||||
"""Production-grade pipeline orchestrator"""
|
||||
|
||||
def __init__(self, config: Dict):
|
||||
self.config = config
|
||||
self.results = {
|
||||
'status': 'initialized',
|
||||
'start_time': datetime.now().isoformat(),
|
||||
'processed_items': 0
|
||||
}
|
||||
logger.info(f"Initialized {self.__class__.__name__}")
|
||||
|
||||
def validate_config(self) -> bool:
|
||||
"""Validate configuration"""
|
||||
logger.info("Validating configuration...")
|
||||
# Add validation logic
|
||||
logger.info("Configuration validated")
|
||||
|
||||
# ============================================================================
|
||||
# Data Classes
|
||||
# ============================================================================
|
||||
|
||||
@dataclass
|
||||
class SourceConfig:
|
||||
"""Source system configuration."""
|
||||
type: str # postgres, mysql, s3, kafka, api
|
||||
connection_id: str
|
||||
schema: Optional[str] = None
|
||||
tables: List[str] = field(default_factory=list)
|
||||
query: Optional[str] = None
|
||||
incremental_column: Optional[str] = None
|
||||
incremental_strategy: str = "timestamp" # timestamp, id, cdc
|
||||
|
||||
@dataclass
|
||||
class DestinationConfig:
|
||||
"""Destination system configuration."""
|
||||
type: str # snowflake, bigquery, redshift, s3, delta
|
||||
connection_id: str
|
||||
schema: str = "raw"
|
||||
write_mode: str = "append" # append, overwrite, merge
|
||||
partition_by: Optional[str] = None
|
||||
cluster_by: List[str] = field(default_factory=list)
|
||||
|
||||
@dataclass
|
||||
class TaskConfig:
|
||||
"""Individual task configuration."""
|
||||
task_id: str
|
||||
operator: str
|
||||
dependencies: List[str] = field(default_factory=list)
|
||||
params: Dict[str, Any] = field(default_factory=dict)
|
||||
retries: int = 2
|
||||
retry_delay_minutes: int = 5
|
||||
timeout_minutes: int = 60
|
||||
pool: Optional[str] = None
|
||||
priority_weight: int = 1
|
||||
|
||||
@dataclass
|
||||
class PipelineConfig:
|
||||
"""Complete pipeline configuration."""
|
||||
name: str
|
||||
description: str
|
||||
schedule: str # cron expression or @daily, @hourly
|
||||
owner: str = "data-team"
|
||||
tags: List[str] = field(default_factory=list)
|
||||
catchup: bool = False
|
||||
max_active_runs: int = 1
|
||||
default_retries: int = 2
|
||||
source: Optional[SourceConfig] = None
|
||||
destination: Optional[DestinationConfig] = None
|
||||
tasks: List[TaskConfig] = field(default_factory=list)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Pipeline Generators
|
||||
# ============================================================================
|
||||
|
||||
class PipelineGenerator(ABC):
|
||||
"""Abstract base class for pipeline generators."""
|
||||
|
||||
@abstractmethod
|
||||
def generate(self, config: PipelineConfig) -> str:
|
||||
"""Generate pipeline code from config."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def validate(self, code: str) -> Dict[str, Any]:
|
||||
"""Validate generated pipeline code."""
|
||||
pass
|
||||
|
||||
|
||||
class AirflowGenerator(PipelineGenerator):
|
||||
"""Generate Airflow DAG code."""
|
||||
|
||||
OPERATOR_IMPORTS = {
|
||||
'python': 'from airflow.operators.python import PythonOperator',
|
||||
'bash': 'from airflow.operators.bash import BashOperator',
|
||||
'postgres': 'from airflow.providers.postgres.operators.postgres import PostgresOperator',
|
||||
'snowflake': 'from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator',
|
||||
's3': 'from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator',
|
||||
's3_to_snowflake': 'from airflow.providers.snowflake.transfers.s3_to_snowflake import S3ToSnowflakeOperator',
|
||||
'sensor': 'from airflow.sensors.base import BaseSensorOperator',
|
||||
'trigger': 'from airflow.operators.trigger_dagrun import TriggerDagRunOperator',
|
||||
'email': 'from airflow.operators.email import EmailOperator',
|
||||
'slack': 'from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator',
|
||||
}
|
||||
|
||||
def generate(self, config: PipelineConfig) -> str:
|
||||
"""Generate Airflow DAG from configuration."""
|
||||
|
||||
# Collect required imports
|
||||
imports = self._collect_imports(config)
|
||||
|
||||
# Generate DAG code
|
||||
code = self._generate_header(imports)
|
||||
code += self._generate_default_args(config)
|
||||
code += self._generate_dag_definition(config)
|
||||
code += self._generate_tasks(config)
|
||||
code += self._generate_dependencies(config)
|
||||
|
||||
return code
|
||||
|
||||
def _collect_imports(self, config: PipelineConfig) -> List[str]:
|
||||
"""Collect required import statements."""
|
||||
imports = [
|
||||
"from airflow import DAG",
|
||||
"from airflow.utils.dates import days_ago",
|
||||
"from datetime import datetime, timedelta",
|
||||
]
|
||||
|
||||
operators_used = set()
|
||||
for task in config.tasks:
|
||||
op_type = task.operator.split('_')[0].lower()
|
||||
if op_type in self.OPERATOR_IMPORTS:
|
||||
operators_used.add(op_type)
|
||||
|
||||
# Add source/destination specific imports
|
||||
if config.source:
|
||||
if config.source.type == 'postgres':
|
||||
operators_used.add('postgres')
|
||||
elif config.source.type == 's3':
|
||||
operators_used.add('s3')
|
||||
|
||||
if config.destination:
|
||||
if config.destination.type == 'snowflake':
|
||||
operators_used.add('snowflake')
|
||||
operators_used.add('s3_to_snowflake')
|
||||
|
||||
for op in operators_used:
|
||||
if op in self.OPERATOR_IMPORTS:
|
||||
imports.append(self.OPERATOR_IMPORTS[op])
|
||||
|
||||
return imports
|
||||
|
||||
def _generate_header(self, imports: List[str]) -> str:
|
||||
"""Generate file header with imports."""
|
||||
header = '''"""
|
||||
Auto-generated Airflow DAG
|
||||
Generated by Pipeline Orchestrator
|
||||
"""
|
||||
|
||||
'''
|
||||
header += '\n'.join(imports)
|
||||
header += '\n\n'
|
||||
return header
|
||||
|
||||
def _generate_default_args(self, config: PipelineConfig) -> str:
|
||||
"""Generate default_args dictionary."""
|
||||
return f'''
|
||||
default_args = {{
|
||||
'owner': '{config.owner}',
|
||||
'depends_on_past': False,
|
||||
'email_on_failure': True,
|
||||
'email_on_retry': False,
|
||||
'retries': {config.default_retries},
|
||||
'retry_delay': timedelta(minutes=5),
|
||||
}}
|
||||
|
||||
'''
|
||||
|
||||
def _generate_dag_definition(self, config: PipelineConfig) -> str:
|
||||
"""Generate DAG definition."""
|
||||
tags_str = str(config.tags) if config.tags else "[]"
|
||||
|
||||
return f'''
|
||||
with DAG(
|
||||
dag_id='{config.name}',
|
||||
default_args=default_args,
|
||||
description='{config.description}',
|
||||
schedule_interval='{config.schedule}',
|
||||
start_date=days_ago(1),
|
||||
catchup={config.catchup},
|
||||
max_active_runs={config.max_active_runs},
|
||||
tags={tags_str},
|
||||
) as dag:
|
||||
|
||||
'''
|
||||
|
||||
def _generate_tasks(self, config: PipelineConfig) -> str:
|
||||
"""Generate task definitions."""
|
||||
tasks_code = ""
|
||||
|
||||
for task in config.tasks:
|
||||
if 'python' in task.operator.lower():
|
||||
tasks_code += self._generate_python_task(task)
|
||||
elif 'bash' in task.operator.lower():
|
||||
tasks_code += self._generate_bash_task(task)
|
||||
elif 'sql' in task.operator.lower() or 'postgres' in task.operator.lower():
|
||||
tasks_code += self._generate_sql_task(task, config)
|
||||
elif 'snowflake' in task.operator.lower():
|
||||
tasks_code += self._generate_snowflake_task(task)
|
||||
else:
|
||||
tasks_code += self._generate_generic_task(task)
|
||||
|
||||
return tasks_code
|
||||
|
||||
def _generate_python_task(self, task: TaskConfig) -> str:
|
||||
"""Generate PythonOperator task."""
|
||||
callable_name = task.params.get('callable', 'process_data')
|
||||
return f'''
|
||||
def {callable_name}(**kwargs):
|
||||
"""Task: {task.task_id}"""
|
||||
# Add your processing logic here
|
||||
execution_date = kwargs.get('ds')
|
||||
print(f"Processing data for {{execution_date}}")
|
||||
return True
|
||||
|
||||
def process(self) -> Dict:
|
||||
"""Main processing logic"""
|
||||
logger.info("Starting processing...")
|
||||
|
||||
|
||||
{task.task_id} = PythonOperator(
|
||||
task_id='{task.task_id}',
|
||||
python_callable={callable_name},
|
||||
retries={task.retries},
|
||||
retry_delay=timedelta(minutes={task.retry_delay_minutes}),
|
||||
execution_timeout=timedelta(minutes={task.timeout_minutes}),
|
||||
)
|
||||
|
||||
'''
|
||||
|
||||
def _generate_bash_task(self, task: TaskConfig) -> str:
|
||||
"""Generate BashOperator task."""
|
||||
command = task.params.get('command', 'echo "Hello World"')
|
||||
return f'''
|
||||
{task.task_id} = BashOperator(
|
||||
task_id='{task.task_id}',
|
||||
bash_command='{command}',
|
||||
retries={task.retries},
|
||||
retry_delay=timedelta(minutes={task.retry_delay_minutes}),
|
||||
execution_timeout=timedelta(minutes={task.timeout_minutes}),
|
||||
)
|
||||
|
||||
'''
|
||||
|
||||
def _generate_sql_task(self, task: TaskConfig, config: PipelineConfig) -> str:
|
||||
"""Generate SQL operator task."""
|
||||
sql = task.params.get('sql', 'SELECT 1')
|
||||
conn_id = config.source.connection_id if config.source else 'default_conn'
|
||||
|
||||
return f'''
|
||||
{task.task_id} = PostgresOperator(
|
||||
task_id='{task.task_id}',
|
||||
postgres_conn_id='{conn_id}',
|
||||
sql="""{sql}""",
|
||||
retries={task.retries},
|
||||
retry_delay=timedelta(minutes={task.retry_delay_minutes}),
|
||||
)
|
||||
|
||||
'''
|
||||
|
||||
def _generate_snowflake_task(self, task: TaskConfig) -> str:
|
||||
"""Generate SnowflakeOperator task."""
|
||||
sql = task.params.get('sql', 'SELECT 1')
|
||||
return f'''
|
||||
{task.task_id} = SnowflakeOperator(
|
||||
task_id='{task.task_id}',
|
||||
snowflake_conn_id='snowflake_default',
|
||||
sql="""{sql}""",
|
||||
retries={task.retries},
|
||||
retry_delay=timedelta(minutes={task.retry_delay_minutes}),
|
||||
)
|
||||
|
||||
'''
|
||||
|
||||
def _generate_generic_task(self, task: TaskConfig) -> str:
|
||||
"""Generate generic task placeholder."""
|
||||
return f'''
|
||||
# TODO: Implement {task.operator} for {task.task_id}
|
||||
{task.task_id} = PythonOperator(
|
||||
task_id='{task.task_id}',
|
||||
python_callable=lambda: print("{task.task_id}"),
|
||||
)
|
||||
|
||||
'''
|
||||
|
||||
def _generate_dependencies(self, config: PipelineConfig) -> str:
|
||||
"""Generate task dependencies."""
|
||||
deps_code = "\n # Task dependencies\n"
|
||||
|
||||
for task in config.tasks:
|
||||
if task.dependencies:
|
||||
for dep in task.dependencies:
|
||||
deps_code += f" {dep} >> {task.task_id}\n"
|
||||
|
||||
return deps_code
|
||||
|
||||
def validate(self, code: str) -> Dict[str, Any]:
|
||||
"""Validate generated DAG code."""
|
||||
issues = []
|
||||
warnings = []
|
||||
|
||||
# Check for common issues
|
||||
if 'default_args' not in code:
|
||||
issues.append("Missing default_args definition")
|
||||
|
||||
if 'with DAG' not in code:
|
||||
issues.append("Missing DAG context manager")
|
||||
|
||||
if 'schedule_interval' not in code:
|
||||
warnings.append("No schedule_interval defined, DAG won't run automatically")
|
||||
|
||||
# Try to parse the code
|
||||
try:
|
||||
self.validate_config()
|
||||
|
||||
# Main processing
|
||||
result = self._execute()
|
||||
|
||||
self.results['status'] = 'completed'
|
||||
self.results['end_time'] = datetime.now().isoformat()
|
||||
|
||||
logger.info("Processing completed successfully")
|
||||
return self.results
|
||||
|
||||
except Exception as e:
|
||||
self.results['status'] = 'failed'
|
||||
self.results['error'] = str(e)
|
||||
logger.error(f"Processing failed: {e}")
|
||||
raise
|
||||
|
||||
def _execute(self) -> Dict:
|
||||
"""Execute main logic"""
|
||||
# Implementation here
|
||||
return {'success': True}
|
||||
compile(code, '<string>', 'exec')
|
||||
except SyntaxError as e:
|
||||
issues.append(f"Syntax error: {e}")
|
||||
|
||||
return {
|
||||
'valid': len(issues) == 0,
|
||||
'issues': issues,
|
||||
'warnings': warnings
|
||||
}
|
||||
|
||||
|
||||
class PrefectGenerator(PipelineGenerator):
|
||||
"""Generate Prefect flow code."""
|
||||
|
||||
def generate(self, config: PipelineConfig) -> str:
|
||||
"""Generate Prefect flow from configuration."""
|
||||
|
||||
code = self._generate_header()
|
||||
code += self._generate_tasks(config)
|
||||
code += self._generate_flow(config)
|
||||
|
||||
return code
|
||||
|
||||
def _generate_header(self) -> str:
|
||||
"""Generate file header."""
|
||||
return '''"""
|
||||
Auto-generated Prefect Flow
|
||||
Generated by Pipeline Orchestrator
|
||||
"""
|
||||
|
||||
from prefect import flow, task, get_run_logger
|
||||
from prefect.tasks import task_input_hash
|
||||
from datetime import timedelta
|
||||
import pandas as pd
|
||||
|
||||
'''
|
||||
|
||||
def _generate_tasks(self, config: PipelineConfig) -> str:
|
||||
"""Generate Prefect tasks."""
|
||||
tasks_code = ""
|
||||
|
||||
for task_config in config.tasks:
|
||||
cache_expiration = task_config.params.get('cache_hours', 1)
|
||||
tasks_code += f'''
|
||||
@task(
|
||||
name="{task_config.task_id}",
|
||||
retries={task_config.retries},
|
||||
retry_delay_seconds={task_config.retry_delay_minutes * 60},
|
||||
cache_key_fn=task_input_hash,
|
||||
cache_expiration=timedelta(hours={cache_expiration}),
|
||||
)
|
||||
def {task_config.task_id}(input_data=None):
|
||||
"""Task: {task_config.task_id}"""
|
||||
logger = get_run_logger()
|
||||
logger.info(f"Executing {task_config.task_id}")
|
||||
|
||||
# Add processing logic here
|
||||
result = input_data
|
||||
|
||||
return result
|
||||
|
||||
'''
|
||||
return tasks_code
|
||||
|
||||
def _generate_flow(self, config: PipelineConfig) -> str:
|
||||
"""Generate Prefect flow."""
|
||||
flow_code = f'''
|
||||
@flow(
|
||||
name="{config.name}",
|
||||
description="{config.description}",
|
||||
version="1.0.0",
|
||||
)
|
||||
def {config.name.replace('-', '_')}_flow():
|
||||
"""Main flow orchestrating all tasks."""
|
||||
logger = get_run_logger()
|
||||
logger.info("Starting flow: {config.name}")
|
||||
|
||||
'''
|
||||
# Generate task calls with dependencies
|
||||
task_vars = {}
|
||||
for i, task_config in enumerate(config.tasks):
|
||||
task_name = task_config.task_id
|
||||
var_name = f"result_{i}"
|
||||
task_vars[task_name] = var_name
|
||||
|
||||
if task_config.dependencies:
|
||||
# Get input from first dependency
|
||||
dep_var = task_vars.get(task_config.dependencies[0], "None")
|
||||
flow_code += f" {var_name} = {task_name}({dep_var})\n"
|
||||
else:
|
||||
flow_code += f" {var_name} = {task_name}()\n"
|
||||
|
||||
flow_code += '''
|
||||
logger.info("Flow completed successfully")
|
||||
return True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
''' + f'{config.name.replace("-", "_")}_flow()' + '\n'
|
||||
|
||||
return flow_code
|
||||
|
||||
def validate(self, code: str) -> Dict[str, Any]:
|
||||
"""Validate Prefect flow code."""
|
||||
issues = []
|
||||
|
||||
if '@flow' not in code:
|
||||
issues.append("Missing @flow decorator")
|
||||
|
||||
if '@task' not in code:
|
||||
issues.append("No tasks defined with @task decorator")
|
||||
|
||||
try:
|
||||
compile(code, '<string>', 'exec')
|
||||
except SyntaxError as e:
|
||||
issues.append(f"Syntax error: {e}")
|
||||
|
||||
return {
|
||||
'valid': len(issues) == 0,
|
||||
'issues': issues,
|
||||
'warnings': []
|
||||
}
|
||||
|
||||
|
||||
class DagsterGenerator(PipelineGenerator):
|
||||
"""Generate Dagster job code."""
|
||||
|
||||
def generate(self, config: PipelineConfig) -> str:
|
||||
"""Generate Dagster job from configuration."""
|
||||
|
||||
code = self._generate_header()
|
||||
code += self._generate_ops(config)
|
||||
code += self._generate_job(config)
|
||||
|
||||
return code
|
||||
|
||||
def _generate_header(self) -> str:
|
||||
"""Generate file header."""
|
||||
return '''"""
|
||||
Auto-generated Dagster Job
|
||||
Generated by Pipeline Orchestrator
|
||||
"""
|
||||
|
||||
from dagster import op, job, In, Out, Output, DynamicOut, graph
|
||||
from dagster import AssetMaterialization, MetadataValue
|
||||
import pandas as pd
|
||||
|
||||
'''
|
||||
|
||||
def _generate_ops(self, config: PipelineConfig) -> str:
|
||||
"""Generate Dagster ops."""
|
||||
ops_code = ""
|
||||
|
||||
for task_config in config.tasks:
|
||||
has_input = len(task_config.dependencies) > 0
|
||||
|
||||
if has_input:
|
||||
ops_code += f'''
|
||||
@op(
|
||||
ins={{"input_data": In()}},
|
||||
out=Out(),
|
||||
)
|
||||
def {task_config.task_id}(context, input_data):
|
||||
"""Op: {task_config.task_id}"""
|
||||
context.log.info(f"Executing {task_config.task_id}")
|
||||
|
||||
# Add processing logic here
|
||||
result = input_data
|
||||
|
||||
# Log asset materialization
|
||||
yield AssetMaterialization(
|
||||
asset_key="{task_config.task_id}",
|
||||
metadata={{
|
||||
"row_count": MetadataValue.int(len(result) if hasattr(result, '__len__') else 0),
|
||||
}}
|
||||
)
|
||||
yield Output(result)
|
||||
|
||||
'''
|
||||
else:
|
||||
ops_code += f'''
|
||||
@op(out=Out())
|
||||
def {task_config.task_id}(context):
|
||||
"""Op: {task_config.task_id}"""
|
||||
context.log.info(f"Executing {task_config.task_id}")
|
||||
|
||||
# Add processing logic here
|
||||
result = {{}}
|
||||
|
||||
yield AssetMaterialization(
|
||||
asset_key="{task_config.task_id}",
|
||||
)
|
||||
yield Output(result)
|
||||
|
||||
'''
|
||||
return ops_code
|
||||
|
||||
def _generate_job(self, config: PipelineConfig) -> str:
|
||||
"""Generate Dagster job."""
|
||||
job_code = f'''
|
||||
@job(
|
||||
name="{config.name}",
|
||||
description="{config.description}",
|
||||
tags={{
|
||||
"owner": "{config.owner}",
|
||||
"schedule": "{config.schedule}",
|
||||
}},
|
||||
)
|
||||
def {config.name.replace('-', '_')}_job():
|
||||
"""Main job orchestrating all ops."""
|
||||
'''
|
||||
# Build dependency graph
|
||||
task_outputs = {}
|
||||
for task_config in config.tasks:
|
||||
task_name = task_config.task_id
|
||||
|
||||
if task_config.dependencies:
|
||||
dep_output = task_outputs.get(task_config.dependencies[0], None)
|
||||
if dep_output:
|
||||
job_code += f" {task_name}_output = {task_name}({dep_output})\n"
|
||||
else:
|
||||
job_code += f" {task_name}_output = {task_name}()\n"
|
||||
else:
|
||||
job_code += f" {task_name}_output = {task_name}()\n"
|
||||
|
||||
task_outputs[task_name] = f"{task_name}_output"
|
||||
|
||||
return job_code
|
||||
|
||||
def validate(self, code: str) -> Dict[str, Any]:
|
||||
"""Validate Dagster job code."""
|
||||
issues = []
|
||||
|
||||
if '@job' not in code:
|
||||
issues.append("Missing @job decorator")
|
||||
|
||||
if '@op' not in code:
|
||||
issues.append("No ops defined with @op decorator")
|
||||
|
||||
try:
|
||||
compile(code, '<string>', 'exec')
|
||||
except SyntaxError as e:
|
||||
issues.append(f"Syntax error: {e}")
|
||||
|
||||
return {
|
||||
'valid': len(issues) == 0,
|
||||
'issues': issues,
|
||||
'warnings': []
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ETL Pattern Templates
|
||||
# ============================================================================
|
||||
|
||||
class ETLPatternGenerator:
|
||||
"""Generate common ETL patterns."""
|
||||
|
||||
@staticmethod
|
||||
def generate_extract_load(
|
||||
source_type: str,
|
||||
destination_type: str,
|
||||
tables: List[str],
|
||||
mode: str = "incremental"
|
||||
) -> PipelineConfig:
|
||||
"""Generate extract-load pipeline configuration."""
|
||||
|
||||
tasks = []
|
||||
|
||||
# Extract tasks
|
||||
for table in tables:
|
||||
extract_task = TaskConfig(
|
||||
task_id=f"extract_{table}",
|
||||
operator="python_operator",
|
||||
params={
|
||||
'callable': f'extract_{table}',
|
||||
'sql': f'SELECT * FROM {table}' + (
|
||||
' WHERE updated_at > {{{{ prev_ds }}}}' if mode == 'incremental' else ''
|
||||
)
|
||||
}
|
||||
)
|
||||
tasks.append(extract_task)
|
||||
|
||||
# Load tasks with dependencies
|
||||
for table in tables:
|
||||
load_task = TaskConfig(
|
||||
task_id=f"load_{table}",
|
||||
operator="python_operator",
|
||||
dependencies=[f"extract_{table}"],
|
||||
params={'callable': f'load_{table}'}
|
||||
)
|
||||
tasks.append(load_task)
|
||||
|
||||
# Quality check task
|
||||
quality_task = TaskConfig(
|
||||
task_id="quality_check",
|
||||
operator="python_operator",
|
||||
dependencies=[f"load_{table}" for table in tables],
|
||||
params={'callable': 'run_quality_checks'}
|
||||
)
|
||||
tasks.append(quality_task)
|
||||
|
||||
return PipelineConfig(
|
||||
name=f"el_{source_type}_to_{destination_type}",
|
||||
description=f"Extract from {source_type}, load to {destination_type}",
|
||||
schedule="0 5 * * *", # Daily at 5 AM
|
||||
tags=["etl", source_type, destination_type],
|
||||
source=SourceConfig(
|
||||
type=source_type,
|
||||
connection_id=f"{source_type}_default",
|
||||
tables=tables,
|
||||
incremental_strategy="timestamp" if mode == "incremental" else "full"
|
||||
),
|
||||
destination=DestinationConfig(
|
||||
type=destination_type,
|
||||
connection_id=f"{destination_type}_default",
|
||||
write_mode="append" if mode == "incremental" else "overwrite"
|
||||
),
|
||||
tasks=tasks
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def generate_transform_pipeline(
|
||||
source_tables: List[str],
|
||||
target_table: str,
|
||||
dbt_models: List[str]
|
||||
) -> PipelineConfig:
|
||||
"""Generate transformation pipeline with dbt."""
|
||||
|
||||
tasks = []
|
||||
|
||||
# Sensor for source freshness
|
||||
for table in source_tables:
|
||||
sensor_task = TaskConfig(
|
||||
task_id=f"wait_for_{table}",
|
||||
operator="sql_sensor",
|
||||
params={
|
||||
'sql': f"SELECT MAX(updated_at) FROM {table} WHERE updated_at > '{{{{ ds }}}}'"
|
||||
}
|
||||
)
|
||||
tasks.append(sensor_task)
|
||||
|
||||
# dbt run task
|
||||
dbt_run = TaskConfig(
|
||||
task_id="dbt_run",
|
||||
operator="bash_operator",
|
||||
dependencies=[f"wait_for_{t}" for t in source_tables],
|
||||
params={
|
||||
'command': f'cd /opt/dbt && dbt run --select {" ".join(dbt_models)}'
|
||||
},
|
||||
timeout_minutes=120
|
||||
)
|
||||
tasks.append(dbt_run)
|
||||
|
||||
# dbt test task
|
||||
dbt_test = TaskConfig(
|
||||
task_id="dbt_test",
|
||||
operator="bash_operator",
|
||||
dependencies=["dbt_run"],
|
||||
params={
|
||||
'command': f'cd /opt/dbt && dbt test --select {" ".join(dbt_models)}'
|
||||
}
|
||||
)
|
||||
tasks.append(dbt_test)
|
||||
|
||||
return PipelineConfig(
|
||||
name=f"transform_{target_table}",
|
||||
description=f"Transform data into {target_table} using dbt",
|
||||
schedule="0 6 * * *", # Daily at 6 AM (after extraction)
|
||||
tags=["transform", "dbt"],
|
||||
tasks=tasks
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# CLI Interface
|
||||
# ============================================================================
|
||||
|
||||
def main():
|
||||
"""Main entry point"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Pipeline Orchestrator"
|
||||
description="Pipeline Orchestrator - Generate and manage data pipeline configurations",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
Generate Airflow DAG:
|
||||
python pipeline_orchestrator.py generate --type airflow --source postgres --destination snowflake --tables orders,customers
|
||||
|
||||
Generate from config file:
|
||||
python pipeline_orchestrator.py generate --config pipeline.yaml --type prefect
|
||||
|
||||
Validate existing DAG:
|
||||
python pipeline_orchestrator.py validate --dag dags/my_dag.py --type airflow
|
||||
"""
|
||||
)
|
||||
parser.add_argument('--input', '-i', required=True, help='Input path')
|
||||
parser.add_argument('--output', '-o', required=True, help='Output path')
|
||||
parser.add_argument('--config', '-c', help='Configuration file')
|
||||
parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
|
||||
|
||||
|
||||
subparsers = parser.add_subparsers(dest='command', help='Command to run')
|
||||
|
||||
# Generate command
|
||||
gen_parser = subparsers.add_parser('generate', help='Generate pipeline code')
|
||||
gen_parser.add_argument('--type', '-t', required=True,
|
||||
choices=['airflow', 'prefect', 'dagster'],
|
||||
help='Pipeline framework type')
|
||||
gen_parser.add_argument('--source', '-s', help='Source system type')
|
||||
gen_parser.add_argument('--destination', '-d', help='Destination system type')
|
||||
gen_parser.add_argument('--tables', help='Comma-separated list of tables')
|
||||
gen_parser.add_argument('--config', '-c', help='Configuration YAML file')
|
||||
gen_parser.add_argument('--output', '-o', help='Output file path')
|
||||
gen_parser.add_argument('--name', '-n', help='Pipeline name')
|
||||
gen_parser.add_argument('--schedule', default='0 5 * * *', help='Cron schedule')
|
||||
gen_parser.add_argument('--mode', default='incremental',
|
||||
choices=['incremental', 'full'],
|
||||
help='Load mode')
|
||||
|
||||
# Validate command
|
||||
val_parser = subparsers.add_parser('validate', help='Validate pipeline code')
|
||||
val_parser.add_argument('--dag', required=True, help='DAG file to validate')
|
||||
val_parser.add_argument('--type', '-t', required=True,
|
||||
choices=['airflow', 'prefect', 'dagster'])
|
||||
|
||||
# Template command
|
||||
tmpl_parser = subparsers.add_parser('template', help='Generate from template')
|
||||
tmpl_parser.add_argument('--pattern', '-p', required=True,
|
||||
choices=['extract-load', 'transform', 'cdc'],
|
||||
help='ETL pattern to generate')
|
||||
tmpl_parser.add_argument('--type', '-t', required=True,
|
||||
choices=['airflow', 'prefect', 'dagster'])
|
||||
tmpl_parser.add_argument('--source', '-s', required=True)
|
||||
tmpl_parser.add_argument('--destination', '-d', required=True)
|
||||
tmpl_parser.add_argument('--tables', required=True)
|
||||
tmpl_parser.add_argument('--output', '-o', help='Output file path')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.verbose:
|
||||
logging.getLogger().setLevel(logging.DEBUG)
|
||||
|
||||
try:
|
||||
config = {
|
||||
'input': args.input,
|
||||
'output': args.output
|
||||
}
|
||||
|
||||
processor = PipelineOrchestrator(config)
|
||||
results = processor.process()
|
||||
|
||||
print(json.dumps(results, indent=2))
|
||||
sys.exit(0)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Fatal error: {e}")
|
||||
|
||||
if args.command is None:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
if args.command == 'generate':
|
||||
# Load config if provided
|
||||
if args.config:
|
||||
with open(args.config) as f:
|
||||
config_data = yaml.safe_load(f)
|
||||
config = PipelineConfig(**config_data)
|
||||
else:
|
||||
# Build config from arguments
|
||||
tables = args.tables.split(',') if args.tables else []
|
||||
|
||||
config = ETLPatternGenerator.generate_extract_load(
|
||||
source_type=args.source or 'postgres',
|
||||
destination_type=args.destination or 'snowflake',
|
||||
tables=tables,
|
||||
mode=args.mode
|
||||
)
|
||||
|
||||
if args.name:
|
||||
config.name = args.name
|
||||
config.schedule = args.schedule
|
||||
|
||||
# Generate code
|
||||
generators = {
|
||||
'airflow': AirflowGenerator(),
|
||||
'prefect': PrefectGenerator(),
|
||||
'dagster': DagsterGenerator()
|
||||
}
|
||||
|
||||
generator = generators[args.type]
|
||||
code = generator.generate(config)
|
||||
|
||||
# Validate
|
||||
validation = generator.validate(code)
|
||||
if not validation['valid']:
|
||||
logger.warning(f"Validation issues: {validation['issues']}")
|
||||
|
||||
# Output
|
||||
if args.output:
|
||||
with open(args.output, 'w') as f:
|
||||
f.write(code)
|
||||
logger.info(f"Generated pipeline saved to {args.output}")
|
||||
else:
|
||||
print(code)
|
||||
|
||||
elif args.command == 'validate':
|
||||
with open(args.dag) as f:
|
||||
code = f.read()
|
||||
|
||||
generators = {
|
||||
'airflow': AirflowGenerator(),
|
||||
'prefect': PrefectGenerator(),
|
||||
'dagster': DagsterGenerator()
|
||||
}
|
||||
|
||||
generator = generators[args.type]
|
||||
result = generator.validate(code)
|
||||
|
||||
print(json.dumps(result, indent=2))
|
||||
sys.exit(0 if result['valid'] else 1)
|
||||
|
||||
elif args.command == 'template':
|
||||
tables = args.tables.split(',')
|
||||
|
||||
if args.pattern == 'extract-load':
|
||||
config = ETLPatternGenerator.generate_extract_load(
|
||||
source_type=args.source,
|
||||
destination_type=args.destination,
|
||||
tables=tables
|
||||
)
|
||||
elif args.pattern == 'transform':
|
||||
config = ETLPatternGenerator.generate_transform_pipeline(
|
||||
source_tables=tables,
|
||||
target_table='fct_output',
|
||||
dbt_models=['stg_*', 'fct_*']
|
||||
)
|
||||
else:
|
||||
logger.error(f"Pattern {args.pattern} not yet implemented")
|
||||
sys.exit(1)
|
||||
|
||||
generators = {
|
||||
'airflow': AirflowGenerator(),
|
||||
'prefect': PrefectGenerator(),
|
||||
'dagster': DagsterGenerator()
|
||||
}
|
||||
|
||||
generator = generators[args.type]
|
||||
code = generator.generate(config)
|
||||
|
||||
if args.output:
|
||||
with open(args.output, 'w') as f:
|
||||
f.write(code)
|
||||
logger.info(f"Generated {args.pattern} pipeline saved to {args.output}")
|
||||
else:
|
||||
print(code)
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user