import time, json, traceback
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

# Keep it minimal
spark.conf.set("spark.sql.adaptive.enabled", "true")

# -------- Conn --------
host     = "dz-olake-benchmark-postgres.postgres.database.azure.com"
port     = "5432"
database = "database_name"
user     = "postgres_username"
password = "password"  # e.g., dbutils.secrets.get("scope","key")

jdbc_url = (
    f"jdbc:postgresql://{host}:{port}/{database}"
    "?sslmode=require&connectTimeout=30&socketTimeout=28800&queryTimeout=27000"
    "&tcpKeepAlive=true&keepalives_idle=120&keepalives_interval=30&keepalives_count=5"
)

JDBC_PROPS = {
  "driver": "org.postgresql.Driver",
  "user": user,
  "password": password,
  "fetchsize": "200000"   # single-stream batches
}

SRC_SCHEMA = "tpch"
TARGET_SCHEMA = "datazip.testing3"   # adjust if needed
TABLES = ['lineitem','orders','partsupp','customer','part','supplier','nation','region']

def log(msg):
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{ts}] {msg}", flush=True)

def rows_bytes_from_history(fq):
    try:
        h = (spark.sql(f"DESCRIBE HISTORY {fq}")
               .where("operation in ('WRITE','MERGE','CREATE TABLE AS SELECT')")
               .orderBy("timestamp", ascending=False)
               .limit(1)
               .collect()[0]["operationMetrics"])
        def g(k, default="0"): return int(h.get(k, default))
        num_rows = g("numOutputRows", h.get("numInsertedRows","0"))
        bytes_written = g("bytesWritten")
        files_added = g("numFiles", h.get("numAddedFiles","0"))
        return num_rows, bytes_written, files_added
    except Exception:
        return 0, 0, 0

def ensure_table_with_props(fq, src_table):
    # Build empty DF (0 rows) to carry schema
    empty_df = (spark.read.format("jdbc")
                  .option("url", jdbc_url)
                  .option("dbtable", f"(SELECT * FROM {SRC_SCHEMA}.{src_table} WHERE 1=0) s")
                  .options(**JDBC_PROPS)
                  .load())
    view = f"v_{src_table}_empty"
    empty_df.createOrReplaceTempView(view)

    # Create once with table-level Auto Optimize before first real write
    spark.sql(f"""
      CREATE TABLE IF NOT EXISTS {fq}
      USING DELTA
      TBLPROPERTIES (
        'delta.autoOptimize.optimizeWrite'='true',
        'delta.autoOptimize.autoCompact'='true'
      )
      AS SELECT * FROM {view} WHERE 1=0
    """)
    # Ensure props in case table existed without them
    spark.sql(f"""
      ALTER TABLE {fq} SET TBLPROPERTIES (
        'delta.autoOptimize.optimizeWrite'='true',
        'delta.autoOptimize.autoCompact'='true'
      )
    """)
    spark.catalog.dropTempView(view)

def truncate_table(fq):
    spark.sql(f"TRUNCATE TABLE {fq}")

def read_whole_table(src_table):
    t0 = time.time()
    df = (spark.read.format("jdbc")
            .option("url", jdbc_url)
            .option("dbtable", f"{SRC_SCHEMA}.{src_table}")
            .options(**JDBC_PROPS)
            .load())
    return df, time.time() - t0

# -------- MAIN --------
overall_t0 = time.time()
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {TARGET_SCHEMA}")

log("="*80)
log("🚀 SIMPLE TPCH COPY (single JDBC stream, fetchsize=200000, create-first with table-level Auto Optimize)")
log("="*80)

summary = []

for t in TABLES:
    fq = f"{TARGET_SCHEMA}.{t}"
    start = time.time()
    log(f"\n========== {t} ==========")
    try:
        # 1) Ensure empty table exists with desired properties, then truncate for a fresh load
        ensure_table_with_props(fq, t)
        truncate_table(fq)

        # 2) Single-stream READ
        df, read_sec = read_whole_table(t)
        log(f"Read {t} finished in {read_sec:.2f}s")

        # 3) Single WRITE (append into truncated table)
        w0 = time.time()
        (df.write
            .format("delta")
            .mode("append")
            .saveAsTable(fq))
        write_sec = time.time() - w0

        rows, bytes_written, files = rows_bytes_from_history(fq)
        total = time.time() - start
        log(f"✅ {t} written in {write_sec:.2f}s | total {total:.2f}s")
        log(f"   rows={rows:,} | bytes={bytes_written:,} | files={files}")

        summary.append({
            "table": t,
            "read_sec": round(read_sec,2),
            "write_sec": round(write_sec,2),
            "total_sec": round(total,2),
            "rows": rows,
            "bytes": bytes_written,
            "files": files
        })

    except Exception as e:
        log(f"❌ FAILED {t}: {e}")
        log(traceback.format_exc())
        summary.append({"table": t, "error": str(e)})

overall = time.time() - overall_t0
log("\n" + "="*80)
log("🎉 PIPELINE COMPLETED")
log("="*80)
log(f"Total elapsed: {overall/60:.2f} min")
log("\n📋 FINAL SUMMARY:")
log(json.dumps(summary, indent=2))

import time
import logging

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger('tpch_benchmark')

def configure_spark_settings():
    """Configure Spark settings to optimize query execution"""
    spark.conf.set("spark.sql.adaptive.enabled", "true")
    spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
    spark.conf.set("spark.sql.shuffle.partitions", "64")
    logger.info("Spark configurations set: Adaptive Query Execution enabled, Coalesce Partitions enabled, shuffle partitions set to 64")

