# Example 1: Spark-Snowflake read with configuration options
from pyspark.sql import SparkSession

# Spark session with Snowflake connector
# The connector JAR handles all Snowflake communication
spark = SparkSession.builder \
    .appName("SnowflakeIntegration") \
    .config("spark.jars.packages", 
            "net.snowflake:spark-snowflake_2.12:2.13.0-spark_3.3") \
    .getOrCreate()

# Read from Snowflake
# Parameters:
#   url: Snowflake account URL
#   user/password: Authentication credentials
#   db/schema/warehouse: Snowflake context
#   dbtable: Source table (mutually exclusive with query)
#   query: Custom SQL query (mutually exclusive with dbtable)
df = spark.read \
    .format("snowflake") \
    .option("url", "myaccount.snowflakecomputing.com") \
    .option("user", "myuser") \
    .option("password", "mypassword") \
    .option("db", "ANALYTICS_DB") \
    .option("schema", "PUBLIC") \
    .option("warehouse", "ANALYTICS_WH") \
    .option("dbtable", "sales_data") \
    .load()

# Transform and write back
# mode options: "overwrite" (replace), "append" (add), "ignore" (skip if exists), "error" (fail)
transformed_df = df.filter("amount > 1000") \
    .groupBy("region") \
    .agg({"amount": "sum", "quantity": "count"})

transformed_df.write \
    .format("snowflake") \
    .option("dbtable", "sales_summary") \
    .mode("overwrite") \
    .save()

# Alternative: Use query instead of dbtable for custom reads
df_custom = spark.read \
    .format("snowflake") \
    .option("query", "SELECT * FROM sales WHERE amount > 1000") \
    .load()

# Pushdown predicates (Snowflake executes these filters)
df_filtered = spark.read \
    .format("snowflake") \
    .option("dbtable", "sales_data") \
    .load()
df_filtered = df_filtered.filter(df_filtered["amount"] > 5000)  # Pushed to Snowflake

# Example 2: Airflow DAG for Snowflake pipeline
from airflow import DAG
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

default_args = {
    'owner': 'data_team',
    'depends_on_past': False,           # Don't wait for previous run
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': True,
    'email': ['data-team@company.com'],
    'retries': 2,                        # Retry twice on failure
    'retry_delay': timedelta(minutes=5),
    'execution_timeout': timedelta(hours=2),  # Kill after 2 hours
}

with DAG('snowflake_etl_pipeline',
         default_args=default_args,
         schedule_interval='@daily',     # Run daily at midnight
         catchup=False,                   # Don't backfill missed runs
         max_active_runs=1) as dag:      # One run at a time

    # Extract: Load raw data from staging
    extract_data = SnowflakeOperator(
        task_id='extract_raw_data',
        sql="""
            INSERT INTO raw_sales 
            SELECT * FROM staging_sales 
            WHERE load_date = '{{ ds }}'
        """,
        snowflake_conn_id='snowflake_default',  # Connection ID from Airflow
        warehouse='ETL_WH',
        database='RAW_DB',
        schema='PUBLIC'
    )

    # Validate: Check row count
    validate_data = SnowflakeOperator(
        task_id='validate_data',
        sql="""
            SELECT COUNT(*) as row_count 
            FROM raw_sales 
            WHERE load_date = '{{ ds }}'
        """,
        snowflake_conn_id='snowflake_default'
    )

    # Transform: Run dbt or custom SQL
    transform_data = SnowflakeOperator(
        task_id='transform_data',
        sql="transform_query.sql",        # External SQL file
        parameters={'execution_date': '{{ ds }}'},
        snowflake_conn_id='snowflake_default'
    )

    # Define dependencies (task ordering)
    extract_data >> validate_data >> transform_data

# Example 3: Airflow custom operator for Snowflake data quality
from airflow.models import BaseOperator
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.utils.decorators import apply_defaults

