# Snowflake warehouse is set as a connector option
snowflake_options = {
    "sfURL": "myorg-myaccount.snowflakecomputing.com",
    "sfDatabase": "ANALYTICS",
    "sfSchema": "PUBLIC",
    "sfWarehouse": "SPARK_WH",  # which warehouse runs the load
    "sfRole": "DATA_ENGINEER",
}

# Full Snowflake path: DATABASE.SCHEMA.TABLE
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ANALYTICS.PUBLIC.SALES") \
    .load()

# Or use sfDatabase + sfSchema options + short table name
snowflake_options["sfDatabase"] = "ANALYTICS"
snowflake_options["sfSchema"]   = "PUBLIC"
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "SALES") \  # short name works with sfDatabase+sfSchema set
    .load()

spark-submit \
  --packages net.snowflake:spark-snowflake_2.12:2.12.0-spark_3.4,\
             net.snowflake:snowflake-jdbc:3.14.4 \
  my_pyspark_job.py

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("SnowflakeIntegration") \
    .config("spark.jars.packages",
            "net.snowflake:spark-snowflake_2.12:2.12.0-spark_3.4,"
            "net.snowflake:snowflake-jdbc:3.14.4") \
    .getOrCreate()

# On Databricks — connector is pre-installed, just install via cluster libraries
# On EMR — add to bootstrap or use --packages flag

snowflake_options = {
    # Connection
    "sfURL"      : "myorg-myaccount.snowflakecomputing.com",  # account URL
    "sfUser"     : "spark_user",
    "sfPassword" : "s3cr3t",          # ← BAD for prod, see section 30.9

    # Target
    "sfDatabase" : "ANALYTICS",
    "sfSchema"   : "PUBLIC",
    "sfWarehouse": "SPARK_WH",

    # Optional — role to assume (Snowflake RBAC)
    "sfRole"     : "DATA_ENGINEER",
}

# Use it in reads/writes
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .load()

# Read entire ORDERS table into a Spark DataFrame
df_orders = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .load()

df_orders.printSchema()
df_orders.show(5)

# Only read 2024 orders — Snowflake filters before sending to Spark
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("query",
            "SELECT order_id, customer_id, amount, order_date "
            "FROM ORDERS WHERE order_date >= '2024-01-01'") \
    .load()

# Can also use multi-line string
sql = """
    SELECT
        c.customer_name,
        SUM(o.amount) AS total_spend
    FROM ORDERS o
    JOIN CUSTOMERS c ON o.customer_id = c.customer_id
    GROUP BY c.customer_name
"""

df_summary = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("query", sql) \
    .load()

from pyspark.sql.types import StructType, StructField, StringType, DoubleType, DateType

# Option 1: Let connector infer schema (default)
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .load()

# Option 2: Provide explicit schema — faster, no extra metadata call
schema = StructType([
    StructField("ORDER_ID",    StringType(), nullable=False),
    StructField("AMOUNT",      DoubleType(), nullable=True),
    StructField("ORDER_DATE",  DateType(),   nullable=True),
])

df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .schema(schema) \
    .load()

from pyspark.sql.functions import from_json, col, get_json_object
from pyspark.sql.types import StructType, StructField, StringType, DoubleType

# Snowflake table has: ORDER_ID (STRING), METADATA (VARIANT)
# METADATA looks like: {"channel": "web", "discount_pct": 0.10}

df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .load()

# METADATA arrives as StringType — parse it
metadata_schema = StructType([
    StructField("channel",      StringType(), True),
    StructField("discount_pct", DoubleType(), True),
])

df_parsed = df.withColumn(
    "metadata_parsed",
    from_json(col("METADATA"), metadata_schema)
).select(
    col("ORDER_ID"),
    col("metadata_parsed.channel").alias("channel"),
    col("metadata_parsed.discount_pct").alias("discount_pct")
)

df_result.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_PROCESSED") \
    .mode("append") \     # append / overwrite / ignore / error
    .save()