def run_tpch_query(query_num, sql_query):
    """Run a single TPC-H query with timing and logging"""
    query_name = f"TPC-H Query {query_num}"
    logger.info(f"Starting {query_name}")
    
    start_time = time.perf_counter()

    try:
        # Execute the exact SQL query as provided
        if query_num == 15 and isinstance(sql_query, list):
            for i, sql_stmt in enumerate(sql_query):
                logger.info(f"{query_name} - Executing statement {i+1}/3")
                if i == 1:
                    # SELECT statement: return/display results to force execution
                    df = spark.sql(sql_stmt)
                    df.show()
                else:
                    # CREATE VIEW and DROP VIEW statements
                    spark.sql(sql_stmt)
        else:
            # All other queries are single-statement
            df = spark.sql(sql_query)
            df.show()
        
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        
        logger.info(f"{query_name} completed successfully in {execution_time:.2f} seconds")
        
        return {
            'query_number': query_num,
            'time_seconds': execution_time,
            'status': 'SUCCESS'
        }
        
    
    except Exception as e:
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        logger.error(f"{query_name} failed after {execution_time:.2f} seconds - Error: {str(e)}")
        
        return {
            'query_number': query_num,
            'time_seconds': execution_time,
            'status': 'FAILED',
            'error': str(e)
        }

