[新機能]Snowflake の Pipeline Builder で SQL + Python パイプラインを GUI で組んで Deploy してみた

USE ROLE ACCOUNTADMIN;

CREATE ROLE PB_ENGINEER;
GRANT ROLE PB_ENGINEER TO ROLE SYSADMIN;
GRANT ROLE PB_ENGINEER TO USER <自分のユーザー名>;

CREATE WAREHOUSE PB_WH WAREHOUSE_SIZE = XSMALL AUTO_SUSPEND = 60 AUTO_RESUME = TRUE;
GRANT USAGE ON WAREHOUSE PB_WH TO ROLE PB_ENGINEER;

CREATE DATABASE PB_DEMO;
CREATE SCHEMA PB_DEMO.DEV;
CREATE SCHEMA PB_DEMO.PROD;

-- dbt Labs 公開の jaffle shop データ(6 年分)
CREATE STAGE PB_DEMO.PUBLIC.JAFFLE_S3
  URL = 's3://dbt-tutorial-public/long_term_dataset/'
  FILE_FORMAT = (TYPE = CSV SKIP_HEADER = 1 FIELD_OPTIONALLY_ENCLOSED_BY = '"' NULL_IF = ('', 'NULL'));

-- ドキュメントの Access control requirements に沿って個別に付与
GRANT USAGE ON DATABASE PB_DEMO TO ROLE PB_ENGINEER;
GRANT USAGE ON SCHEMA PB_DEMO.PUBLIC TO ROLE PB_ENGINEER;
GRANT USAGE ON STAGE PB_DEMO.PUBLIC.JAFFLE_S3 TO ROLE PB_ENGINEER;
GRANT USAGE, CREATE TABLE, CREATE VIEW, CREATE TASK, CREATE CODE BUNDLE, CREATE DCM PROJECT
  ON SCHEMA PB_DEMO.DEV TO ROLE PB_ENGINEER;
GRANT USAGE, CREATE TABLE, CREATE VIEW, CREATE TASK, CREATE CODE BUNDLE, CREATE DCM PROJECT
  ON SCHEMA PB_DEMO.PROD TO ROLE PB_ENGINEER;
GRANT EXECUTE TASK ON ACCOUNT TO ROLE PB_ENGINEER;

-- パイプラインを置く共有 Workspace を PB_ENGINEER が作れるようにする
GRANT CREATE WORKSPACE ON SCHEMA PB_DEMO.PUBLIC TO ROLE PB_ENGINEER;

-- 失敗時 Task が書き込むログテーブル(環境ごと)
USE ROLE PB_ENGINEER;
CREATE TABLE PB_DEMO.DEV.PIPELINE_RUN_LOG (logged_at TIMESTAMP_LTZ, graph_name VARCHAR, message VARCHAR);
CREATE TABLE PB_DEMO.PROD.PIPELINE_RUN_LOG (logged_at TIMESTAMP_LTZ, graph_name VARCHAR, message VARCHAR);

CREATE TABLE IF NOT EXISTS raw_customers (
    id VARCHAR,
    name VARCHAR
);
CREATE TABLE IF NOT EXISTS raw_orders (
    id VARCHAR,
    customer VARCHAR,
    ordered_at TIMESTAMP_NTZ,
    store_id VARCHAR,
    subtotal NUMBER,
    tax_paid NUMBER,
    order_total NUMBER
);
CREATE TABLE IF NOT EXISTS raw_order_items (
    id VARCHAR,
    order_id VARCHAR,
    sku VARCHAR
);
CREATE TABLE IF NOT EXISTS raw_products (
    sku VARCHAR,
    name VARCHAR,
    type VARCHAR,
    price NUMBER,
    description VARCHAR
);
CREATE TABLE IF NOT EXISTS raw_supplies (
    id VARCHAR,
    name VARCHAR,
    cost NUMBER,
    perishable BOOLEAN,
    sku VARCHAR
);
CREATE TABLE IF NOT EXISTS raw_stores (
    id VARCHAR,
    name VARCHAR,
    opened_at TIMESTAMP_NTZ,
    tax_rate FLOAT
);

COPY INTO raw_customers   FROM @PB_DEMO.PUBLIC.JAFFLE_S3/raw_customers.csv;
COPY INTO raw_orders      FROM @PB_DEMO.PUBLIC.JAFFLE_S3/raw_orders.csv;
COPY INTO raw_order_items FROM @PB_DEMO.PUBLIC.JAFFLE_S3/raw_order_items.csv;
COPY INTO raw_products    FROM @PB_DEMO.PUBLIC.JAFFLE_S3/raw_products.csv;
COPY INTO raw_supplies    FROM @PB_DEMO.PUBLIC.JAFFLE_S3/raw_supplies.csv;
COPY INTO raw_stores      FROM @PB_DEMO.PUBLIC.JAFFLE_S3/raw_stores.csv;