# Safer "overwrite" — keeps table structure, just empties rows first
df_result.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_PROCESSED") \
    .option("preactions", "TRUNCATE TABLE ORDERS_PROCESSED") \
    .mode("append") \   # append after truncate = logical overwrite
    .save()

import json

# Spark col name → Snowflake col name mapping
# Use when your DataFrame cols don't match Snowflake cols exactly
column_map = {
    "order_id"   : "ORDER_ID",
    "cust_id"    : "CUSTOMER_ID",    # renaming
    "amt"        : "AMOUNT",          # renaming
    "order_date" : "ORDER_DATE",
}

df_result.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .option("columnmap", json.dumps(column_map)) \
    .mode("append") \
    .save()

# Scenario: Load to staging table, then MERGE into production table

df_result.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_STAGING") \
    .option("preactions",
            "TRUNCATE TABLE ORDERS_STAGING") \
    .option("postactions", """
        MERGE INTO ORDERS AS target
        USING ORDERS_STAGING AS src
            ON target.ORDER_ID = src.ORDER_ID
        WHEN MATCHED THEN UPDATE SET
            target.AMOUNT = src.AMOUNT,
            target.STATUS = src.STATUS
        WHEN NOT MATCHED THEN INSERT (ORDER_ID, AMOUNT, STATUS)
            VALUES (src.ORDER_ID, src.AMOUNT, src.STATUS)
    """) \
    .mode("append") \
    .save()

# Multiple statements: separate with semicolons
.option("preactions", "TRUNCATE TABLE STAGING; ALTER SESSION SET TIMEZONE='UTC'")

-- Step 1: Create a storage integration (one-time, by admin)
CREATE STORAGE INTEGRATION s3_int
    TYPE = EXTERNAL_STAGE
    STORAGE_PROVIDER = 'S3'
    ENABLED = TRUE
    STORAGE_AWS_ROLE_ARN = 'arn:aws:iam::123456789:role/SnowflakeS3Role'
    STORAGE_ALLOWED_LOCATIONS = ('s3://my-data-lake/snowflake-stage/');

-- Step 2: Create the external stage
CREATE STAGE my_s3_stage
    URL = 's3://my-data-lake/snowflake-stage/'
    STORAGE_INTEGRATION = s3_int
    FILE_FORMAT = (TYPE = PARQUET);

# Step 1: Spark writes Parquet to S3
output_path = "s3://my-data-lake/snowflake-stage/orders/"
df_result.write \
    .mode("overwrite") \
    .parquet(output_path)

# Step 2: Load from S3 stage into Snowflake using boto3 or Snowflake session
# Option A: Use connector with external stage option
df_result.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .option("tempdir", "s3://my-data-lake/snowflake-stage/") \
    .option("aws_access_key", "...") \   # or use IAM role
    .option("aws_secret_key", "...") \
    .mode("append") \
    .save()

import snowflake.connector

# Step 1: Write Spark output to S3
df_result.write.mode("overwrite").parquet("s3://bucket/stage/orders/")

# Step 2: Connect to Snowflake and run COPY INTO
conn = snowflake.connector.connect(
    user="spark_user",
    account="myorg-myaccount",
    private_key=load_private_key(),   # see section 30.9
    database="ANALYTICS",
    schema="PUBLIC",
    warehouse="SPARK_WH",
)
cursor = conn.cursor()
cursor.execute("""
    COPY INTO ORDERS
    FROM @my_s3_stage/orders/
    FILE_FORMAT = (TYPE = PARQUET)
    MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE
    PURGE = TRUE
""")
cursor.close()
conn.close()

df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .load()

# Apply a filter
df_filtered = df.filter(df["ORDER_DATE"] >= "2024-01-01") \
                .select("ORDER_ID", "AMOUNT")

# Check the physical plan
# Look for "SnowflakeRDD" or "SnowflakeRelation" in the plan
# If you see a Spark Filter on top of SnowflakeRDD → pushdown FAILED
# If you see SnowflakeRDD with the filter baked in → pushdown SUCCESS
df_filtered.explain(True)

