helpers.py

import os

FILE_DIR = os.path.dirname(__file__)
FIXTURES_DIR = os.path.join(FILE_DIR, "fixtures")


def _filter_comments(s):
    return "\n".join([line for line in s.splitlines() if line and not line.startswith("--")])


def _extract_meta(sql):
    meta = {}
    sql_lines = sql.split("\n")
    i = 0
    while sql_lines[i].startswith("#"):
        key, val = sql_lines[i].split(":", maxsplit=1)
        meta[key.lstrip("#").strip()] = val.strip()
        i += 1
    sql = "\n".join(sql_lines[i:])
    return sql, meta


def assert_logger_contains(message, logger, level="error"):
    output = "\n".join(str(args[0][0]) for args in getattr(logger, level).call_args_list)
    if message not in output:
        print(f"Expected '{message}' not in {output}")
        raise


def load_sql_fixtures(filename):
    with open(os.path.join(FIXTURES_DIR, filename), encoding="utf-8") as f:
        yield from _filter_comments(f.read()).splitlines()


def load_sql_fixture_pairs(filename):
    with open(os.path.join(FIXTURES_DIR, filename), encoding="utf-8") as f:
        statements = _filter_comments(f.read()).split(";")

        size = len(statements)

        for i in range(0, size, 2):
            if i + 1 < size:
                sql = statements[i].strip()
                sql, meta = _extract_meta(sql)
                expected = statements[i + 1].strip()
                yield meta, sql, expected


def string_to_bool(string):
    if string is None:
        return False
    if string in (True, False):
        return string
    return string and string.lower() in ("true", "1")


SKIP_INTEGRATION = string_to_bool(os.environ.get("SKIP_INTEGRATION", "0").lower())

TPCH_SCHEMA = {
    "lineitem": {
        "l_orderkey": "bigint",
        "l_partkey": "bigint",
        "l_suppkey": "bigint",
        "l_linenumber": "bigint",
        "l_quantity": "double",
        "l_extendedprice": "double",
        "l_discount": "double",
        "l_tax": "double",
        "l_returnflag": "string",
        "l_linestatus": "string",
        "l_shipdate": "string",
        "l_commitdate": "string",
        "l_receiptdate": "string",
        "l_shipinstruct": "string",
        "l_shipmode": "string",
        "l_comment": "string",
    },
    "orders": {
        "o_orderkey": "bigint",
        "o_custkey": "bigint",
        "o_orderstatus": "string",
        "o_totalprice": "double",
        "o_orderdate": "string",
        "o_orderpriority": "string",
        "o_clerk": "string",
        "o_shippriority": "int",
        "o_comment": "string",
    },
    "customer": {
        "c_custkey": "bigint",
        "c_name": "string",
        "c_address": "string",
        "c_nationkey": "bigint",
        "c_phone": "string",
        "c_acctbal": "double",
        "c_mktsegment": "string",
        "c_comment": "string",
    },
    "part": {
        "p_partkey": "bigint",
        "p_name": "string",
        "p_mfgr": "string",
        "p_brand": "string",
        "p_type": "string",
        "p_size": "int",
        "p_container": "string",
        "p_retailprice": "double",
        "p_comment": "string",
    },
    "supplier": {
        "s_suppkey": "bigint",
        "s_name": "string",
        "s_address": "string",
        "s_nationkey": "bigint",
        "s_phone": "string",
        "s_acctbal": "double",
        "s_comment": "string",
    },
    "partsupp": {
        "ps_partkey": "bigint",
        "ps_suppkey": "bigint",
        "ps_availqty": "int",
        "ps_supplycost": "double",
        "ps_comment": "string",
    },
    "nation": {
        "n_nationkey": "bigint",
        "n_name": "string",
        "n_regionkey": "bigint",
        "n_comment": "string",
    },
    "region": {
        "r_regionkey": "bigint",
        "r_name": "string",
        "r_comment": "string",
    },
}