CREATE OR REPLACE VIEW stg_customers AS
SELECT
    id   AS customer_id,
    name AS customer_name
FROM raw_customers;

CREATE OR REPLACE VIEW stg_locations AS
SELECT
    id   AS location_id,
    name AS location_name,
    tax_rate,
    DATE_TRUNC('day', opened_at) AS opened_date
FROM raw_stores;

CREATE OR REPLACE VIEW stg_order_items AS
SELECT
    id       AS order_item_id,
    order_id,
    sku      AS product_id
FROM raw_order_items;

CREATE OR REPLACE VIEW stg_orders AS
SELECT
    id          AS order_id,
    store_id    AS location_id,
    customer    AS customer_id,
    subtotal    AS subtotal_cents,
    tax_paid    AS tax_paid_cents,
    order_total AS order_total_cents,
    (subtotal / 100)::NUMBER(16, 2)    AS subtotal,
    (tax_paid / 100)::NUMBER(16, 2)    AS tax_paid,
    (order_total / 100)::NUMBER(16, 2) AS order_total,
    DATE_TRUNC('day', ordered_at)      AS ordered_at
FROM raw_orders;

CREATE OR REPLACE VIEW stg_products AS
SELECT
    sku         AS product_id,
    name        AS product_name,
    type        AS product_type,
    description AS product_description,
    (price / 100)::NUMBER(16, 2)        AS product_price,
    COALESCE(type = 'jaffle', FALSE)    AS is_food_item,
    COALESCE(type = 'beverage', FALSE)  AS is_drink_item
FROM raw_products;

CREATE OR REPLACE VIEW stg_supplies AS
SELECT
    MD5(COALESCE(id, '') || '-' || COALESCE(sku, '')) AS supply_uuid,
    id         AS supply_id,
    sku        AS product_id,
    name       AS supply_name,
    (cost / 100)::NUMBER(16, 2) AS supply_cost,
    perishable AS is_perishable_supply
FROM raw_supplies;

CREATE OR REPLACE TABLE order_items AS
WITH order_supplies_summary AS (
    SELECT
        product_id,
        SUM(supply_cost) AS supply_cost
    FROM stg_supplies
    GROUP BY 1
)
SELECT
    oi.*,
    o.ordered_at,
    p.product_name,
    p.product_price,
    p.is_food_item,
    p.is_drink_item,
    s.supply_cost
FROM stg_order_items AS oi
LEFT JOIN stg_orders   AS o ON oi.order_id = o.order_id
LEFT JOIN stg_products AS p ON oi.product_id = p.product_id
LEFT JOIN order_supplies_summary AS s ON oi.product_id = s.product_id;

CREATE OR REPLACE TABLE orders AS
WITH order_items_summary AS (
    SELECT
        order_id,
        SUM(supply_cost)     AS order_cost,
        SUM(product_price)   AS order_items_subtotal,
        COUNT(order_item_id) AS count_order_items,
        SUM(CASE WHEN is_food_item  THEN 1 ELSE 0 END) AS count_food_items,
        SUM(CASE WHEN is_drink_item THEN 1 ELSE 0 END) AS count_drink_items
    FROM order_items
    GROUP BY 1
)
SELECT
    o.*,
    s.order_cost,
    s.order_items_subtotal,
    s.count_food_items,
    s.count_drink_items,
    s.count_order_items,
    s.count_food_items  > 0 AS is_food_order,
    s.count_drink_items > 0 AS is_drink_order,
    ROW_NUMBER() OVER (PARTITION BY o.customer_id ORDER BY o.ordered_at ASC) AS customer_order_number
FROM stg_orders AS o
LEFT JOIN order_items_summary AS s ON o.order_id = s.order_id;

CREATE OR REPLACE TABLE customers AS
WITH customer_orders_summary AS (
    SELECT
        customer_id,
        COUNT(DISTINCT order_id)     AS count_lifetime_orders,
        COUNT(DISTINCT order_id) > 1 AS is_repeat_buyer,
        MIN(ordered_at)  AS first_ordered_at,
        MAX(ordered_at)  AS last_ordered_at,
        SUM(subtotal)    AS lifetime_spend_pretax,
        SUM(tax_paid)    AS lifetime_tax_paid,
        SUM(order_total) AS lifetime_spend
    FROM orders
    GROUP BY 1
)
SELECT
    c.*,
    s.count_lifetime_orders,
    s.first_ordered_at,
    s.last_ordered_at,
    s.lifetime_spend_pretax,
    s.lifetime_tax_paid,
    s.lifetime_spend,
    CASE WHEN s.is_repeat_buyer THEN 'returning' ELSE 'new' END AS customer_type
FROM stg_customers AS c
LEFT JOIN customer_orders_summary AS s ON c.customer_id = s.customer_id;

CREATE OR REPLACE TABLE products  AS SELECT * FROM stg_products;
CREATE OR REPLACE TABLE locations AS SELECT * FROM stg_locations;
CREATE OR REPLACE TABLE supplies  AS SELECT * FROM stg_supplies;