# Enable pushdown logging to see the actual SQL sent to Snowflake
spark.conf.set("spark.snowflake.query.pushdown", "true")
# connector logs: "Pushdown query: SELECT ORDER_ID, AMOUNT FROM ORDERS WHERE ORDER_DATE >= '2024-01-01'"

# Disable globally for the session
SnowflakeConnectorUtils = spark._jvm.net.snowflake.spark.snowflake.Utils
SnowflakeConnectorUtils.disablePushdownSession(spark._jvm.org.apache.spark.sql.functions)

# Or per-read via option
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .option("autopushdown", "off") \
    .load()

# Re-enable globally
SnowflakeConnectorUtils.enablePushdownSession(spark._jvm.org.apache.spark.sql.functions)

from datetime import datetime
import boto3, json

# Step 1: Fetch last watermark from DynamoDB control table
dynamodb = boto3.client("dynamodb")
response = dynamodb.get_item(
    TableName="pipeline_watermarks",
    Key={"pipeline_id": {"S": "snowflake_orders_incremental"}}
)
last_watermark = response["Item"]["last_watermark"]["S"]
# e.g., "2024-06-01 00:00:00"

# Step 2: Read only new records from Snowflake
query = f"""
    SELECT *
    FROM ORDERS
    WHERE UPDATED_AT > '{last_watermark}'
    ORDER BY UPDATED_AT
"""

df_incremental = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("query", query) \
    .load()

# Step 3: Process…
# Step 4: Update watermark to max UPDATED_AT in this batch
new_watermark = df_incremental.agg({"UPDATED_AT": "max"}).collect()[0][0]
dynamodb.update_item(
    TableName="pipeline_watermarks",
    Key={"pipeline_id": {"S": "snowflake_orders_incremental"}},
    UpdateExpression="SET last_watermark = :wm",
    ExpressionAttributeValues={":wm": {"S": str(new_watermark)}}
)

# Write incremental data to staging, then MERGE into target
merge_sql = """
    MERGE INTO ORDERS_GOLD AS target
    USING ORDERS_STAGING AS src
        ON target.ORDER_ID = src.ORDER_ID
    WHEN MATCHED AND src.UPDATED_AT > target.UPDATED_AT THEN
        UPDATE SET
            target.STATUS     = src.STATUS,
            target.AMOUNT     = src.AMOUNT,
            target.UPDATED_AT = src.UPDATED_AT
    WHEN NOT MATCHED THEN
        INSERT (ORDER_ID, CUSTOMER_ID, AMOUNT, STATUS, CREATED_AT, UPDATED_AT)
        VALUES (src.ORDER_ID, src.CUSTOMER_ID, src.AMOUNT,
                src.STATUS, src.CREATED_AT, src.UPDATED_AT)
"""

df_incremental.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_STAGING") \
    .option("preactions",  "TRUNCATE TABLE ORDERS_STAGING") \
    .option("postactions", merge_sql) \
    .mode("append") \
    .save()

scd2_merge = """
    -- Step 1: Close changed existing rows
    UPDATE CUSTOMERS_DIM AS target
    SET
        target.EFFECTIVE_END  = src.EFFECTIVE_START,
        target.IS_CURRENT     = FALSE
    FROM CUSTOMERS_STAGING AS src
    WHERE target.CUSTOMER_ID = src.CUSTOMER_ID
      AND target.IS_CURRENT  = TRUE
      AND (target.EMAIL != src.EMAIL OR target.CITY != src.CITY);

    -- Step 2: Insert new rows (both truly new + changed ones)
    INSERT INTO CUSTOMERS_DIM
        (CUSTOMER_ID, EMAIL, CITY, EFFECTIVE_START, EFFECTIVE_END, IS_CURRENT)
    SELECT
        src.CUSTOMER_ID,
        src.EMAIL,
        src.CITY,
        CURRENT_TIMESTAMP()  AS EFFECTIVE_START,
        NULL                 AS EFFECTIVE_END,
        TRUE                 AS IS_CURRENT
    FROM CUSTOMERS_STAGING src
    WHERE NOT EXISTS (
        SELECT 1 FROM CUSTOMERS_DIM tgt
        WHERE tgt.CUSTOMER_ID = src.CUSTOMER_ID
          AND tgt.IS_CURRENT  = TRUE
          AND tgt.EMAIL       = src.EMAIL
          AND tgt.CITY        = src.CITY
    )
"""