TPCDS_SCHEMA = {
    "catalog_sales": {
        "cs_sold_date_sk": "bigint",
        "cs_sold_time_sk": "bigint",
        "cs_ship_date_sk": "bigint",
        "cs_bill_customer_sk": "bigint",
        "cs_bill_cdemo_sk": "bigint",
        "cs_bill_hdemo_sk": "bigint",
        "cs_bill_addr_sk": "bigint",
        "cs_ship_customer_sk": "bigint",
        "cs_ship_cdemo_sk": "bigint",
        "cs_ship_hdemo_sk": "bigint",
        "cs_ship_addr_sk": "bigint",
        "cs_call_center_sk": "bigint",
        "cs_catalog_page_sk": "bigint",
        "cs_ship_mode_sk": "bigint",
        "cs_warehouse_sk": "bigint",
        "cs_item_sk": "bigint",
        "cs_promo_sk": "bigint",
        "cs_order_number": "bigint",
        "cs_quantity": "bigint",
        "cs_wholesale_cost": "double",
        "cs_list_price": "double",
        "cs_sales_price": "double",
        "cs_ext_discount_amt": "double",
        "cs_ext_sales_price": "double",
        "cs_ext_wholesale_cost": "double",
        "cs_ext_list_price": "double",
        "cs_ext_tax": "double",
        "cs_coupon_amt": "double",
        "cs_ext_ship_cost": "double",
        "cs_net_paid": "double",
        "cs_net_paid_inc_tax": "double",
        "cs_net_paid_inc_ship": "double",
        "cs_net_paid_inc_ship_tax": "double",
        "cs_net_profit": "double",
    },
    "catalog_returns": {
        "cr_returned_date_sk": "bigint",
        "cr_returned_time_sk": "bigint",
        "cr_item_sk": "bigint",
        "cr_refunded_customer_sk": "bigint",
        "cr_refunded_cdemo_sk": "bigint",
        "cr_refunded_hdemo_sk": "bigint",
        "cr_refunded_addr_sk": "bigint",
        "cr_returning_customer_sk": "bigint",
        "cr_returning_cdemo_sk": "bigint",
        "cr_returning_hdemo_sk": "bigint",
        "cr_returning_addr_sk": "bigint",
        "cr_call_center_sk": "bigint",
        "cr_catalog_page_sk": "bigint",
        "cr_ship_mode_sk": "bigint",
        "cr_warehouse_sk": "bigint",
        "cr_reason_sk": "bigint",
        "cr_order_number": "bigint",
        "cr_return_quantity": "bigint",
        "cr_return_amount": "double",
        "cr_return_tax": "double",
        "cr_return_amt_inc_tax": "double",
        "cr_fee": "double",
        "cr_return_ship_cost": "double",
        "cr_refunded_cash": "double",
        "cr_reversed_charge": "double",
        "cr_store_credit": "double",
        "cr_net_loss": "double",
    },
    "inventory": {
        "inv_date_sk": "bigint",
        "inv_item_sk": "bigint",
        "inv_warehouse_sk": "bigint",
        "inv_quantity_on_hand": "bigint",
    },
    "store_sales": {
        "ss_sold_date_sk": "bigint",
        "ss_sold_time_sk": "bigint",
        "ss_item_sk": "bigint",
        "ss_customer_sk": "bigint",
        "ss_cdemo_sk": "bigint",
        "ss_hdemo_sk": "bigint",
        "ss_addr_sk": "bigint",
        "ss_store_sk": "bigint",
        "ss_promo_sk": "bigint",
        "ss_ticket_number": "bigint",
        "ss_quantity": "bigint",
        "ss_wholesale_cost": "double",
        "ss_list_price": "double",
        "ss_sales_price": "double",
        "ss_ext_discount_amt": "double",
        "ss_ext_sales_price": "double",
        "ss_ext_wholesale_cost": "double",
        "ss_ext_list_price": "double",
        "ss_ext_tax": "double",
        "ss_coupon_amt": "double",
        "ss_net_paid": "double",
        "ss_net_paid_inc_tax": "double",
        "ss_net_profit": "double",
    },
    "store_returns": {
        "sr_returned_date_sk": "bigint",
        "sr_return_time_sk": "bigint",
        "sr_item_sk": "bigint",
        "sr_customer_sk": "bigint",
        "sr_cdemo_sk": "bigint",
        "sr_hdemo_sk": "bigint",
        "sr_addr_sk": "bigint",
        "sr_store_sk": "bigint",
        "sr_reason_sk": "bigint",
        "sr_ticket_number": "bigint",
        "sr_return_quantity": "bigint",
        "sr_return_amt": "double",
        "sr_return_tax": "double",
        "sr_return_amt_inc_tax": "double",
        "sr_fee": "double",
        "sr_return_ship_cost": "double",
        "sr_refunded_cash": "double",
        "sr_reversed_charge": "double",
        "sr_store_credit": "double",
        "sr_net_loss": "double",
    },
    "web_sales": {
        "ws_sold_date_sk": "bigint",
        "ws_sold_time_sk": "bigint",
        "ws_ship_date_sk": "bigint",
        "ws_item_sk": "bigint",
        "ws_bill_customer_sk": "bigint",
        "ws_bill_cdemo_sk": "bigint",
        "ws_bill_hdemo_sk": "bigint",
        "ws_bill_addr_sk": "bigint",
        "ws_ship_customer_sk": "bigint",
        "ws_ship_cdemo_sk": "bigint",
        "ws_ship_hdemo_sk": "bigint",
        "ws_ship_addr_sk": "bigint",
        "ws_web_page_sk": "bigint",
        "ws_web_site_sk": "bigint",
        "ws_ship_mode_sk": "bigint",
        "ws_warehouse_sk": "bigint",
        "ws_promo_sk": "bigint",
        "ws_order_number": "bigint",
        "ws_quantity": "bigint",
        "ws_wholesale_cost": "double",
        "ws_list_price": "double",
        "ws_sales_price": "double",
        "ws_ext_discount_amt": "double",
        "ws_ext_sales_price": "double",
        "ws_ext_wholesale_cost": "double",
        "ws_ext_list_price": "double",
        "ws_ext_tax": "double",
        "ws_coupon_amt": "double",
        "ws_ext_ship_cost": "double",
        "ws_net_paid": "double",
        "ws_net_paid_inc_tax": "double",
        "ws_net_paid_inc_ship": "double",
        "ws_net_paid_inc_ship_tax": "double",
        "ws_net_profit": "double",
    },
    "web_returns": {
        "wr_returned_date_sk": "bigint",
        "wr_returned_time_sk": "bigint",
        "wr_item_sk": "bigint",
        "wr_refunded_customer_sk": "bigint",
        "wr_refunded_cdemo_sk": "bigint",
        "wr_refunded_hdemo_sk": "bigint",
        "wr_refunded_addr_sk": "bigint",
        "wr_returning_customer_sk": "bigint",
        "wr_returning_cdemo_sk": "bigint",
        "wr_returning_hdemo_sk": "bigint",
        "wr_returning_addr_sk": "bigint",
        "wr_web_page_sk": "bigint",
        "wr_reason_sk": "bigint",
        "wr_order_number": "bigint",
        "wr_return_quantity": "bigint",
        "wr_return_amt": "double",
        "wr_return_tax": "double",
        "wr_return_amt_inc_tax": "double",
        "wr_fee": "double",
        "wr_return_ship_cost": "double",
        "wr_refunded_cash": "double",
        "wr_reversed_charge": "double",
        "wr_account_credit": "double",
        "wr_net_loss": "double",
    },
    "call_center": {
        "cc_call_center_sk": "bigint",
        "cc_call_center_id": "string",
        "cc_rec_start_date": "string",
        "cc_rec_end_date": "string",
        "cc_closed_date_sk": "bigint",
        "cc_open_date_sk": "bigint",
        "cc_name": "string",
        "cc_class": "string",
        "cc_employees": "bigint",
        "cc_sq_ft": "bigint",
        "cc_hours": "string",
        "cc_manager": "string",
        "cc_mkt_id": "bigint",
        "cc_mkt_class": "string",
        "cc_mkt_desc": "string",
        "cc_market_manager": "string",
        "cc_division": "bigint",
        "cc_division_name": "string",
        "cc_company": "bigint",
        "cc_company_name": "string",
        "cc_street_number": "string",
        "cc_street_name": "string",
        "cc_street_type": "string",
        "cc_suite_number": "string",
        "cc_city": "string",
        "cc_county": "string",
        "cc_state": "string",
        "cc_zip": "string",
        "cc_country": "string",
        "cc_gmt_offset": "double",
        "cc_tax_percentage": "double",
    },
    "catalog_page": {
        "cp_catalog_page_sk": "bigint",
        "cp_catalog_page_id": "string",
        "cp_start_date_sk": "bigint",
        "cp_end_date_sk": "bigint",
        "cp_department": "string",
        "cp_catalog_number": "bigint",
        "cp_catalog_page_number": "bigint",
        "cp_description": "string",
        "cp_type": "string",
    },
    "customer": {
        "c_customer_sk": "bigint",
        "c_customer_id": "string",
        "c_current_cdemo_sk": "bigint",
        "c_current_hdemo_sk": "bigint",
        "c_current_addr_sk": "bigint",
        "c_first_shipto_date_sk": "bigint",
        "c_first_sales_date_sk": "bigint",
        "c_salutation": "string",
        "c_first_name": "string",
        "c_last_name": "string",
        "c_preferred_cust_flag": "string",
        "c_birth_day": "bigint",
        "c_birth_month": "bigint",
        "c_birth_year": "bigint",
        "c_birth_country": "string",
        "c_login": "string",
        "c_email_address": "string",
        "c_last_review_date": "string",
    },
    "customer_address": {
        "ca_address_sk": "bigint",
        "ca_address_id": "string",
        "ca_street_number": "string",
        "ca_street_name": "string",
        "ca_street_type": "string",
        "ca_suite_number": "string",
        "ca_city": "string",
        "ca_county": "string",
        "ca_state": "string",
        "ca_zip": "string",
        "ca_country": "string",
        "ca_gmt_offset": "double",
        "ca_location_type": "string",
    },
    "customer_demographics": {
        "cd_demo_sk": "bigint",
        "cd_gender": "string",
        "cd_marital_status": "string",
        "cd_education_status": "string",
        "cd_purchase_estimate": "bigint",
        "cd_credit_rating": "string",
        "cd_dep_count": "bigint",
        "cd_dep_employed_count": "bigint",
        "cd_dep_college_count": "bigint",
    },
    "date_dim": {
        "d_date_sk": "bigint",
        "d_date_id": "string",
        "d_date": "string",
        "d_month_seq": "bigint",
        "d_week_seq": "bigint",
        "d_quarter_seq": "bigint",
        "d_year": "bigint",
        "d_dow": "bigint",
        "d_moy": "bigint",
        "d_dom": "bigint",
        "d_qoy": "bigint",
        "d_fy_year": "bigint",
        "d_fy_quarter_seq": "bigint",
        "d_fy_week_seq": "bigint",
        "d_day_name": "string",
        "d_quarter_name": "string",
        "d_holiday": "string",
        "d_weekend": "string",
        "d_following_holiday": "string",
        "d_first_dom": "bigint",
        "d_last_dom": "bigint",
        "d_same_day_ly": "bigint",
        "d_same_day_lq": "bigint",
        "d_current_day": "string",
        "d_current_week": "string",
        "d_current_month": "string",
        "d_current_quarter": "string",
        "d_current_year": "string",
    },
    "household_demographics": {
        "hd_demo_sk": "bigint",
        "hd_income_band_sk": "bigint",
        "hd_buy_potential": "string",
        "hd_dep_count": "bigint",
        "hd_vehicle_count": "bigint",
    },
    "income_band": {
        "ib_income_band_sk": "bigint",
        "ib_lower_bound": "bigint",
        "ib_upper_bound": "bigint",
    },
    "item": {
        "i_item_sk": "bigint",
        "i_item_id": "string",
        "i_rec_start_date": "string",
        "i_rec_end_date": "string",
        "i_item_desc": "string",
        "i_current_price": "double",
        "i_wholesale_cost": "double",
        "i_brand_id": "bigint",
        "i_brand": "string",
        "i_class_id": "bigint",
        "i_class": "string",
        "i_category_id": "bigint",
        "i_category": "string",
        "i_manufact_id": "bigint",
        "i_manufact": "string",
        "i_size": "string",
        "i_formulation": "string",
        "i_color": "string",
        "i_units": "string",
        "i_container": "string",
        "i_manager_id": "bigint",
        "i_product_name": "string",
    },
    "promotion": {
        "p_promo_sk": "bigint",
        "p_promo_id": "string",
        "p_start_date_sk": "bigint",
        "p_end_date_sk": "bigint",
        "p_item_sk": "bigint",
        "p_cost": "double",
        "p_response_target": "bigint",
        "p_promo_name": "string",
        "p_channel_dmail": "string",
        "p_channel_email": "string",
        "p_channel_catalog": "string",
        "p_channel_tv": "string",
        "p_channel_radio": "string",
        "p_channel_press": "string",
        "p_channel_event": "string",
        "p_channel_demo": "string",
        "p_channel_details": "string",
        "p_purpose": "string",
        "p_discount_active": "string",
    },
    "reason": {"r_reason_sk": "bigint", "r_reason_id": "string", "r_reason_desc": "string"},
    "ship_mode": {
        "sm_ship_mode_sk": "bigint",
        "sm_ship_mode_id": "string",
        "sm_type": "string",
        "sm_code": "string",
        "sm_carrier": "string",
        "sm_contract": "string",
    },
    "store": {
        "s_store_sk": "bigint",
        "s_store_id": "string",
        "s_rec_start_date": "string",
        "s_rec_end_date": "string",
        "s_closed_date_sk": "bigint",
        "s_store_name": "string",
        "s_number_employees": "bigint",
        "s_floor_space": "bigint",
        "s_hours": "string",
        "s_manager": "string",
        "s_market_id": "bigint",
        "s_geography_class": "string",
        "s_market_desc": "string",
        "s_market_manager": "string",
        "s_division_id": "bigint",
        "s_division_name": "string",
        "s_company_id": "bigint",
        "s_company_name": "string",
        "s_street_number": "string",
        "s_street_name": "string",
        "s_street_type": "string",
        "s_suite_number": "string",
        "s_city": "string",
        "s_county": "string",
        "s_state": "string",
        "s_zip": "string",
        "s_country": "string",
        "s_gmt_offset": "double",
        "s_tax_precentage": "double",
    },
    "time_dim": {
        "t_time_sk": "bigint",
        "t_time_id": "string",
        "t_time": "bigint",
        "t_hour": "bigint",
        "t_minute": "bigint",
        "t_second": "bigint",
        "t_am_pm": "string",
        "t_shift": "string",
        "t_sub_shift": "string",
        "t_meal_time": "string",
    },
    "warehouse": {
        "w_warehouse_sk": "bigint",
        "w_warehouse_id": "string",
        "w_warehouse_name": "string",
        "w_warehouse_sq_ft": "bigint",
        "w_street_number": "string",
        "w_street_name": "string",
        "w_street_type": "string",
        "w_suite_number": "string",
        "w_city": "string",
        "w_county": "string",
        "w_state": "string",
        "w_zip": "string",
        "w_country": "string",
        "w_gmt_offset": "double",
    },
    "web_page": {
        "wp_web_page_sk": "bigint",
        "wp_web_page_id": "string",
        "wp_rec_start_date": "string",
        "wp_rec_end_date": "string",
        "wp_creation_date_sk": "bigint",
        "wp_access_date_sk": "bigint",
        "wp_autogen_flag": "string",
        "wp_customer_sk": "bigint",
        "wp_url": "string",
        "wp_type": "string",
        "wp_char_count": "bigint",
        "wp_link_count": "bigint",
        "wp_image_count": "bigint",
        "wp_max_ad_count": "bigint",
    },
    "web_site": {
        "web_site_sk": "bigint",
        "web_site_id": "string",
        "web_rec_start_date": "string",
        "web_rec_end_date": "string",
        "web_name": "string",
        "web_open_date_sk": "bigint",
        "web_close_date_sk": "bigint",
        "web_class": "string",
        "web_manager": "string",
        "web_mkt_id": "bigint",
        "web_mkt_class": "string",
        "web_mkt_desc": "string",
        "web_market_manager": "string",
        "web_company_id": "bigint",
        "web_company_name": "string",
        "web_street_number": "string",
        "web_street_name": "string",
        "web_street_type": "string",
        "web_suite_number": "string",
        "web_city": "string",
        "web_county": "string",
        "web_state": "string",
        "web_zip": "string",
        "web_country": "string",
        "web_gmt_offset": "string",
        "web_tax_percentage": "double",
    },
}