"""marts の件数チェック (Python タスク)。

ORDERS の件数が --min-orders 未満なら例外を送出し、タスクを失敗させる。
成功時は PIPELINE_CHECK_LOG に件数を記録する。
"""
import argparse

from snowflake.snowpark.context import get_active_session

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--min-orders", type=int, default=1)
    parser.add_argument("--schema", default=None, help="database.schema (default: current schema of the session)")
    args = parser.parse_args()

    session = get_active_session()
    schema = args.schema or f"{session.get_current_database()}.{session.get_current_schema()}"

    counts = {t: session.table(f"{schema}.{t}").count() for t in ("ORDERS", "ORDER_ITEMS", "CUSTOMERS")}
    print(f"schema={schema} counts={counts} min_orders={args.min_orders}")

    if counts["ORDERS"] < args.min_orders:
        raise RuntimeError(f"ORDERS has {counts['ORDERS']} rows, less than --min-orders {args.min_orders}")

    session.sql(
        f"CREATE TABLE IF NOT EXISTS {schema}.PIPELINE_CHECK_LOG "
        "(checked_at TIMESTAMP_LTZ, orders NUMBER, order_items NUMBER, customers NUMBER)"
    ).collect()
    session.sql(
        f"INSERT INTO {schema}.PIPELINE_CHECK_LOG VALUES "
        f"(CURRENT_TIMESTAMP(), {counts['ORDERS']}, {counts['ORDER_ITEMS']}, {counts['CUSTOMERS']})"
    ).collect()
    print("check passed")

if __name__ == "__main__":
    main()

definitions:
  - resource: task
    name: PB_DEMO.DEV.load_raw_task
    definition:
      type: sql_file
      location:
        type: stage_uri
        uri: >-
          snow://workspace/PB_DEMO.PUBLIC.PB_SHARED/versions/last/JAFFLE_PIPELINE/01_load_raw.sql
  - resource: task
    name: PB_DEMO.DEV.staging_task
    definition:
      type: sql_file
      location:
        type: stage_uri
        uri: >-
          snow://workspace/PB_DEMO.PUBLIC.PB_SHARED/versions/last/JAFFLE_PIPELINE/02_staging.sql
    predecessors:
      - PB_DEMO.DEV.LOAD_RAW_TASK
  - resource: task_graph_object
    name: PB_DEMO.DEV.JAFFLE_PIPELINE
    roots:
      - PB_DEMO.DEV.load_raw_task
    schedule:
      type: cron
      value: 0 9 * * * UTC
    state: suspended
    task_defaults:
      warehouse: PB_WH
  - resource: task
    name: PB_DEMO.DEV.marts_task
    definition:
      type: sql_file
      location:
        type: stage_uri
        uri: >-
          snow://workspace/PB_DEMO.PUBLIC.PB_SHARED/versions/last/JAFFLE_PIPELINE/03_marts.sql
    predecessors:
      - PB_DEMO.DEV.STAGING_TASK
  - resource: code_bundle
    name: PB_DEMO.DEV.jaffle_pipeline_waqvc
    from:
      type: asset
      asset: project
  - resource: task
    name: PB_DEMO.DEV.check_marts_task
    definition:
      type: code_bundle
      name: PB_DEMO.DEV.jaffle_pipeline_waqvc
      entrypoint: check_marts.py
      specification:
        bundle:
          type: custom
          compute_type: warehouse
          language: python
          compute_options:
            runtime_version: '3.11'
          properties:
            requirements_file: ''
      arguments:
        - '--min-orders'
        - '1'
    predecessors:
      - PB_DEMO.DEV.MARTS_TASK
    warehouse: PB_WH
  - resource: task
    name: PB_DEMO.DEV.PUBLISH_SUMMARY
    definition:
      type: sql
      value: |-
        CREATE OR REPLACE TABLE daily_summary AS
        SELECT
            TO_DATE(ordered_at) AS order_date,
            COUNT(*)            AS order_count,
            SUM(order_total)    AS revenue
        FROM orders
        GROUP BY 1;
    predecessors:
      - PB_DEMO.DEV.CHECK_MARTS_TASK
    warehouse: PB_WH
  - resource: task
    name: PB_DEMO.DEV.LOG_FAILURE
    definition:
      type: sql
      value: |-
        INSERT INTO pipeline_run_log (logged_at, graph_name, message)
        SELECT
            CURRENT_TIMESTAMP(),
            SYSTEM$TASK_RUNTIME_INFO('CURRENT_TASK_GRAPH_NAME'),
            'CHECK_MARTS failed';
    predecessors:
      - PB_DEMO.DEV.CHECK_MARTS_TASK
    warehouse: PB_WH
    condition: SYSTEM$PREDECESSORS_STATE('AT_LEAST_ONE_FAILED')