df_customers.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "CUSTOMERS_STAGING") \
    .option("preactions",  "TRUNCATE TABLE CUSTOMERS_STAGING") \
    .option("postactions", scd2_merge) \
    .mode("append") \
    .save()

-- Create a stream on the ORDERS table
CREATE STREAM orders_stream ON TABLE ORDERS;

-- Stream shows: METADATA$ACTION (INSERT/DELETE), METADATA$ISUPDATE (TRUE/FALSE)
SELECT *, METADATA$ACTION, METADATA$ISUPDATE FROM orders_stream;

# Read the stream — gets all changes since last consume
df_changes = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "orders_stream") \
    .load()

# Separate inserts/updates from deletes
df_upserts = df_changes.filter(df_changes["METADATA$ACTION"] == "INSERT")
df_deletes = df_changes.filter(df_changes["METADATA$ACTION"] == "DELETE")

# Process upserts and deletes separately...
# Note: Stream is consumed (advanced) when you commit the stream read via a DML

# Option: Run DDL pre-action to resize warehouse before load
df_large.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "LARGE_FACT_TABLE") \
    .option("preactions",
            "ALTER WAREHOUSE SPARK_WH SET WAREHOUSE_SIZE = 'X-LARGE'") \
    .option("postactions",
            "ALTER WAREHOUSE SPARK_WH SET WAREHOUSE_SIZE = 'X-SMALL'") \
    .mode("append") \
    .save()

# Partition by ORDER_ID into 10 parallel reads
# Connector splits: ORDER_ID % 10 = 0, 1, 2 ... 9
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .option("partition_size_in_mb", "256") \   # target size per partition
    .load()

# Alternative: manually split with query ranges
# Read partition 1 of 4: ORDER_DATE in Q1
queries = [
    "SELECT * FROM ORDERS WHERE ORDER_DATE BETWEEN '2024-01-01' AND '2024-03-31'",
    "SELECT * FROM ORDERS WHERE ORDER_DATE BETWEEN '2024-04-01' AND '2024-06-30'",
    "SELECT * FROM ORDERS WHERE ORDER_DATE BETWEEN '2024-07-01' AND '2024-09-30'",
    "SELECT * FROM ORDERS WHERE ORDER_DATE BETWEEN '2024-10-01' AND '2024-12-31'",
]

from functools import reduce

dfs = [
    spark.read
        .format("net.snowflake.spark.snowflake")
        .options(**snowflake_options)
        .option("query", q)
        .load()
    for q in queries
]

df_full = reduce(lambda a, b: a.union(b), dfs)

CREATE WAREHOUSE SPARK_WH
    WAREHOUSE_SIZE = 'SMALL'
    AUTO_SUSPEND = 60         -- suspend after 60s of inactivity
    AUTO_RESUME = TRUE        -- automatically wake on next query
    INITIALLY_SUSPENDED = TRUE;  -- don't start until needed

# Generate RSA private key (2048-bit)
openssl genrsa 2048 | openssl pkcs8 -topk8 -inform PEM -out rsa_key.p8 -nocrypt

# Extract public key
openssl rsa -in rsa_key.p8 -pubout -out rsa_key.pub

-- Paste the public key content (without header/footer lines)
ALTER USER spark_user
    SET RSA_PUBLIC_KEY='MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A...(base64 content)...';

import boto3, json
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend

# Step 1: Fetch private key from AWS Secrets Manager
sm = boto3.client("secretsmanager")
secret = json.loads(sm.get_secret_value(SecretId="snowflake/spark_user")["SecretString"])
private_key_pem = secret["private_key"]  # PEM string