# All 22 TPC-H Queries exactly as provided in your file
tpch_queries = {
    1: """SELECT
l_returnflag,
l_linestatus,
SUM(l_quantity) AS sum_qty,
SUM(l_extendedprice) AS sum_base_price,
SUM(l_extendedprice * (1 - l_discount)) AS sum_disc_price,
SUM(l_extendedprice * (1 - l_discount) * (1 + l_tax)) AS sum_charge,
AVG(l_quantity) AS avg_qty,
AVG(l_extendedprice) AS avg_price,
AVG(l_discount) AS avg_disc,
COUNT(*) AS count_order
FROM datazip.testing3.lineitem
WHERE l_shipdate <= DATE '1998-12-01' - INTERVAL '90' DAY
GROUP BY l_returnflag, l_linestatus
ORDER BY l_returnflag, l_linestatus""",

    2: """SELECT
s_acctbal,
s_name,
n_name,
p_partkey,
p_mfgr,
s_address,
s_phone,
s_comment
FROM
datazip.testing3.part,
datazip.testing3.supplier,
datazip.testing3.partsupp,
datazip.testing3.nation,
datazip.testing3.region
WHERE
p_partkey = ps_partkey
AND s_suppkey = ps_suppkey
AND p_size = 15
AND p_type LIKE '%BRASS'
AND s_nationkey = n_nationkey
AND n_regionkey = r_regionkey
AND r_name = 'EUROPE'
AND ps_supplycost = (
SELECT MIN(ps_supplycost)
FROM datazip.testing3.partsupp, datazip.testing3.supplier, datazip.testing3.nation, datazip.testing3.region
WHERE p_partkey = ps_partkey
AND s_suppkey = ps_suppkey
AND s_nationkey = n_nationkey
AND n_regionkey = r_regionkey
AND r_name = 'EUROPE'
)
ORDER BY
s_acctbal DESC,
n_name,
s_name,
p_partkey
LIMIT 100""",

    3: """SELECT
l_orderkey,
SUM(l_extendedprice * (1 - l_discount)) AS revenue,
o_orderdate,
o_shippriority
FROM
datazip.testing3.customer,
datazip.testing3.orders,
datazip.testing3.lineitem
WHERE
c_mktsegment = 'BUILDING'
AND c_custkey = o_custkey
AND l_orderkey = o_orderkey
AND o_orderdate < DATE '1995-03-15'
AND l_shipdate > DATE '1995-03-15'
GROUP BY
l_orderkey,
o_orderdate,
o_shippriority
ORDER BY
revenue DESC,
o_orderdate
LIMIT 10""",

    4: """SELECT
o_orderpriority,
COUNT(*) AS order_count
FROM datazip.testing3.orders
WHERE
o_orderdate >= DATE '1993-07-01'
AND o_orderdate < DATE '1993-07-01' + INTERVAL '3' MONTH
AND EXISTS (
SELECT *
FROM datazip.testing3.lineitem
WHERE l_orderkey = o_orderkey
AND l_commitdate < l_receiptdate
)
GROUP BY o_orderpriority
ORDER BY o_orderpriority""",

    5: """SELECT
n_name,
SUM(l_extendedprice * (1 - l_discount)) AS revenue
FROM
datazip.testing3.customer,
datazip.testing3.orders,
datazip.testing3.lineitem,
datazip.testing3.supplier,
datazip.testing3.nation,
datazip.testing3.region
WHERE
c_custkey = o_custkey
AND l_orderkey = o_orderkey
AND l_suppkey = s_suppkey
AND c_nationkey = s_nationkey
AND s_nationkey = n_nationkey
AND n_regionkey = r_regionkey
AND r_name = 'ASIA'
AND o_orderdate >= DATE '1994-01-01'
AND o_orderdate < DATE '1994-01-01' + INTERVAL '1' YEAR
GROUP BY n_name
ORDER BY revenue DESC""",

    6: """SELECT SUM(l_extendedprice * l_discount) AS revenue
FROM datazip.testing3.lineitem
WHERE
l_shipdate >= DATE '1994-01-01'
AND l_shipdate < DATE '1994-01-01' + INTERVAL '1' YEAR
AND l_discount BETWEEN 0.06 - 0.01 AND 0.06 + 0.01
AND l_quantity < 24""",

    7: """SELECT
supp_nation,
cust_nation,
l_year,
SUM(volume) AS revenue
FROM (
SELECT
n1.n_name AS supp_nation,
n2.n_name AS cust_nation,
EXTRACT(YEAR FROM l_shipdate) AS l_year,
l_extendedprice * (1 - l_discount) AS volume
FROM
datazip.testing3.supplier,
datazip.testing3.lineitem,
datazip.testing3.orders,
datazip.testing3.customer,
datazip.testing3.nation n1,
datazip.testing3.nation n2
WHERE
s_suppkey = l_suppkey
AND o_orderkey = l_orderkey
AND c_custkey = o_custkey
AND s_nationkey = n1.n_nationkey
AND c_nationkey = n2.n_nationkey
AND ((n1.n_name = 'FRANCE' AND n2.n_name = 'GERMANY')
OR (n1.n_name = 'GERMANY' AND n2.n_name = 'FRANCE'))
AND l_shipdate BETWEEN DATE '1995-01-01' AND DATE '1996-12-31'
) AS shipping
GROUP BY supp_nation, cust_nation, l_year
ORDER BY supp_nation, cust_nation, l_year""",

    8: """SELECT
o_year,
SUM(CASE WHEN nation = 'BRAZIL' THEN volume ELSE 0 END) / SUM(volume) AS mkt_share
FROM (
SELECT
EXTRACT(YEAR FROM o_orderdate) AS o_year,
l_extendedprice * (1 - l_discount) AS volume,
n2.n_name AS nation
FROM
datazip.testing3.part,
datazip.testing3.supplier,
datazip.testing3.lineitem,
datazip.testing3.orders,
datazip.testing3.customer,
datazip.testing3.nation n1,
datazip.testing3.nation n2,
datazip.testing3.region
WHERE
p_partkey = l_partkey
AND s_suppkey = l_suppkey
AND l_orderkey = o_orderkey
AND o_custkey = c_custkey
AND c_nationkey = n1.n_nationkey
AND n1.n_regionkey = r_regionkey
AND r_name = 'AMERICA'
AND s_nationkey = n2.n_nationkey
AND o_orderdate BETWEEN DATE '1995-01-01' AND DATE '1996-12-31'
AND p_type = 'ECONOMY ANODIZED STEEL'
) AS all_nations
GROUP BY o_year
ORDER BY o_year""",

    9: """SELECT
nation,
o_year,
SUM(amount) AS sum_profit
FROM (
SELECT
n_name AS nation,
EXTRACT(YEAR FROM o_orderdate) AS o_year,
l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity AS amount
FROM
datazip.testing3.part,
datazip.testing3.supplier,
datazip.testing3.lineitem,
datazip.testing3.partsupp,
datazip.testing3.orders,
datazip.testing3.nation
WHERE
s_suppkey = l_suppkey
AND ps_suppkey = l_suppkey
AND ps_partkey = l_partkey
AND p_partkey = l_partkey
AND o_orderkey = l_orderkey
AND s_nationkey = n_nationkey
AND p_name LIKE '%green%'
) AS profit
GROUP BY nation, o_year
ORDER BY nation, o_year DESC""",

    10: """SELECT
c_custkey,
c_name,
SUM(l_extendedprice * (1 - l_discount)) AS revenue,
c_acctbal,
n_name,
c_address,
c_phone,
c_comment
FROM
datazip.testing3.customer,
datazip.testing3.orders,
datazip.testing3.lineitem,
datazip.testing3.nation
WHERE
c_custkey = o_custkey
AND l_orderkey = o_orderkey
AND o_orderdate >= DATE '1993-10-01'
AND o_orderdate < DATE '1993-10-01' + INTERVAL '3' MONTH
AND l_returnflag = 'R'
AND c_nationkey = n_nationkey
GROUP BY
c_custkey,
c_name,
c_acctbal,
c_phone,
n_name,
c_address,
c_comment
ORDER BY revenue DESC
LIMIT 20""",

    11: """SELECT
ps_partkey,
SUM(ps_supplycost * ps_availqty) AS value
FROM
datazip.testing3.partsupp,
datazip.testing3.supplier,
datazip.testing3.nation
WHERE
ps_suppkey = s_suppkey
AND s_nationkey = n_nationkey
AND n_name = 'GERMANY'
GROUP BY ps_partkey
HAVING SUM(ps_supplycost * ps_availqty) > (
SELECT SUM(ps_supplycost * ps_availqty) * 0.0001
FROM
datazip.testing3.partsupp,
datazip.testing3.supplier,
datazip.testing3.nation
WHERE
ps_suppkey = s_suppkey
AND s_nationkey = n_nationkey
AND n_name = 'GERMANY'
)
ORDER BY value DESC""",

    12: """SELECT
l_shipmode,
SUM(CASE
WHEN o_orderpriority = '1-URGENT'
OR o_orderpriority = '2-HIGH'
THEN 1
ELSE 0
END) AS high_line_count,
SUM(CASE
WHEN o_orderpriority <> '1-URGENT'
AND o_orderpriority <> '2-HIGH'
THEN 1
ELSE 0
END) AS low_line_count
FROM
datazip.testing3.orders,
datazip.testing3.lineitem
WHERE
o_orderkey = l_orderkey
AND l_shipmode IN ('MAIL', 'SHIP')
AND l_commitdate < l_receiptdate
AND l_shipdate < l_commitdate
AND l_receiptdate >= DATE '1994-01-01'
AND l_receiptdate < DATE '1994-01-01' + INTERVAL '1' YEAR
GROUP BY l_shipmode
ORDER BY l_shipmode""",

    13: """SELECT
c_count,
COUNT(*) AS custdist
FROM (
SELECT
c_custkey,
COUNT(o_orderkey) AS c_count
FROM
datazip.testing3.customer LEFT OUTER JOIN datazip.testing3.orders ON
c_custkey = o_custkey
AND o_comment NOT LIKE '%special%requests%'
GROUP BY c_custkey
) c_orders
GROUP BY c_count
ORDER BY custdist DESC, c_count DESC""",

    14: """SELECT
100.00 * SUM(CASE
WHEN p_type LIKE 'PROMO%'
THEN l_extendedprice * (1 - l_discount)
ELSE 0
END) / SUM(l_extendedprice * (1 - l_discount)) AS promo_revenue
FROM
datazip.testing3.lineitem,
datazip.testing3.part
WHERE
l_partkey = p_partkey
AND l_shipdate >= DATE '1995-09-01'
AND l_shipdate < DATE '1995-09-01' + INTERVAL '1' MONTH""",

    15: [
        """CREATE OR REPLACE TEMPORARY VIEW revenue0 AS
        SELECT
          l_suppkey AS supplier_no,
          SUM(l_extendedprice * (1 - l_discount)) AS total_revenue
        FROM datazip.testing3.lineitem
        WHERE
          l_shipdate >= DATE '1996-01-01'
          AND l_shipdate < DATE '1996-01-01' + INTERVAL '3' MONTH
        GROUP BY l_suppkey""",

        """SELECT
          s_suppkey,
          s_name,
          s_address,
          s_phone,
          total_revenue
        FROM
          datazip.testing3.supplier,
          revenue0
        WHERE
          s_suppkey = supplier_no
          AND total_revenue = (
            SELECT MAX(total_revenue) FROM revenue0
          )
        ORDER BY s_suppkey""",

        """DROP VIEW revenue0"""
    ]
,

    16: """SELECT
p_brand,
p_type,
p_size,
COUNT(DISTINCT ps_suppkey) AS supplier_cnt
FROM
datazip.testing3.partsupp,
datazip.testing3.part
WHERE
p_partkey = ps_partkey
AND p_brand <> 'Brand#45'
AND p_type NOT LIKE 'MEDIUM POLISHED%'
AND p_size IN (49, 14, 23, 45, 19, 3, 36, 9)
AND ps_suppkey NOT IN (
SELECT s_suppkey
FROM datazip.testing3.supplier
WHERE s_comment LIKE '%Customer%Complaints%'
)
GROUP BY p_brand, p_type, p_size
ORDER BY supplier_cnt DESC, p_brand, p_type, p_size""",

    17: """SELECT SUM(l_extendedprice) / 7.0 AS avg_yearly
FROM
datazip.testing3.lineitem,
datazip.testing3.part
WHERE
p_partkey = l_partkey
AND p_brand = 'Brand#23'
AND p_container = 'MED BOX'
AND l_quantity < (
SELECT 0.2 * AVG(l_quantity)
FROM datazip.testing3.lineitem
WHERE l_partkey = p_partkey
)""",

    18: """SELECT
c_name,
c_custkey,
o_orderkey,
o_orderdate,
o_totalprice,
SUM(l_quantity) AS sum_qty
FROM
datazip.testing3.customer,
datazip.testing3.orders,
datazip.testing3.lineitem
WHERE
o_orderkey IN (
SELECT l_orderkey
FROM datazip.testing3.lineitem
GROUP BY l_orderkey
HAVING SUM(l_quantity) > 300
)
AND c_custkey = o_custkey
AND o_orderkey = l_orderkey
GROUP BY
c_name,
c_custkey,
o_orderkey,
o_orderdate,
o_totalprice
ORDER BY o_totalprice DESC, o_orderdate
LIMIT 100""",

    19: """SELECT SUM(l_extendedprice * (1 - l_discount)) AS revenue
FROM
datazip.testing3.lineitem,
datazip.testing3.part
WHERE
(
p_partkey = l_partkey
AND p_brand = 'Brand#12'
AND p_container IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')
AND l_quantity >= 1 AND l_quantity <= 1 + 10
AND p_size BETWEEN 1 AND 5
AND l_shipmode IN ('AIR', 'AIR REG')
AND l_shipinstruct = 'DELIVER IN PERSON'
)
OR
(
p_partkey = l_partkey
AND p_brand = 'Brand#23'
AND p_container IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')
AND l_quantity >= 10 AND l_quantity <= 10 + 10
AND p_size BETWEEN 1 AND 10
AND l_shipmode IN ('AIR', 'AIR REG')
AND l_shipinstruct = 'DELIVER IN PERSON'
)
OR
(
p_partkey = l_partkey
AND p_brand = 'Brand#34'
AND p_container IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')
AND l_quantity >= 20 AND l_quantity <= 20 + 10
AND p_size BETWEEN 1 AND 15
AND l_shipmode IN ('AIR', 'AIR REG')
AND l_shipinstruct = 'DELIVER IN PERSON'
)""",

    20: """SELECT
s_name,
s_address
FROM
datazip.testing3.supplier,
datazip.testing3.nation
WHERE
s_suppkey IN (
SELECT ps_suppkey
FROM datazip.testing3.partsupp
WHERE ps_partkey IN (
SELECT p_partkey
FROM datazip.testing3.part
WHERE p_name LIKE 'forest%'
)
AND ps_availqty > (
SELECT 0.5 * SUM(l_quantity)
FROM datazip.testing3.lineitem
WHERE l_partkey = ps_partkey
AND l_suppkey = ps_suppkey
AND l_shipdate >= DATE '1994-01-01'
AND l_shipdate < DATE '1994-01-01' + INTERVAL '1' YEAR
)
)
AND s_nationkey = n_nationkey
AND n_name = 'CANADA'
ORDER BY s_name""",

    21: """SELECT
s_name,
COUNT(*) AS numwait
FROM
datazip.testing3.supplier,
datazip.testing3.lineitem l1,
datazip.testing3.orders,
datazip.testing3.nation
WHERE
s_suppkey = l1.l_suppkey
AND o_orderkey = l1.l_orderkey
AND o_orderstatus = 'F'
AND l1.l_receiptdate > l1.l_commitdate
AND EXISTS (
SELECT *
FROM datazip.testing3.lineitem l2
WHERE l2.l_orderkey = l1.l_orderkey
AND l2.l_suppkey <> l1.l_suppkey
)
AND NOT EXISTS (
SELECT *
FROM datazip.testing3.lineitem l3
WHERE l3.l_orderkey = l1.l_orderkey
AND l3.l_suppkey <> l1.l_suppkey
AND l3.l_receiptdate > l3.l_commitdate
)
AND s_nationkey = n_nationkey
AND n_name = 'SAUDI ARABIA'
GROUP BY s_name
ORDER BY numwait DESC, s_name
LIMIT 100""",

    22: """SELECT
cntrycode,
COUNT(*) AS numcust,
SUM(c_acctbal) AS totacctbal
FROM (
SELECT
SUBSTRING(c_phone, 1, 2) AS cntrycode,
c_acctbal
FROM datazip.testing3.customer
WHERE SUBSTRING(c_phone, 1, 2) IN ('13', '31', '23', '29', '30', '18', '17')
AND c_acctbal > (
SELECT AVG(c_acctbal)
FROM datazip.testing3.customer
WHERE c_acctbal > 0.00
AND SUBSTRING(c_phone, 1, 2) IN ('13', '31', '23', '29', '30', '18', '17')
)
AND NOT EXISTS (
SELECT *
FROM datazip.testing3.orders
WHERE o_custkey = c_custkey
)
) AS custsale
GROUP BY cntrycode
ORDER BY cntrycode"""
}