class SnowflakeDataQualityOperator(BaseOperator):
    """Custom operator to validate Snowflake data quality."""
    
    @apply_defaults
    def __init__(
        self,
        sql,
        expected_result,
        comparison='equals',       # 'equals', 'greater_than', 'less_than'
        snowflake_conn_id='snowflake_default',
        *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.sql = sql
        self.expected_result = expected_result
        self.comparison = comparison
        self.snowflake_conn_id = snowflake_conn_id
    
    def execute(self, context):
        hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
        result = hook.get_first(self.sql)
        actual = result[0]
        
        if self.comparison == 'equals':
            if actual != self.expected_result:
                raise ValueError(
                    f"Data quality check failed: {actual} != {self.expected_result}"
                )
        elif self.comparison == 'greater_than':
            if actual <= self.expected_result:
                raise ValueError(
                    f"Data quality check failed: {actual} <= {self.expected_result}"
                )
        
        self.log.info(f"Quality check passed: {actual} {self.comparison} {self.expected_result}")
        return actual

# Usage in DAG
quality_check = SnowflakeDataQualityOperator(
    task_id='check_row_count',
    sql="SELECT COUNT(*) FROM fct_orders WHERE order_date = CURRENT_DATE()",
    expected_result=1000,
    comparison='greater_than',
    snowflake_conn_id='snowflake_default'
)

-- Example 4: dbt staging model (models/staging/stg_orders.sql)
-- Staging models clean and standardize raw data
-- {{ source() }} macro references the source definition in sources.yml
with source as (
    select * from {{ source('raw', 'orders') }}
),

renamed as (
    select
        order_id,                                              -- Keep as-is
        customer_id,                                           -- Keep as-is
        order_date,                                            -- Keep as-is
        amount as order_amount,                                -- Rename for clarity
        status as order_status,                                -- Rename for clarity
        created_at,
        updated_at
    from source
)

select * from renamed

-- Example 5: dbt mart model (models/marts/fct_orders.sql)
-- Mart models create final analytical tables
-- {{ ref() }} macro references other models (builds dependency graph)
with orders as (
    select * from {{ ref('stg_orders') }}
),

customers as (
    select * from {{ ref('dim_customers') }}
),

final as (
    select
        orders.order_id,
        orders.customer_id,
        customers.customer_name,
        customers.customer_segment,
        orders.order_date,
        orders.order_amount,
        orders.order_status,
        DATE_TRUNC('month', orders.order_date) as order_month
    from orders
    left join customers 
        on orders.customer_id = customers.customer_id
)

select * from final

# Example 6: dbt source definition (models/sources.yml)
version: 2

sources:
  - name: raw
    database: RAW_DB
    schema: PUBLIC
    tables:
      - name: orders
        loaded_at_field: updated_at       # Freshness check column
        freshness:
          warn_after: {count: 12, period: hour}   # Warn if > 12 hours stale
          error_after: {count: 24, period: hour}  # Error if > 24 hours stale
        columns:
          - name: order_id
            tests:
              - unique
              - not_null
          - name: customer_id
            tests:
              - not_null
              - relationships:
                  to: ref('dim_customers')
                  field: customer_id
      - name: customers
        columns:
          - name: customer_id
            tests:
              - unique
              - not_null

-- Example 7: dbt test (tests/not_null_order_id.sql)
-- Singular tests are custom SQL queries that return failing rows
SELECT *
FROM {{ ref('fct_orders') }}
WHERE order_id IS NULL

-- Example 8: dbt generic test (tests/greater_than_zero.sql)
-- Generic tests accept arguments and can be reused
{% test greater_than_zero(model, column_name) %}

select *
from {{ model }}
where {{ column_name }} <= 0

{% endtest %}

-- Usage in schema YAML:
-- tests:
--   - greater_than_zero

-- Example 9: External function for API integration
-- External functions call external APIs from SQL
CREATE OR REPLACE EXTERNAL FUNCTION call_enrichment_api(data VARIANT)
RETURNS VARIANT
LANGUAGE JAVA
HANDLER = 'EnrichmentHandler'
ENDPOINT = 'https://api.example.com/enrich'
API_INTEGRATION = my_api_integration;

-- Use external function in queries
SELECT 
    customer_id,
    name,
    call_enrichment_api(OBJECT_CONSTRUCT('email', email)) AS enriched_data
FROM customers;

-- Example 10: Snowpipe for continuous data loading
-- Snowpipe auto-ingests files from cloud storage
CREATE PIPE orders_pipe
    AUTO_INGEST = TRUE
    INTEGRATION = 's3_notification'
    AS
    COPY INTO orders_raw
    FROM @orders_stage
    FILE_FORMAT = (TYPE = 'PARQUET');

-- Monitor pipe status
SELECT * FROM TABLE(INFORMATION_SCHEMA.COPY_LOAD_HISTORY(
    PIPE_NAME => 'orders_pipe',
    START_TIME => DATEADD(day, -1, CURRENT_TIMESTAMP())
));

# Example 11: Python connector for application integration
import snowflake.connector
from snowflake.connector import DictCursor

# Connect to Snowflake
conn = snowflake.connector.connect(
    user='your_user',
    password='your_password',
    account='your_account',
    warehouse='COMPUTE_WH',
    database='ANALYTICS_DB',
    schema='PUBLIC'
)

# Execute queries
cursor = conn.cursor(DictCursor)
cursor.execute("SELECT * FROM sales_data WHERE amount > %s", (1000,))
for row in cursor:
    print(f"Order: {row['ORDER_ID']}, Amount: {row['AMOUNT']}")

# Batch insert
cursor.executemany(
    "INSERT INTO staging_data (id, value) VALUES (%s, %s)",
    [(1, 'a'), (2, 'b'), (3, 'c')]
)

conn.close()