# Step 2: Load key and serialize to DER bytes (connector needs this format)
private_key = serialization.load_pem_private_key(
    private_key_pem.encode(),
    password=None,
    backend=default_backend()
)
pk_bytes = private_key.private_bytes(
    encoding=serialization.Encoding.DER,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption()
)

# Step 3: Pass to Snowflake connector
snowflake_options = {
    "sfURL"       : secret["sfURL"],
    "sfUser"      : secret["sfUser"],
    "pem_private_key": pk_bytes.decode("latin-1"),  # connector accepts this
    "sfDatabase"  : "ANALYTICS",
    "sfSchema"    : "PUBLIC",
    "sfWarehouse" : "SPARK_WH",
}

def get_snowflake_options(secret_name: str) -> dict:
    """Load Snowflake connection options from AWS Secrets Manager."""
    sm = boto3.client("secretsmanager")
    secret = json.loads(
        sm.get_secret_value(SecretId=secret_name)["SecretString"]
    )
    return {
        "sfURL"      : secret["sfURL"],
        "sfUser"     : secret["sfUser"],
        "sfPassword" : secret["sfPassword"],  # or use pem_private_key
        "sfDatabase" : secret["sfDatabase"],
        "sfSchema"   : secret["sfSchema"],
        "sfWarehouse": secret["sfWarehouse"],
        "sfRole"     : secret.get("sfRole", ""),
    }

# Usage — no secrets in code anywhere
sf_opts = get_snowflake_options("prod/snowflake/spark_user")
df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**sf_opts) \
    .option("dbtable", "ORDERS") \
    .load()

# In Databricks notebooks/jobs — dbutils is available
snowflake_options = {
    "sfURL"      : dbutils.secrets.get("snowflake", "sfURL"),
    "sfUser"     : dbutils.secrets.get("snowflake", "sfUser"),
    "sfPassword" : dbutils.secrets.get("snowflake", "sfPassword"),
    "sfDatabase" : "ANALYTICS",
    "sfSchema"   : "PUBLIC",
    "sfWarehouse": "SPARK_WH",
}
# Secrets never appear in notebook output — Databricks redacts them

net.snowflake:spark-snowflake_2.12:2.12.0-spark_3.4
net.snowflake:snowflake-jdbc:3.14.4

# Step 1: Auto Loader — continuously ingest files to Delta Bronze
df_stream = spark.readStream \
    .format("cloudFiles") \
    .option("cloudFiles.format", "json") \
    .option("cloudFiles.schemaLocation", "/mnt/checkpoints/orders_schema") \
    .load("s3://my-bucket/landing/orders/")

df_stream.writeStream \
    .format("delta") \
    .option("checkpointLocation", "/mnt/checkpoints/orders_bronze") \
    .outputMode("append") \
    .start("/mnt/delta/bronze/orders")

# Step 2 (scheduled job): Read Delta Gold, write to Snowflake
df_gold = spark.read.format("delta").load("/mnt/delta/gold/orders")

df_gold.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_DW") \
    .mode("overwrite") \
    .save()

# Enable Change Data Feed on Delta table
spark.sql("ALTER TABLE delta.`/mnt/delta/gold/orders` SET TBLPROPERTIES (delta.enableChangeDataFeed = true)")

# Read changes since last sync version
last_version = get_last_synced_version()  # from control table

df_changes = spark.read \
    .format("delta") \
    .option("readChangeFeed", "true") \
    .option("startingVersion", last_version) \
    .load("/mnt/delta/gold/orders")

# Filter to only inserts/updates (ignore deletes for simple sync)
df_upserts = df_changes.filter(df_changes["_change_type"].isin(["insert", "update_postimage"]))

# Write to Snowflake with MERGE
df_upserts.drop("_change_type", "_commit_version", "_commit_timestamp") \
    .write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_STAGING") \
    .option("postactions", merge_sql) \
    .mode("overwrite") \
    .save()

import boto3

emr = boto3.client("emr")