# Run all 22 TPC-H queries with timing and logging
def run_all_tpch_queries():
    all_results = []
    total_start_time = time.perf_counter()
    
    logger.info("Starting complete TPC-H benchmark run (22 queries)")
    
    for query_num in range(1, 23):  # 1 to 22
        if query_num in tpch_queries:
            result = run_tpch_query(query_num, tpch_queries[query_num])
            all_results.append(result)
            
            # Small pause between queries
            time.sleep(0.5)
    
    total_end_time = time.perf_counter()
    total_time = total_end_time - total_start_time
    
    logger.info(f"Complete TPC-H benchmark finished in {total_time:.2f} seconds")
    
    # Print summary
    print("\n" + "="*60)
    print("TPC-H BENCHMARK RESULTS SUMMARY")
    print("="*60)
    
    successful_queries = 0
    failed_queries = 0
    total_query_time = 0
    
    for result in all_results:
        status_symbol = "✓" if result['status'] == 'SUCCESS' else "✗"
        print(f"{status_symbol} Query {result['query_number']:2d}: {result['time_seconds']:8.2f}s - {result['status']}")
        
        if result['status'] == 'SUCCESS':
            successful_queries += 1
            total_query_time += result['time_seconds']
        else:
            failed_queries += 1
    
    print("="*60)
    print(f"Successful queries: {successful_queries}/22")
    print(f"Failed queries: {failed_queries}/22")
    print(f"Total query execution time: {total_query_time:.2f} seconds")
    print(f"Total benchmark time: {total_time:.2f} seconds")
    print("="*60)
    
    return all_results

