# /// script
# requires-python = ">=3.11"
# dependencies = ["pyspark==4.1.0"]
# ///
"""Create, change, and drop an Iceberg table in Unity Catalog from Spark, through Iceberg REST.

Start the docs' local server from `uc-docs-env/uc-aws`:

    docker compose up -d --wait

Then run this script, with `_setup.py` from the same page beside it:

    AWS_ENDPOINT_URL=http://localhost:9000 AWS_ALLOW_HTTP=true uv run spark_iceberg.py

Needs Java 17 and Maven access on first run. Behind a Maven mirror, export
PYSPARK_SUBMIT_ARGS="--repositories <mirror-url> pyspark-shell" first.
"""

from pyspark.sql import SparkSession

ICEBERG_SPARK = "org.apache.iceberg:iceberg-spark-runtime-4.1_2.13:1.11.0"
ICEBERG_AWS = "org.apache.iceberg:iceberg-aws-bundle:1.11.0"

spark = (
    SparkSession.builder.appName("uc-iceberg")
    .config("spark.jars.packages", f"{ICEBERG_SPARK},{ICEBERG_AWS}")
    .config(
        "spark.sql.extensions",
        "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions",
    )
    .config("spark.sql.catalog.retail", "org.apache.iceberg.spark.SparkCatalog")
    .config("spark.sql.catalog.retail.type", "rest")
    .config(
        "spark.sql.catalog.retail.uri",
        "http://localhost:8080/api/2.1/unity-catalog/iceberg",
    )
    .config("spark.sql.catalog.retail.warehouse", "retail")
    .config("spark.sql.catalog.retail.token", "not-used")
    .config("spark.sql.defaultCatalog", "retail")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

from _setup import setup_retail, table_info  # noqa: E402

setup_retail()

spark.sql("CREATE NAMESPACE staging")
spark.sql("SHOW NAMESPACES").show()
assert {r.namespace for r in spark.sql("SHOW NAMESPACES").collect()} == {
    "sales",
    "staging",
}

spark.sql("""
    CREATE TABLE sales.orders (
        order_id BIGINT,
        customer STRING,
        amount DOUBLE
    ) USING iceberg
""")
info = table_info("retail.sales.orders")
assert (info["data_source_format"], info["table_type"]) == ("ICEBERG", "MANAGED"), info
assert info["storage_location"].startswith("s3://uc-docs/retail/managed/"), info

spark.sql("""
    INSERT INTO sales.orders VALUES
        (1, 'ada', 9.5),
        (2, 'grace', 20.0),
        (3, 'linus', 3.25)
""")

spark.sql("SELECT * FROM sales.orders WHERE amount > 5 ORDER BY order_id").show()
assert spark.table("sales.orders").count() == 3

spark.sql("UPDATE sales.orders SET amount = 21.0 WHERE order_id = 2")
spark.sql("DELETE FROM sales.orders WHERE customer = 'linus'")
spark.sql("""
    MERGE INTO sales.orders AS t
    USING (SELECT 4 AS order_id, 'barbara' AS customer, 12.0 AS amount) AS s
    ON t.order_id = s.order_id
    WHEN MATCHED THEN UPDATE SET *
    WHEN NOT MATCHED THEN INSERT *
""")
rows = spark.sql("SELECT * FROM sales.orders ORDER BY order_id").collect()
assert [tuple(r) for r in rows] == [
    (1, "ada", 9.5),
    (2, "grace", 21.0),
    (4, "barbara", 12.0),
], rows

spark.sql("ALTER TABLE sales.orders RENAME TO sales.order_history")
assert [r.tableName for r in spark.sql("SHOW TABLES IN sales").collect()] == [
    "order_history"
]

spark.sql("DROP TABLE sales.order_history")
spark.sql("DROP NAMESPACE staging")
assert spark.sql("SHOW TABLES IN sales").collect() == []

spark.stop()