response = emr.run_job_flow(
    Name="Snowflake-Spark-Pipeline",
    ReleaseLabel="emr-7.0.0",
    Applications=[{"Name": "Spark"}],
    Instances={
        "InstanceGroups": [
            {"Name": "Master", "InstanceRole": "MASTER",
             "InstanceType": "m5.xlarge", "InstanceCount": 1},
            {"Name": "Core",   "InstanceRole": "CORE",
             "InstanceType": "m5.xlarge", "InstanceCount": 4},
        ],
        "KeepJobFlowAliveWhenNoSteps": False,
    },
    Steps=[{
        "Name": "SnowflakeLoad",
        "ActionOnFailure": "TERMINATE_CLUSTER",
        "HadoopJarStep": {
            "Jar": "command-runner.jar",
            "Args": [
                "spark-submit",
                "--packages",
                "net.snowflake:spark-snowflake_2.12:2.12.0-spark_3.4,"
                "net.snowflake:snowflake-jdbc:3.14.4",
                "s3://my-bucket/scripts/snowflake_load.py",
            ]
        }
    }],
    JobFlowRole="EMR_EC2_DefaultRole",
    ServiceRole="EMR_DefaultRole",
    AutoTerminateAfterLastStep=True,
)

# On EMR: write Parquet to S3 (EMR IAM role handles S3 auth)
df_processed.write \
    .mode("overwrite") \
    .parquet("s3://my-data-lake/snowflake-load/orders/dt=2024-06-01/")

# Connect to Snowflake (credentials from Secrets Manager)
import snowflake.connector

secret = json.loads(
    boto3.client("secretsmanager")
        .get_secret_value(SecretId="prod/snowflake/emr")["SecretString"]
)
conn = snowflake.connector.connect(**secret)
cursor = conn.cursor()

# Trigger COPY INTO from external S3 stage
cursor.execute("""
    COPY INTO ANALYTICS.PUBLIC.ORDERS
    FROM @my_s3_external_stage/orders/dt=2024-06-01/
    FILE_FORMAT = (TYPE = PARQUET, SNAPPY_COMPRESSION = TRUE)
    MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE
    PURGE = FALSE
""")
cursor.close()
conn.close()

from pyspark.sql.functions import from_json, col
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, DoubleType

# Snowflake ARRAY column "TAGS" → arrives as StringType like '["vip","b2b"]'
tags_schema = ArrayType(StringType())
df = df.withColumn("TAGS_PARSED", from_json(col("TAGS"), tags_schema))

# Snowflake OBJECT column "ADDRESS" → arrives as '{"city":"NYC","zip":"10001"}'
addr_schema = StructType([
    StructField("city", StringType(), True),
    StructField("zip",  StringType(), True),
])
df = df.withColumn("ADDRESS_PARSED", from_json(col("ADDRESS"), addr_schema))

# Writing complex types to Snowflake → convert to STRING first
from pyspark.sql.functions import to_json
df_to_write = df.withColumn("TAGS",    to_json(col("tags_array"))) \
                .withColumn("ADDRESS", to_json(col("address_struct")))

from pyspark.sql.functions import from_json, col
from pyspark.sql.types import StructType, StructField, StringType, DoubleType

# Step 1: Read from Kafka
df_kafka = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "broker1:9092") \
    .option("subscribe", "orders-topic") \
    .option("startingOffsets", "latest") \
    .load()

# Step 2: Parse the message value
order_schema = StructType([
    StructField("order_id",   StringType(), True),
    StructField("customer_id", StringType(), True),
    StructField("amount",      DoubleType(), True),
])
df_parsed = df_kafka.select(
    from_json(col("value").cast("string"), order_schema).alias("data")
).select("data.*")

# Step 3: foreachBatch — write each micro-batch to Snowflake
def write_to_snowflake(batch_df, batch_id):
    if batch_df.isEmpty():
        return
    batch_df.write \
        .format("net.snowflake.spark.snowflake") \
        .options(**snowflake_options) \
        .option("dbtable", "ORDERS_REALTIME") \
        .mode("append") \
        .save()
    print(f"Batch {batch_id}: wrote {batch_df.count()} rows to Snowflake")