# Execute the benchmark
if __name__ == "__main__":
    benchmark_results = run_all_tpch_queries()

import time
import logging
import argparse
import os
from pyspark.sql import SparkSession
from typing import Optional, Union, List, Dict, Any

# Logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("tpch_iceberg_benchmark")

def build_spark_session(app_name: str, catalog: str, warehouse: Optional[str], region: Optional[str]) -> SparkSession:
    builder = (
        SparkSession.builder
        .appName(app_name)
        .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
        .config(f"spark.sql.catalog.{catalog}", "org.apache.iceberg.spark.SparkCatalog")
        .config(f"spark.sql.catalog.{catalog}.catalog-impl", "org.apache.iceberg.aws.glue.GlueCatalog")
        .config(f"spark.sql.catalog.{catalog}.io-impl", "org.apache.iceberg.aws.s3.S3FileIO")
        .config("spark.sql.defaultCatalog", catalog)
    )
    if warehouse:
        builder = builder.config(f"spark.sql.catalog.{catalog}.warehouse", warehouse)
    if region:
        builder = builder.config(f"spark.sql.catalog.{catalog}.client.region", region)

    # Reasonable Spark defaults for scans/joins
    builder = (
        builder
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        .config("spark.sql.broadcastTimeout", "1800")
        .config("spark.network.timeout", "300s")
        .config("spark.executor.heartbeatInterval", "60s")
    )
    spark = builder.getOrCreate()
    return spark

def configure_spark_settings(spark: SparkSession, shuffle_partitions: int) -> None:
    spark.conf.set("spark.sql.shuffle.partitions", str(shuffle_partitions))
    logger.info(
        f"Spark configured: AQE ON, Coalesce ON, shuffle partitions={shuffle_partitions}"
    )

def run_tpch_query(spark: SparkSession, query_num: int, sql_query) -> Dict[str, Any]:
    query_name = f"TPC-H Query {query_num}"
    logger.info(f"Starting {query_name}")
    start_time = time.perf_counter()

    try:
        if query_num == 15 and isinstance(sql_query, list):
            for i, sql_stmt in enumerate(sql_query):
                logger.info(f"{query_name} - Executing statement {i+1}/3")
                if i == 1:
                    df = spark.sql(sql_stmt)
                    df.show()
                else:
                    spark.sql(sql_stmt)
        else:
            df = spark.sql(sql_query)
            df.show()

        end_time = time.perf_counter()
        execution_time = end_time - start_time
        logger.info(f"{query_name} completed successfully in {execution_time:.2f} seconds")
        return {"query_number": query_num, "time_seconds": execution_time, "status": "SUCCESS"}
    except Exception as e:
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        logger.error(f"{query_name} failed after {execution_time:.2f} seconds - Error: {str(e)}")
        return {
            "query_number": query_num,
            "time_seconds": execution_time,
            "status": "FAILED",
            "error": str(e),
        }

