sqlmesh-dag-generator 0.1.0__py3-none-any.whl

sqlmesh_dag_generator/__init__.py

@@ -0,0 +1,14 @@
+"""
+SQLMesh DAG Generator - Open Source Airflow Integration for SQLMesh
+"""
+
+__version__ = "0.1.0"
+
+from sqlmesh_dag_generator.generator import SQLMeshDAGGenerator
+from sqlmesh_dag_generator.config import DAGGeneratorConfig
+
+__all__ = [
+    "SQLMeshDAGGenerator",
+    "DAGGeneratorConfig",
+]
+

sqlmesh_dag_generator/cli.py

@@ -0,0 +1,202 @@
+"""
+Command-line interface for SQLMesh DAG Generator
+"""
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+from sqlmesh_dag_generator import SQLMeshDAGGenerator
+from sqlmesh_dag_generator.config import DAGGeneratorConfig
+
+
+def setup_logging(verbose: bool = False):
+    """Setup logging configuration"""
+    level = logging.DEBUG if verbose else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    )
+
+
+def main():
+    """Main CLI entry point"""
+    parser = argparse.ArgumentParser(
+        description="Generate Airflow DAGs from SQLMesh projects",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Generate DAG from SQLMesh project
+  sqlmesh-dag-gen --project-path /path/to/sqlmesh --dag-id my_dag
+
+  # Generate with custom output directory
+  sqlmesh-dag-gen --project-path /path/to/sqlmesh --output-dir /path/to/dags
+
+  # Use configuration file
+  sqlmesh-dag-gen --config config.yaml
+
+  # Validate without generating
+  sqlmesh-dag-gen --project-path /path/to/sqlmesh --validate-only
+        """
+    )
+
+    # Configuration source
+    config_group = parser.add_mutually_exclusive_group()
+    config_group.add_argument(
+        '--config', '-c',
+        help='Path to configuration YAML file'
+    )
+    config_group.add_argument(
+        '--project-path', '-p',
+        help='Path to SQLMesh project'
+    )
+
+    # SQLMesh options
+    parser.add_argument(
+        '--environment', '-e',
+        default='prod',
+        help='SQLMesh environment (default: prod)'
+    )
+    parser.add_argument(
+        '--gateway', '-g',
+        help='SQLMesh gateway name'
+    )
+
+    # Airflow DAG options
+    parser.add_argument(
+        '--dag-id',
+        help='Airflow DAG ID'
+    )
+    parser.add_argument(
+        '--schedule',
+        help='Airflow schedule interval (cron expression or preset)'
+    )
+    parser.add_argument(
+        '--tags',
+        nargs='+',
+        default=['sqlmesh'],
+        help='Airflow DAG tags'
+    )
+
+    # Generation options
+    parser.add_argument(
+        '--output-dir', '-o',
+        default='./dags',
+        help='Output directory for generated DAG files (default: ./dags)'
+    )
+    parser.add_argument(
+        '--operator-type',
+        choices=['python', 'bash', 'kubernetes'],
+        default='python',
+        help='Airflow operator type to use (default: python)'
+    )
+    parser.add_argument(
+        '--include-models',
+        nargs='+',
+        help='Only include these models'
+    )
+    parser.add_argument(
+        '--exclude-models',
+        nargs='+',
+        help='Exclude these models'
+    )
+
+    # Actions
+    parser.add_argument(
+        '--validate-only',
+        action='store_true',
+        help='Only validate, do not generate DAG'
+    )
+    parser.add_argument(
+        '--dry-run',
+        action='store_true',
+        help='Generate DAG code but do not write to file'
+    )
+
+    # Other options
+    parser.add_argument(
+        '--verbose', '-v',
+        action='store_true',
+        help='Enable verbose logging'
+    )
+
+    args = parser.parse_args()
+
+    # Setup logging
+    setup_logging(args.verbose)
+    logger = logging.getLogger(__name__)
+
+    try:
+        # Load or build configuration
+        if args.config:
+            logger.info(f"Loading configuration from: {args.config}")
+            config = DAGGeneratorConfig.from_file(args.config)
+        else:
+            if not args.project_path:
+                parser.error("Either --config or --project-path must be provided")
+
+            if not args.dag_id:
+                # Generate DAG ID from project path
+                project_name = Path(args.project_path).name
+                args.dag_id = f"sqlmesh_{project_name}"
+
+            logger.info("Building configuration from command-line arguments")
+            config = DAGGeneratorConfig.from_dict({
+                "sqlmesh": {
+                    "project_path": args.project_path,
+                    "environment": args.environment,
+                    "gateway": args.gateway,
+                },
+                "airflow": {
+                    "dag_id": args.dag_id,
+                    "schedule_interval": args.schedule,
+                    "tags": args.tags,
+                },
+                "generation": {
+                    "output_dir": args.output_dir,
+                    "operator_type": args.operator_type,
+                    "include_models": args.include_models,
+                    "exclude_models": args.exclude_models,
+                    "dry_run": args.dry_run,
+                },
+            })
+
+        # Create generator
+        generator = SQLMeshDAGGenerator(config=config)
+
+        # Validate
+        logger.info("Validating SQLMesh project...")
+        if not generator.validate():
+            logger.error("Validation failed")
+            sys.exit(1)
+
+        logger.info("✓ Validation passed")
+
+        if args.validate_only:
+            logger.info("Validation complete (--validate-only specified)")
+            sys.exit(0)
+
+        # Generate DAG
+        logger.info(f"Generating Airflow DAG: {config.airflow.dag_id}")
+        dag_code = generator.generate_dag()
+
+        if args.dry_run:
+            logger.info("=" * 60)
+            logger.info("Generated DAG (dry-run mode):")
+            logger.info("=" * 60)
+            print(dag_code)
+        else:
+            output_path = generator._get_output_path()
+            logger.info(f"✓ DAG file generated: {output_path}")
+
+        logger.info("Success!")
+        sys.exit(0)
+
+    except Exception as e:
+        logger.error(f"Error: {e}", exc_info=args.verbose)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
+

sqlmesh_dag_generator/config.py

@@ -0,0 +1,116 @@
+"""
+Configuration module for SQLMesh DAG Generator
+"""
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any
+from pathlib import Path
+import yaml
+
+
+@dataclass
+class SQLMeshConfig:
+    """SQLMesh project configuration"""
+    project_path: str
+    environment: str = "prod"
+    gateway: Optional[str] = None
+    config_path: Optional[str] = None
+
+
+@dataclass
+class AirflowConfig:
+    """Airflow DAG configuration"""
+    dag_id: str
+    schedule_interval: Optional[str] = None
+    default_args: Dict[str, Any] = field(default_factory=dict)
+    tags: List[str] = field(default_factory=list)
+    catchup: bool = False
+    max_active_runs: int = 1
+    description: Optional[str] = None
+
+
+@dataclass
+class GenerationConfig:
+    """DAG generation settings"""
+    output_dir: str = "./dags"
+    mode: str = "dynamic"  # "static" or "dynamic" - dynamic is default (fire & forget!)
+    operator_type: str = "python"  # python, bash, or kubernetes
+    include_tests: bool = False
+    parallel_tasks: bool = True
+    max_parallel_tasks: Optional[int] = None
+    include_models: Optional[List[str]] = None
+    exclude_models: Optional[List[str]] = None
+    model_pattern: Optional[str] = None
+    dry_run: bool = False
+
+
+@dataclass
+class DAGGeneratorConfig:
+    """Complete configuration for DAG generator"""
+    sqlmesh: SQLMeshConfig
+    airflow: AirflowConfig
+    generation: GenerationConfig = field(default_factory=GenerationConfig)
+
+    @classmethod
+    def from_file(cls, config_path: str) -> "DAGGeneratorConfig":
+        """Load configuration from YAML file"""
+        config_file = Path(config_path)
+        if not config_file.exists():
+            raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+        with open(config_file, "r") as f:
+            config_data = yaml.safe_load(f)
+
+        return cls(
+            sqlmesh=SQLMeshConfig(**config_data.get("sqlmesh", {})),
+            airflow=AirflowConfig(**config_data.get("airflow", {})),
+            generation=GenerationConfig(**config_data.get("generation", {})),
+        )
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any]) -> "DAGGeneratorConfig":
+        """Load configuration from dictionary"""
+        return cls(
+            sqlmesh=SQLMeshConfig(**config_dict.get("sqlmesh", {})),
+            airflow=AirflowConfig(**config_dict.get("airflow", {})),
+            generation=GenerationConfig(**config_dict.get("generation", {})),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert configuration to dictionary"""
+        return {
+            "sqlmesh": {
+                "project_path": self.sqlmesh.project_path,
+                "environment": self.sqlmesh.environment,
+                "gateway": self.sqlmesh.gateway,
+                "config_path": self.sqlmesh.config_path,
+            },
+            "airflow": {
+                "dag_id": self.airflow.dag_id,
+                "schedule_interval": self.airflow.schedule_interval,
+                "default_args": self.airflow.default_args,
+                "tags": self.airflow.tags,
+                "catchup": self.airflow.catchup,
+                "max_active_runs": self.airflow.max_active_runs,
+                "description": self.airflow.description,
+            },
+            "generation": {
+                "output_dir": self.generation.output_dir,
+                "operator_type": self.generation.operator_type,
+                "include_tests": self.generation.include_tests,
+                "parallel_tasks": self.generation.parallel_tasks,
+                "max_parallel_tasks": self.generation.max_parallel_tasks,
+                "include_models": self.generation.include_models,
+                "exclude_models": self.generation.exclude_models,
+                "model_pattern": self.generation.model_pattern,
+                "dry_run": self.generation.dry_run,
+            },
+        }
+
+    def save(self, output_path: str) -> None:
+        """Save configuration to YAML file"""
+        output_file = Path(output_path)
+        output_file.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_file, "w") as f:
+            yaml.dump(self.to_dict(), f, default_flow_style=False, sort_keys=False)
+