# Step 4: Start the streaming query
query = df_parsed.writeStream \
    .foreachBatch(write_to_snowflake) \
    .option("checkpointLocation", "s3://bucket/checkpoints/kafka-snowflake") \
    .trigger(processingTime="5 minutes") \  # write to Snowflake every 5 min
    .start()

query.awaitTermination()

# In foreachBatch: write to S3, then notify Snowpipe REST API
import requests

def write_via_snowpipe(batch_df, batch_id):
    if batch_df.isEmpty(): return

    # Write to S3 (path that Snowpipe pipe watches)
    s3_path = f"s3://bucket/snowpipe-landing/orders/batch_{batch_id}/"
    batch_df.write.mode("overwrite").parquet(s3_path)

    # Notify Snowpipe to ingest (using Snowpipe REST API)
    # In practice: S3 event notification → SQS → Snowpipe handles this automatically
    # with AUTO_INGEST = TRUE on the pipe definition

# Snowflake pipe definition (one-time setup)
"""
CREATE PIPE orders_pipe
    AUTO_INGEST = TRUE
    AS
    COPY INTO ORDERS_REALTIME
    FROM @my_s3_stage/snowpipe-landing/orders/
    FILE_FORMAT = (TYPE = PARQUET);
"""
# S3 → SQS notification → Snowpipe auto-triggers when Spark writes new files

# Read Delta Gold layer (fast, metadata-driven)
df_gold = spark.read.format("delta").load("s3://datalake/gold/orders/")

# Write to Snowflake (full refresh of DW layer)
df_gold.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS_DW") \
    .option("preactions", "TRUNCATE TABLE ORDERS_DW") \
    .mode("append") \
    .save()

-- Create external table pointing to Delta Lake Parquet files
CREATE EXTERNAL TABLE ORDERS_DELTA_EXT (
    ORDER_ID     STRING  AS (VALUE:c1::STRING),
    AMOUNT       FLOAT   AS (VALUE:c2::FLOAT),
    ORDER_DATE   DATE    AS (VALUE:c3::DATE)
)
PARTITION BY (ORDER_DATE)
LOCATION = @my_s3_external_stage/gold/orders/
FILE_FORMAT = (TYPE = PARQUET)
AUTO_REFRESH = TRUE;  -- auto-detects new Parquet files via S3 events

-- Snowflake natively reads Iceberg tables on S3
-- Create Iceberg table integration (Snowflake Open Catalog / Polaris)
CREATE CATALOG INTEGRATION iceberg_catalog
    CATALOG_SOURCE = ICEBERG_REST
    TABLE_FORMAT = ICEBERG
    CATALOG_NAMESPACE = 'ANALYTICS'
    REST_CONFIG = (
        CATALOG_URI = 'https://my-iceberg-rest-catalog.com'
    )
    ENABLED = TRUE;

-- Reference existing Iceberg table
CREATE ICEBERG TABLE ORDERS_ICE
    CATALOG = 'iceberg_catalog'
    CATALOG_NAMESPACE = 'default'
    CATALOG_TABLE_NAME = 'orders';

-- Now Snowflake and Spark BOTH query the same Iceberg table on S3
-- No ETL, no copies — true lakehouse convergence
SELECT * FROM ORDERS_ICE WHERE order_date >= '2024-01-01';

# Add to snowflake options
snowflake_options["sfConnectStr"] = "CLIENT_SESSION_KEEP_ALIVE=TRUE"
# OR via preaction
snowflake_options["preactions"] = "ALTER SESSION SET CLIENT_SESSION_KEEP_ALIVE = TRUE"

df = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .option("autopushdown", "off") \   # ← disable pushdown
    .load()

# Then apply Spark-side filter manually
df_filtered = df.filter(df.ORDER_DATE >= "2024-01-01")

from pyspark.sql.functions import col
from pyspark.sql.types import DoubleType, DateType, LongType

# Cast before writing to match Snowflake target schema
df_clean = df \
    .withColumn("AMOUNT",     col("AMOUNT").cast(DoubleType())) \
    .withColumn("ORDER_DATE", col("ORDER_DATE").cast(DateType())) \
    .withColumn("QUANTITY",   col("QUANTITY").cast(LongType()))