def build_queries(catalog: str, database: str) -> Dict[int, Union[str, List[str]]]:
    def tbl(name: str) -> str:
        # backtick database and table because the db contains underscores and to be safe for any special chars
        return f"{catalog}.`{database}`.`{name}`"

    return {
        1: f"""SELECT
l_returnflag,
l_linestatus,
SUM(l_quantity) AS sum_qty,
SUM(l_extendedprice) AS sum_base_price,
SUM(l_extendedprice * (1 - l_discount)) AS sum_disc_price,
SUM(l_extendedprice * (1 - l_discount) * (1 + l_tax)) AS sum_charge,
AVG(l_quantity) AS avg_qty,
AVG(l_extendedprice) AS avg_price,
AVG(l_discount) AS avg_disc,
COUNT(*) AS count_order
FROM {tbl('lineitem')}
WHERE l_shipdate <= DATE '1998-12-01' - INTERVAL '90' DAY
GROUP BY l_returnflag, l_linestatus
ORDER BY l_returnflag, l_linestatus""",

        2: f"""SELECT
s_acctbal,
s_name,
n_name,
p_partkey,
p_mfgr,
s_address,
s_phone,
s_comment
FROM
{tbl('part')},
{tbl('supplier')},
{tbl('partsupp')},
{tbl('nation')},
{tbl('region')}
WHERE
p_partkey = ps_partkey
AND s_suppkey = ps_suppkey
AND p_size = 15
AND p_type LIKE '%BRASS'
AND s_nationkey = n_nationkey
AND n_regionkey = r_regionkey
AND r_name = 'EUROPE'
AND ps_supplycost = (
SELECT MIN(ps_supplycost)
FROM {tbl('partsupp')}, {tbl('supplier')}, {tbl('nation')}, {tbl('region')}
WHERE p_partkey = ps_partkey
AND s_suppkey = ps_suppkey
AND s_nationkey = n_nationkey
AND n_regionkey = r_regionkey
AND r_name = 'EUROPE'
)
ORDER BY
s_acctbal DESC,
n_name,
s_name,
p_partkey
LIMIT 100""",

        3: f"""SELECT
l_orderkey,
SUM(l_extendedprice * (1 - l_discount)) AS revenue,
o_orderdate,
o_shippriority
FROM
{tbl('customer')},
{tbl('orders')},
{tbl('lineitem')}
WHERE
c_mktsegment = 'BUILDING'
AND c_custkey = o_custkey
AND l_orderkey = o_orderkey
AND o_orderdate < DATE '1995-03-15'
AND l_shipdate > DATE '1995-03-15'
GROUP BY
l_orderkey,
o_orderdate,
o_shippriority
ORDER BY
revenue DESC,
o_orderdate
LIMIT 10""",

        4: f"""SELECT
o_orderpriority,
COUNT(*) AS order_count
FROM {tbl('orders')}
WHERE
o_orderdate >= DATE '1993-07-01'
AND o_orderdate < DATE '1993-07-01' + INTERVAL '3' MONTH
AND EXISTS (
SELECT *
FROM {tbl('lineitem')}
WHERE l_orderkey = o_orderkey
AND l_commitdate < l_receiptdate
)
GROUP BY o_orderpriority
ORDER BY o_orderpriority""",

        5: f"""SELECT
n_name,
SUM(l_extendedprice * (1 - l_discount)) AS revenue
FROM
{tbl('customer')},
{tbl('orders')},
{tbl('lineitem')},
{tbl('supplier')},
{tbl('nation')},
{tbl('region')}
WHERE
c_custkey = o_custkey
AND l_orderkey = o_orderkey
AND l_suppkey = s_suppkey
AND c_nationkey = s_nationkey
AND s_nationkey = n_nationkey
AND n_regionkey = r_regionkey
AND r_name = 'ASIA'
AND o_orderdate >= DATE '1994-01-01'
AND o_orderdate < DATE '1994-01-01' + INTERVAL '1' YEAR
GROUP BY n_name
ORDER BY revenue DESC""",

        6: f"""SELECT SUM(l_extendedprice * l_discount) AS revenue
FROM {tbl('lineitem')}
WHERE
l_shipdate >= DATE '1994-01-01'
AND l_shipdate < DATE '1994-01-01' + INTERVAL '1' YEAR
AND l_discount BETWEEN 0.06 - 0.01 AND 0.06 + 0.01
AND l_quantity < 24""",

        7: f"""SELECT
supp_nation,
cust_nation,
l_year,
SUM(volume) AS revenue
FROM (
SELECT
n1.n_name AS supp_nation,
n2.n_name AS cust_nation,
EXTRACT(YEAR FROM l_shipdate) AS l_year,
l_extendedprice * (1 - l_discount) AS volume
FROM
{tbl('supplier')},
{tbl('lineitem')},
{tbl('orders')},
{tbl('customer')},
{tbl('nation')} n1,
{tbl('nation')} n2
WHERE
s_suppkey = l_suppkey
AND o_orderkey = l_orderkey
AND c_custkey = o_custkey
AND s_nationkey = n1.n_nationkey
AND c_nationkey = n2.n_nationkey
AND ((n1.n_name = 'FRANCE' AND n2.n_name = 'GERMANY')
OR (n1.n_name = 'GERMANY' AND n2.n_name = 'FRANCE'))
AND l_shipdate BETWEEN DATE '1995-01-01' AND DATE '1996-12-31'
) AS shipping
GROUP BY supp_nation, cust_nation, l_year
ORDER BY supp_nation, cust_nation, l_year""",

        8: f"""SELECT
o_year,
SUM(CASE WHEN nation = 'BRAZIL' THEN volume ELSE 0 END) / SUM(volume) AS mkt_share
FROM (
SELECT
EXTRACT(YEAR FROM o_orderdate) AS o_year,
l_extendedprice * (1 - l_discount) AS volume,
n2.n_name AS nation
FROM
{tbl('part')},
{tbl('supplier')},
{tbl('lineitem')},
{tbl('orders')},
{tbl('customer')},
{tbl('nation')} n1,
{tbl('nation')} n2,
{tbl('region')}
WHERE
p_partkey = l_partkey
AND s_suppkey = l_suppkey
AND l_orderkey = o_orderkey
AND o_custkey = c_custkey
AND c_nationkey = n1.n_nationkey
AND n1.n_regionkey = r_regionkey
AND r_name = 'AMERICA'
AND s_nationkey = n2.n_nationkey
AND o_orderdate BETWEEN DATE '1995-01-01' AND DATE '1996-12-31'
AND p_type = 'ECONOMY ANODIZED STEEL'
) AS all_nations
GROUP BY o_year
ORDER BY o_year""",

        9: f"""SELECT
nation,
o_year,
SUM(amount) AS sum_profit
FROM (
SELECT
n_name AS nation,
EXTRACT(YEAR FROM o_orderdate) AS o_year,
l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity AS amount
FROM
{tbl('part')},
{tbl('supplier')},
{tbl('lineitem')},
{tbl('partsupp')},
{tbl('orders')},
{tbl('nation')}
WHERE
s_suppkey = l_suppkey
AND ps_suppkey = l_suppkey
AND ps_partkey = l_partkey
AND p_partkey = l_partkey
AND o_orderkey = l_orderkey
AND s_nationkey = n_nationkey
AND p_name LIKE '%green%'
) AS profit
GROUP BY nation, o_year
ORDER BY nation, o_year DESC""",

        10: f"""SELECT
c_custkey,
c_name,
SUM(l_extendedprice * (1 - l_discount)) AS revenue,
c_acctbal,
n_name,
c_address,
c_phone,
c_comment
FROM
{tbl('customer')},
{tbl('orders')},
{tbl('lineitem')},
{tbl('nation')}
WHERE
c_custkey = o_custkey
AND l_orderkey = o_orderkey
AND o_orderdate >= DATE '1993-10-01'
AND o_orderdate < DATE '1993-10-01' + INTERVAL '3' MONTH
AND l_returnflag = 'R'
AND c_nationkey = n_nationkey
GROUP BY
c_custkey,
c_name,
c_acctbal,
c_phone,
n_name,
c_address,
c_comment
ORDER BY revenue DESC
LIMIT 20""",

        11: f"""SELECT
ps_partkey,
SUM(ps_supplycost * ps_availqty) AS value
FROM
{tbl('partsupp')},
{tbl('supplier')},
{tbl('nation')}
WHERE
ps_suppkey = s_suppkey
AND s_nationkey = n_nationkey
AND n_name = 'GERMANY'
GROUP BY ps_partkey
HAVING SUM(ps_supplycost * ps_availqty) > (
SELECT SUM(ps_supplycost * ps_availqty) * 0.0001
FROM
{tbl('partsupp')},
{tbl('supplier')},
{tbl('nation')}
WHERE
ps_suppkey = s_suppkey
AND s_nationkey = n_nationkey
AND n_name = 'GERMANY'
)
ORDER BY value DESC""",

        12: f"""SELECT
l_shipmode,
SUM(CASE
WHEN o_orderpriority = '1-URGENT'
OR o_orderpriority = '2-HIGH'
THEN 1
ELSE 0
END) AS high_line_count,
SUM(CASE
WHEN o_orderpriority <> '1-URGENT'
AND o_orderpriority <> '2-HIGH'
THEN 1
ELSE 0
END) AS low_line_count
FROM
{tbl('orders')},
{tbl('lineitem')}
WHERE
o_orderkey = l_orderkey
AND l_shipmode IN ('MAIL', 'SHIP')
AND l_commitdate < l_receiptdate
AND l_shipdate < l_commitdate
AND l_receiptdate >= DATE '1994-01-01'
AND l_receiptdate < DATE '1994-01-01' + INTERVAL '1' YEAR
GROUP BY l_shipmode
ORDER BY l_shipmode""",

        13: f"""SELECT
c_count,
COUNT(*) AS custdist
FROM (
SELECT
c_custkey,
COUNT(o_orderkey) AS c_count
FROM
{tbl('customer')} LEFT OUTER JOIN {tbl('orders')} ON
c_custkey = o_custkey
AND o_comment NOT LIKE '%special%requests%'
GROUP BY c_custkey
) c_orders
GROUP BY c_count
ORDER BY custdist DESC, c_count DESC""",

        14: f"""SELECT
100.00 * SUM(CASE
WHEN p_type LIKE 'PROMO%'
THEN l_extendedprice * (1 - l_discount)
ELSE 0
END) / SUM(l_extendedprice * (1 - l_discount)) AS promo_revenue
FROM
{tbl('lineitem')},
{tbl('part')}
WHERE
l_partkey = p_partkey
AND l_shipdate >= DATE '1995-09-01'
AND l_shipdate < DATE '1995-09-01' + INTERVAL '1' MONTH""",

        15: [
            f"""CREATE OR REPLACE TEMPORARY VIEW revenue0 AS
            SELECT
              l_suppkey AS supplier_no,
              SUM(l_extendedprice * (1 - l_discount)) AS total_revenue
            FROM {tbl('lineitem')}
            WHERE
              l_shipdate >= DATE '1996-01-01'
              AND l_shipdate < DATE '1996-01-01' + INTERVAL '3' MONTH
            GROUP BY l_suppkey""",
            f"""SELECT
              s_suppkey,
              s_name,
              s_address,
              s_phone,
              total_revenue
            FROM
              {tbl('supplier')},
              revenue0
            WHERE
              s_suppkey = supplier_no
              AND total_revenue = (
                SELECT MAX(total_revenue) FROM revenue0
              )
            ORDER BY s_suppkey""",
            """DROP VIEW revenue0"""
        ],

        16: f"""SELECT
p_brand,
p_type,
p_size,
COUNT(DISTINCT ps_suppkey) AS supplier_cnt
FROM
{tbl('partsupp')},
{tbl('part')}
WHERE
p_partkey = ps_partkey
AND p_brand <> 'Brand#45'
AND p_type NOT LIKE 'MEDIUM POLISHED%'
AND p_size IN (49, 14, 23, 45, 19, 3, 36, 9)
AND ps_suppkey NOT IN (
SELECT s_suppkey
FROM {tbl('supplier')}
WHERE s_comment LIKE '%Customer%Complaints%'
)
GROUP BY p_brand, p_type, p_size
ORDER BY supplier_cnt DESC, p_brand, p_type, p_size""",

        17: f"""SELECT SUM(l_extendedprice) / 7.0 AS avg_yearly
FROM
{tbl('lineitem')},
{tbl('part')}
WHERE
p_partkey = l_partkey
AND p_brand = 'Brand#23'
AND p_container = 'MED BOX'
AND l_quantity < (
SELECT 0.2 * AVG(l_quantity)
FROM {tbl('lineitem')}
WHERE l_partkey = p_partkey
)""",

        18: f"""SELECT
c_name,
c_custkey,
o_orderkey,
o_orderdate,
o_totalprice,
SUM(l_quantity) AS sum_qty
FROM
{tbl('customer')},
{tbl('orders')},
{tbl('lineitem')}
WHERE
o_orderkey IN (
SELECT l_orderkey
FROM {tbl('lineitem')}
GROUP BY l_orderkey
HAVING SUM(l_quantity) > 300
)
AND c_custkey = o_custkey
AND o_orderkey = l_orderkey
GROUP BY
c_name,
c_custkey,
o_orderkey,
o_orderdate,
o_totalprice
ORDER BY o_totalprice DESC, o_orderdate
LIMIT 100""",

        19: f"""SELECT SUM(l_extendedprice * (1 - l_discount)) AS revenue
FROM
{tbl('lineitem')},
{tbl('part')}
WHERE
(
p_partkey = l_partkey
AND p_brand = 'Brand#12'
AND p_container IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')
AND l_quantity >= 1 AND l_quantity <= 1 + 10
AND p_size BETWEEN 1 AND 5
AND l_shipmode IN ('AIR', 'AIR REG')
AND l_shipinstruct = 'DELIVER IN PERSON'
)
OR
(
p_partkey = l_partkey
AND p_brand = 'Brand#23'
AND p_container IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')
AND l_quantity >= 10 AND l_quantity <= 10 + 10
AND p_size BETWEEN 1 AND 10
AND l_shipmode IN ('AIR', 'AIR REG')
AND l_shipinstruct = 'DELIVER IN PERSON'
)
OR
(
p_partkey = l_partkey
AND p_brand = 'Brand#34'
AND p_container IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')
AND l_quantity >= 20 AND l_quantity <= 20 + 10
AND p_size BETWEEN 1 AND 15
AND l_shipmode IN ('AIR', 'AIR REG')
AND l_shipinstruct = 'DELIVER IN PERSON'
)""",

        20: f"""SELECT
s_name,
s_address
FROM
{tbl('supplier')},
{tbl('nation')}
WHERE
s_suppkey IN (
SELECT ps_suppkey
FROM {tbl('partsupp')}
WHERE ps_partkey IN (
SELECT p_partkey
FROM {tbl('part')}
WHERE p_name LIKE 'forest%'
)
AND ps_availqty > (
SELECT 0.5 * SUM(l_quantity)
FROM {tbl('lineitem')}
WHERE l_partkey = ps_partkey
AND l_suppkey = ps_suppkey
AND l_shipdate >= DATE '1994-01-01'
AND l_shipdate < DATE '1994-01-01' + INTERVAL '1' YEAR
)
)
AND s_nationkey = n_nationkey
AND n_name = 'CANADA'
ORDER BY s_name""",

        21: f"""SELECT
s_name,
COUNT(*) AS numwait
FROM
{tbl('supplier')},
{tbl('lineitem')} l1,
{tbl('orders')},
{tbl('nation')}
WHERE
s_suppkey = l1.l_suppkey
AND o_orderkey = l1.l_orderkey
AND o_orderstatus = 'F'
AND l1.l_receiptdate > l1.l_commitdate
AND EXISTS (
SELECT *
FROM {tbl('lineitem')} l2
WHERE l2.l_orderkey = l1.l_orderkey
AND l2.l_suppkey <> l1.l_suppkey
)
AND NOT EXISTS (
SELECT *
FROM {tbl('lineitem')} l3
WHERE l3.l_orderkey = l1.l_orderkey
AND l3.l_suppkey <> l1.l_suppkey
AND l3.l_receiptdate > l3.l_commitdate
)
AND s_nationkey = n_nationkey
AND n_name = 'SAUDI ARABIA'
GROUP BY s_name
ORDER BY numwait DESC, s_name
LIMIT 100""",

        22: f"""SELECT
cntrycode,
COUNT(*) AS numcust,
SUM(c_acctbal) AS totacctbal
FROM (
SELECT
SUBSTRING(c_phone, 1, 2) AS cntrycode,
c_acctbal
FROM {tbl('customer')}
WHERE SUBSTRING(c_phone, 1, 2) IN ('13', '31', '23', '29', '30', '18', '17')
AND c_acctbal > (
SELECT AVG(c_acctbal)
FROM {tbl('customer')}
WHERE c_acctbal > 0.00
AND SUBSTRING(c_phone, 1, 2) IN ('13', '31', '23', '29', '30', '18', '17')
)
AND NOT EXISTS (
SELECT *
FROM {tbl('orders')}
WHERE o_custkey = c_custkey
)
) AS custsale
GROUP BY cntrycode
ORDER BY cntrycode"""
    }

