Generating Orchestration Code from Contracts

from airflow import DAG
from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryCreateEmptyDatasetOperator,
    BigQueryCreateEmptyTableOperator,
    BigQueryInsertJobOperator
)
from datetime import datetime, timedelta

default_args = {
    'owner': 'data-engineering',
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    dag_id='gcp_customer_analytics',
    default_args=default_args,
    description='Customer analytics pipeline',
    schedule_interval='@daily',
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=['analytics', 'customers']
) as dag:
    
    create_dataset = BigQueryCreateEmptyDatasetOperator(
        task_id='create_dataset',
        dataset_id='analytics',
        location='US',
        project_id='my-project-id'
    )
    
    create_table = BigQueryCreateEmptyTableOperator(
        task_id='create_table',
        dataset_id='analytics',
        table_id='customers',
        schema_fields=[
            {'name': 'customer_id', 'type': 'INTEGER', 'mode': 'NULLABLE'},
            {'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE'}
        ],
        project_id='my-project-id'
    )
    
    load_data = BigQueryInsertJobOperator(
        task_id='load_data',
        configuration={
            'query': {
                'query': """
                    INSERT INTO analytics.customers
                    SELECT * FROM raw.customer_data
                    WHERE date = CURRENT_DATE()
                """,
                'useLegacySql': False
            }
        },
        project_id='my-project-id'
    )
    
    create_dataset >> create_table >> load_data

from dagster import op, job, resource, In, Out
import boto3

@resource
def aws_s3_client(context):
    return boto3.client('s3', region_name='us-east-1')

@resource
def aws_glue_client(context):
    return boto3.client('glue', region_name='us-east-1')

@op(required_resource_keys={'s3_client'})
def create_bucket(context):
    s3 = context.resources.s3_client
    bucket_name = 'sales-analytics-data'
    
    try:
        s3.create_bucket(Bucket=bucket_name)
        context.log.info(f"Created S3 bucket: {bucket_name}")
        return bucket_name
    except Exception as e:
        context.log.error(f"Error creating bucket: {e}")
        raise

@op(required_resource_keys={'glue_client'}, ins={'bucket': In()})
def create_database(context, bucket):
    glue = context.resources.glue_client
    
    try:
        glue.create_database(
            DatabaseInput={
                'Name': 'sales',
                'Description': 'Sales analytics database'
            }
        )
        context.log.info("Created Glue database: sales")
        return 'sales'
    except Exception as e:
        context.log.error(f"Error creating database: {e}")
        raise

@op(required_resource_keys={'glue_client'}, ins={'database': In()})
def create_table(context, database):
    glue = context.resources.glue_client
    
    try:
        glue.create_table(
            DatabaseName=database,
            TableInput={
                'Name': 'transactions',
                'StorageDescriptor': {
                    'Location': 's3://sales-analytics-data/transactions/',
                    'InputFormat': 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat',
                    'OutputFormat': 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat'
                }
            }
        )
        context.log.info("Created Glue table: transactions")
        return 'transactions'
    except Exception as e:
        context.log.error(f"Error creating table: {e}")
        raise

@job(resource_defs={
    's3_client': aws_s3_client,
    'glue_client': aws_glue_client
})
def aws_sales_analytics():
    database = create_database(create_bucket())
    create_table(database)

from prefect import flow, task
from prefect.deployments import Deployment
from prefect.server.schemas.schedules import CronSchedule
import snowflake.connector

def get_snowflake_connection():
    return snowflake.connector.connect(
        account='xy12345.us-east-1',
        user='...',
        password='...',
        warehouse='COMPUTE_WH',
        database='ANALYTICS'
    )

@task(retries=3, retry_delay_seconds=300)
def create_database():
    conn = get_snowflake_connection()
    cursor = conn.cursor()
    
    try:
        cursor.execute('CREATE DATABASE IF NOT EXISTS ANALYTICS')
        print("Created database: ANALYTICS")
    finally:
        cursor.close()
        conn.close()

@task(retries=3, retry_delay_seconds=300)
def create_schema():
    conn = get_snowflake_connection()
    cursor = conn.cursor()
    
    try:
        cursor.execute('CREATE SCHEMA IF NOT EXISTS INVENTORY')
        print("Created schema: INVENTORY")
    finally:
        cursor.close()
        conn.close()

@task(retries=3, retry_delay_seconds=300)
def create_table():
    conn = get_snowflake_connection()
    cursor = conn.cursor()
    
    try:
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS INVENTORY.STOCK_LEVELS (
                product_id NUMBER,
                quantity NUMBER,
                last_updated TIMESTAMP
            )
        """)
        print("Created table: INVENTORY.STOCK_LEVELS")
    finally:
        cursor.close()
        conn.close()

@task(retries=3, retry_delay_seconds=300)
def load_data():
    conn = get_snowflake_connection()
    cursor = conn.cursor()
    
    try:
        cursor.execute("""
            INSERT INTO INVENTORY.STOCK_LEVELS
            SELECT product_id, SUM(quantity), CURRENT_TIMESTAMP()
            FROM RAW.INVENTORY_UPDATES
            WHERE update_time > DATEADD(hour, -1, CURRENT_TIMESTAMP())
            GROUP BY product_id
        """)
        print(f"Loaded {cursor.rowcount} rows")
    finally:
        cursor.close()
        conn.close()

@flow(name='snowflake_inventory_analytics')
def main():
    create_database()
    create_schema()
    create_table()
    load_data()

# Create deployment
if __name__ == '__main__':
    deployment = Deployment.build_from_flow(
        flow=main,
        name='inventory-analytics-deployment',
        schedule=CronSchedule(cron='0 * * * *'),  # @hourly
        work_queue_name='default'
    )
    deployment.apply()

# Airflow
context.log.info("Processing task...")

# Dagster
context.log.info("Processing op...")

# Prefect
print("Processing task...")  # Captured by Prefect