df_clean.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**snowflake_options) \
    .option("dbtable", "ORDERS") \
    .mode("append") \
    .save()

-- Grant usage on stage to Spark user's role
GRANT USAGE ON STAGE my_s3_stage TO ROLE DATA_ENGINEER;

-- Check storage integration trust
DESC INTEGRATION s3_int;
-- Note the STORAGE_AWS_IAM_USER_ARN and STORAGE_AWS_EXTERNAL_ID
-- Add these to your S3 bucket IAM role trust policy

-- See credit usage by warehouse
SELECT warehouse_name, SUM(credits_used) AS total_credits
FROM SNOWFLAKE.ACCOUNT_USAGE.WAREHOUSE_METERING_HISTORY
WHERE start_time >= DATEADD('day', -7, CURRENT_TIMESTAMP())
GROUP BY 1
ORDER BY 2 DESC;

-- Enforce auto-suspend on your warehouse
ALTER WAREHOUSE SPARK_WH SET AUTO_SUSPEND = 60;

"""
Production Spark → Snowflake Pipeline Template
Covers: secure auth, incremental read, transform, MERGE write, audit logging
"""
import boto3, json
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, current_timestamp
from datetime import datetime

# 1. Init Spark
spark = SparkSession.builder.appName("snowflake_pipeline").getOrCreate()

# 2. Load credentials securely
secret = json.loads(
    boto3.client("secretsmanager")
        .get_secret_value(SecretId="prod/snowflake/spark")["SecretString"]
)
sf_opts = {
    "sfURL"      : secret["sfURL"],
    "sfUser"     : secret["sfUser"],
    "sfPassword" : secret["sfPassword"],
    "sfDatabase" : "ANALYTICS",
    "sfSchema"   : "PUBLIC",
    "sfWarehouse": "SPARK_WH",
}

# 3. Get watermark
wm_resp = boto3.client("dynamodb").get_item(
    TableName="pipeline_watermarks",
    Key={"pipeline_id": {"S": "orders_sync"}}
)
last_wm = wm_resp["Item"]["last_watermark"]["S"]

# 4. Incremental read from Snowflake
df_raw = spark.read \
    .format("net.snowflake.spark.snowflake") \
    .options(**sf_opts) \
    .option("query", f"SELECT * FROM ORDERS WHERE UPDATED_AT > '{last_wm}'") \
    .load()

# 5. Transform
df_clean = df_raw \
    .withColumn("AMOUNT",       col("AMOUNT").cast("double")) \
    .withColumn("LOAD_TS",      current_timestamp()) \
    .filter(col("ORDER_ID").isNotNull())

# 6. Write to staging + MERGE into Gold
df_clean.write \
    .format("net.snowflake.spark.snowflake") \
    .options(**sf_opts) \
    .option("dbtable", "ORDERS_STAGING") \
    .option("preactions",  "TRUNCATE TABLE ORDERS_STAGING") \
    .option("postactions", """
        MERGE INTO ORDERS_GOLD AS t USING ORDERS_STAGING AS s
        ON t.ORDER_ID = s.ORDER_ID
        WHEN MATCHED THEN UPDATE SET t.AMOUNT=s.AMOUNT, t.LOAD_TS=s.LOAD_TS
        WHEN NOT MATCHED THEN INSERT VALUES (s.ORDER_ID,s.AMOUNT,s.LOAD_TS)
    """) \
    .mode("append") \
    .save()

# 7. Update watermark
new_wm = df_clean.agg({"UPDATED_AT": "max"}).collect()[0][0]
boto3.client("dynamodb").update_item(
    TableName="pipeline_watermarks",
    Key={"pipeline_id": {"S": "orders_sync"}},
    UpdateExpression="SET last_watermark = :wm",
    ExpressionAttributeValues={":wm": {"S": str(new_wm)}}
)
print(f"Pipeline complete. New watermark: {new_wm}")