def run_all_tpch_queries(spark: SparkSession, queries: Dict[int, Union[str, List[str]]]) -> List[Dict[str, Any]]:
    all_results = []
    total_start_time = time.perf_counter()
    logger.info("Starting complete TPC-H benchmark run (22 queries)")

    for query_num in range(1, 23):
        if query_num in queries:
            result = run_tpch_query(spark, query_num, queries[query_num])
            all_results.append(result)
            time.sleep(0.5)

    total_end_time = time.perf_counter()
    total_time = total_end_time - total_start_time
    logger.info(f"Complete TPC-H benchmark finished in {total_time:.2f} seconds")

    print("\n" + "=" * 60)
    print("TPC-H BENCHMARK RESULTS SUMMARY")
    print("=" * 60)

    successful_queries = 0
    failed_queries = 0
    total_query_time = 0.0

    for result in all_results:
        status_symbol = "✓" if result["status"] == "SUCCESS" else "✗"
        print(f"{status_symbol} Query {result['query_number']:2d}: {result['time_seconds']:8.2f}s - {result['status']}")
        if result["status"] == "SUCCESS":
            successful_queries += 1
            total_query_time += result["time_seconds"]
        else:
            failed_queries += 1

    print("=" * 60)
    print(f"Successful queries: {successful_queries}/22")
    print(f"Failed queries: {failed_queries}/22")
    print(f"Total query execution time: {total_query_time:.2f} seconds")
    print(f"Total benchmark time: {total_time:.2f} seconds")
    print("=" * 60)

    return all_results

def main():
    parser = argparse.ArgumentParser(description="Run TPC-H on Iceberg (Glue catalog)")
    parser.add_argument("--catalog", default="glue", help="Iceberg catalog name (default: glue)")
    parser.add_argument("--database", default="postgres_postgres_tpch", help="Glue database name containing TPC-H tables")
    parser.add_argument("--warehouse", default=os.environ.get("ICEBERG_WAREHOUSE"), help="S3 warehouse path, e.g., s3://bucket/warehouse/")
    parser.add_argument("--region", default=os.environ.get("AWS_REGION"), help="AWS region for Glue/S3 (optional)")
    parser.add_argument("--shuffle-partitions", type=int, default=64, help="spark.sql.shuffle.partitions")
    args = parser.parse_args()

    spark = build_spark_session(
        app_name="TPC-H Iceberg Benchmark",
        catalog=args.catalog,
        warehouse=args.warehouse,
        region=args.region,
    )
    configure_spark_settings(spark, args.shuffle_partitions)

    queries = build_queries(args.catalog, args.database)
    _ = run_all_tpch_queries(spark, queries)

if __name__ == "__main__":
    main()