@datarecce/ui 0.1.30 → 0.1.31

package/recce-source/integration_tests/sqlmesh/config.py

@@ -0,0 +1,171 @@
+import os
+
+from sqlmesh.core.config import (
+    AirflowSchedulerConfig,
+    AutoCategorizationMode,
+    BigQueryConnectionConfig,
+    CategorizerConfig,
+    Config,
+    DuckDBConnectionConfig,
+    EnvironmentSuffixTarget,
+    GatewayConfig,
+    ModelDefaultsConfig,
+    PlanConfig,
+    PostgresConnectionConfig,
+    SparkConnectionConfig,
+)
+from sqlmesh.core.notification_target import (
+    BasicSMTPNotificationTarget,
+    SlackApiNotificationTarget,
+    SlackWebhookNotificationTarget,
+)
+from sqlmesh.core.user import User, UserRole
+
+CURRENT_FILE_PATH = os.path.abspath(__file__)
+DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
+
+
+# An in memory DuckDB config.
+config = Config(
+    default_connection=DuckDBConnectionConfig(),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)
+
+bigquery_config = Config(
+    gateways={
+        "gcp": GatewayConfig(
+            connection=BigQueryConnectionConfig(),
+            state_connection=DuckDBConnectionConfig(database=f"{DATA_DIR}/bigquery.duckdb"),
+        )
+    },
+    default_gateway="gcp",
+    model_defaults=ModelDefaultsConfig(dialect="bigquery"),
+)
+
+# A configuration used for SQLMesh tests.
+test_config = Config(
+    gateways={"in_memory": GatewayConfig(connection=DuckDBConnectionConfig())},
+    default_gateway="in_memory",
+    plan=PlanConfig(auto_categorize_changes=CategorizerConfig(sql=AutoCategorizationMode.SEMI)),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)
+
+# A stateful DuckDB config.
+local_config = Config(
+    default_connection=DuckDBConnectionConfig(database=f"{DATA_DIR}/local.duckdb"),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)
+
+airflow_config = Config(
+    default_scheduler=AirflowSchedulerConfig(),
+    gateways=GatewayConfig(
+        connection=SparkConnectionConfig(
+            config_dir=os.path.join(CURRENT_FILE_PATH, "..", "airflow", "spark_conf"),
+            config={
+                "spark.hadoop.javax.jdo.option.ConnectionURL": "jdbc:postgresql://localhost:5432/metastore_db"
+            },
+        )
+    ),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb", storage_format="iceberg"),
+)
+
+
+airflow_config_docker = Config(
+    default_scheduler=AirflowSchedulerConfig(airflow_url="http://airflow-webserver:8080/"),
+    gateways=GatewayConfig(connection=SparkConnectionConfig()),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb", storage_format="iceberg"),
+)
+
+# A DuckDB config with a physical schema map.
+map_config = Config(
+    default_connection=DuckDBConnectionConfig(),
+    physical_schema_override={"sushi": "company_internal"},
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)
+
+
+required_approvers_config = Config(
+    default_connection=DuckDBConnectionConfig(),
+    users=[
+        User(
+            username="admin",
+            roles=[UserRole.REQUIRED_APPROVER],
+            notification_targets=[
+                SlackApiNotificationTarget(
+                    notify_on=["apply_start", "apply_failure", "apply_end", "audit_failure"],
+                    token=os.getenv("ADMIN_SLACK_API_TOKEN"),
+                    channel="UXXXXXXXXX",  # User's Slack member ID
+                ),
+            ],
+        )
+    ],
+    notification_targets=[
+        SlackWebhookNotificationTarget(
+            notify_on=["apply_start", "apply_failure", "run_start"],
+            url=os.getenv("SLACK_WEBHOOK_URL"),
+        ),
+        BasicSMTPNotificationTarget(
+            notify_on=["run_failure"],
+            host=os.getenv("SMTP_HOST"),
+            user=os.getenv("SMTP_USER"),
+            password=os.getenv("SMTP_PASSWORD"),
+            sender="sushi@example.com",
+            recipients=[
+                "team@example.com",
+            ],
+        ),
+    ],
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)
+
+
+environment_suffix_config = Config(
+    default_connection=DuckDBConnectionConfig(),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    environment_suffix_target=EnvironmentSuffixTarget.TABLE,
+)
+
+
+CATALOGS = {
+    "in_memory": ":memory:",
+    "other_catalog": ":memory:",
+}
+
+local_catalogs = Config(
+    default_connection=DuckDBConnectionConfig(catalogs=CATALOGS),
+    default_test_connection=DuckDBConnectionConfig(catalogs=CATALOGS),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+)
+
+environment_catalog_mapping_config = Config(
+    default_connection=DuckDBConnectionConfig(
+        catalogs={
+            "physical": ":memory:",
+            "prod_catalog": ":memory:",
+            "dev_catalog": ":memory:",
+        }
+    ),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    environment_suffix_target=EnvironmentSuffixTarget.TABLE,
+    environment_catalog_mapping={
+        "^prod$": "prod_catalog",
+        ".*": "dev_catalog",
+    },
+)
+
+postgresql_config = Config(
+    # default_connection=DuckDBConnectionConfig(database=f"{DATA_DIR}/local.duckdb"),
+    default_connection=DuckDBConnectionConfig(database="gclouddata.duckdb"),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    gateways={
+        "my_gateway": GatewayConfig(
+            state_connection=PostgresConnectionConfig(
+                host="35.222.243.76",
+                port=5432,
+                user="unique-glider",
+                password="1d1dWu[NRYV5XUY3",
+                database="sushi-demo",
+            ),
+        ),
+    },
+)

package/recce-source/integration_tests/sqlmesh/helper.py

@@ -0,0 +1,20 @@
+import random
+import typing as t
+from datetime import datetime, timedelta
+
+import numpy as np
+
+
+def set_seed(dt: datetime) -> None:
+    """Helper function to set a seed so random choices are not so random."""
+    ts = int(dt.timestamp())
+    random.seed(ts)
+    np.random.seed(ts)
+
+
+def iter_dates(start: datetime, end: datetime) -> t.Generator[datetime, None, None]:
+    """Iterate through dates and set a deterministic seed."""
+    for i in range((end - start).days + 1):
+        dt = start + timedelta(days=i)
+        set_seed(dt)
+        yield dt

package/recce-source/integration_tests/sqlmesh/macros/macros.py

@@ -0,0 +1,8 @@
+from macros.utils import between  # type: ignore
+
+from sqlmesh import macro
+
+
+@macro()
+def incremental_by_ds(evaluator, column):
+    return between(evaluator, column, evaluator.locals["start_date"], evaluator.locals["end_date"])

package/recce-source/integration_tests/sqlmesh/macros/macros.sql

@@ -0,0 +1,8 @@
+{% macro alias(expr, alias) %}
+  {{ expr }} AS {{ alias }}
+{% endmacro %}
+
+
+{% macro identity(y) %}
+  {{ y }}
+{% endmacro %}

package/recce-source/integration_tests/sqlmesh/macros/utils.py

@@ -0,0 +1,11 @@
+from sqlmesh import macro
+
+
+@macro()
+def add_one(evaluator, column):
+    return f"{column} + 1"
+
+
+@macro()
+def between(evaluator, value, start, end):
+    return value.between(start, end)

package/recce-source/integration_tests/sqlmesh/metrics/metrics.sql

@@ -0,0 +1,25 @@
+METRIC (
+  name total_ordered_items,
+  owner jen,
+  expression SUM(sushi.order_items.quantity)
+);
+
+METRIC (
+  name total_orders,
+  owner jen,
+  expression COUNT(sushi.orders.id)
+);
+
+
+METRIC (
+  name total_orders_from_active_customers,
+  owner jen,
+  expression COUNT(IF(sushi.customers.status = 'ACTIVE', sushi.orders.id, NULL))
+);
+
+
+METRIC (
+  name items_per_order,
+  owner jen,
+  expression total_ordered_items / total_orders
+);

package/recce-source/integration_tests/sqlmesh/models/customer_revenue_by_day.sql

@@ -0,0 +1,41 @@
+/* Table of revenue from customers by day. */
+MODEL (
+  name sushi.customer_revenue_by_day,
+  kind incremental_by_time_range (
+    time_column event_date,
+    batch_size 10,
+  ),
+  owner jen,
+  cron '@daily',
+  dialect hive,
+  tags expensive,
+  grain (customer_id, event_date),
+);
+
+WITH order_total AS (
+  SELECT
+    oi.order_id AS order_id,
+    SUM(oi.quantity * i.price) AS total,
+    oi.event_date AS event_date
+  FROM sushi.order_items AS oi
+  LEFT JOIN sushi.items AS i
+    ON oi.item_id = i.id AND oi.event_date = i.event_date
+  WHERE
+    oi.event_date BETWEEN CAST('{{ start_ds }}' as DATE) AND CAST('{{ end_ds }}' as DATE)
+  GROUP BY
+    oi.order_id,
+    oi.event_date
+)
+SELECT
+  o.customer_id::INT AS customer_id, /* Customer id */
+  SUM(ot.total)::DOUBLE AS revenue, /* Revenue from orders made by this customer */
+  MAX(0) AS "country code", /* Customer country code, used for testing spaces */
+  o.event_date::DATE AS event_date /* Date */
+FROM sushi.orders AS o
+LEFT JOIN order_total AS ot
+  ON o.id = ot.order_id AND o.event_date = ot.event_date
+WHERE
+  o.event_date BETWEEN CAST('{{ start_ds }}' as DATE) AND CAST('{{ end_ds }}' as DATE)
+GROUP BY
+  o.customer_id,
+  o.event_date

package/recce-source/integration_tests/sqlmesh/models/customer_revenue_lifetime.sql

@@ -0,0 +1,60 @@
+/*
+    Table of lifetime customer revenue.
+    Date is available to get lifetime value up to a certain date.
+    Use latest date to get current lifetime value.
+*/
+MODEL (
+  name sushi.customer_revenue_lifetime,
+  kind incremental_by_time_range (
+    time_column event_date,
+    batch_size 1
+  ),
+  owner jen,
+  cron '@daily',
+  dialect hive,
+  tags expensive,
+  columns (
+    customer_id INT,
+    revenue DOUBLE,
+    event_date DATE
+  ),
+  grain (customer_id, event_date),
+);
+
+WITH order_total AS (
+  SELECT
+    oi.order_id AS order_id,
+    SUM(oi.quantity * i.price) AS total
+  FROM sushi.order_items AS oi
+  LEFT JOIN sushi.items AS i
+    ON oi.item_id = i.id AND oi.event_date = i.event_date
+  WHERE
+    oi.event_date = @end_date
+  GROUP BY
+    oi.order_id
+), incremental_total AS (
+  SELECT
+    o.customer_id::INT AS customer_id,
+    SUM(ot.total)::DOUBLE AS revenue,
+  FROM sushi.orders AS o
+  LEFT JOIN order_total AS ot
+    ON o.id = ot.order_id
+  WHERE
+    o.event_date = @end_date
+  GROUP BY
+    o.customer_id
+), prev_total AS (
+  SELECT
+    crl.customer_id,
+    crl.revenue
+  FROM sushi.customer_revenue_lifetime AS crl
+  WHERE
+    crl.event_date = @end_date - INTERVAL 1 DAY
+)
+SELECT
+  COALESCE(it.customer_id, prev_total.customer_id) AS customer_id, /* Customer id */
+  COALESCE(it.revenue, 0) + COALESCE(prev_total.revenue, 0) AS revenue, /* Lifetime revenue from this customer */
+  @end_date AS event_date /* End date of the lifetime calculation */
+FROM incremental_total AS it
+FULL OUTER JOIN prev_total AS prev_total
+  ON it.customer_id = prev_total.customer_id

package/recce-source/integration_tests/sqlmesh/models/customers.sql

@@ -0,0 +1,32 @@
+MODEL (
+  name sushi.customers,
+  kind FULL,
+  owner jen,
+  cron '@daily',
+  tags (pii, fact),
+  grain customer_id,
+  description 'Sushi customer data'
+);
+
+CREATE SCHEMA IF NOT EXISTS raw;
+DROP TABLE IF EXISTS raw.demographics;
+CREATE TABLE raw.demographics AS (
+  SELECT 1 AS customer_id, '00000' AS zip
+);
+
+WITH current_marketing AS (
+  SELECT
+    customer_id,
+    status
+  FROM sushi.marketing
+  WHERE valid_to is null
+)
+SELECT DISTINCT
+  o.customer_id::INT AS customer_id, -- customer_id uniquely identifies customers
+  m.status,
+  d.zip
+  FROM sushi.orders AS o
+LEFT JOIN current_marketing AS m
+  ON o.customer_id = m.customer_id
+  LEFT JOIN raw.demographics AS d
+  ON o.customer_id = d.customer_id

package/recce-source/integration_tests/sqlmesh/models/items.py

@@ -0,0 +1,95 @@
+import random
+import typing as t
+from datetime import datetime
+
+import numpy as np
+import pandas as pd
+from helper import iter_dates  # type: ignore
+from sqlglot.expressions import to_column
+
+from sqlmesh import ExecutionContext, model
+from sqlmesh.core.model import IncrementalByTimeRangeKind
+from sqlmesh.utils.date import to_date
+
+ITEMS = [
+    "Ahi",
+    "Aji",
+    "Amaebi",
+    "Anago",
+    "Aoyagi",
+    "Bincho",
+    "Katsuo",
+    "Ebi",
+    "Escolar",
+    "Hamachi",
+    "Hamachi Toro",
+    "Hirame",
+    "Hokigai",
+    "Hotate",
+    "Ika",
+    "Ikura",
+    "Iwashi",
+    "Kani",
+    "Kanpachi",
+    "Maguro",
+    "Saba",
+    "Sake",
+    "Sake Toro",
+    "Tai",
+    "Tako",
+    "Tamago",
+    "Tobiko",
+    "Toro",
+    "Tsubugai",
+    "Umi Masu",
+    "Unagi",
+    "Uni",
+]
+
+
+@model(
+    "sushi.items",
+    kind=IncrementalByTimeRangeKind(time_column="event_date", batch_size=30),
+    start="1 week ago",
+    cron="@daily",
+    columns={
+        "id": "int",
+        "name": "text",
+        "price": "double",
+        "event_date": "date",
+    },
+    audits=[
+        ("accepted_values", {"column": to_column("name"), "is_in": ITEMS}),
+        ("not_null", {"columns": [to_column("name"), to_column("price")]}),
+        ("assert_items_price_exceeds_threshold", {"price": 0}),
+    ],
+    session_properties={
+        "string_prop": "some_value",
+        "int_prop": 1,
+        "float_prop": 1.0,
+        "bool_prop": True,
+    },
+)
+def execute(
+    context: ExecutionContext,
+    start: datetime,
+    end: datetime,
+    execution_time: datetime,
+    **kwargs: t.Any,
+) -> pd.DataFrame:
+    dfs = []
+    for dt in iter_dates(start, end):
+        num_items = random.randint(10, len(ITEMS))
+        dfs.append(
+            pd.DataFrame(
+                {
+                    "name": random.sample(ITEMS, num_items),
+                    "price": np.random.uniform(3.0, 10.0, size=num_items).round(2),
+                    "event_date": to_date(dt),
+                }
+            )
+            .reset_index()
+            .rename(columns={"index": "id"})
+        )
+
+    return pd.concat(dfs)

package/recce-source/integration_tests/sqlmesh/models/marketing.sql

@@ -0,0 +1,15 @@
+MODEL (
+  name sushi.marketing,
+  kind SCD_TYPE_2(unique_key customer_id),
+  owner jen,
+  cron '@daily',
+  grain customer_id,
+  description 'Sushi marketing data'
+);
+
+SELECT
+    customer_id::INT AS customer_id, -- customer_id uniquely identifies customers \
+    status::TEXT AS status,
+    updated_at::TIMESTAMP AS updated_at
+FROM
+    sushi.raw_marketing

package/recce-source/integration_tests/sqlmesh/models/order_items.py

@@ -0,0 +1,95 @@
+import random
+import typing as t
+from datetime import datetime
+
+import numpy as np
+import pandas as pd
+from helper import iter_dates  # type: ignore
+from sqlglot import exp
+from sqlglot.expressions import to_column
+
+from sqlmesh import ExecutionContext, model
+from sqlmesh.core.model import IncrementalByTimeRangeKind
+from sqlmesh.utils.date import to_date, to_ds
+
+ITEMS = "sushi.items"
+
+
+def get_items_table(context: ExecutionContext) -> str:
+    return context.table(ITEMS)
+
+
+@model(
+    "sushi.order_items",
+    kind=IncrementalByTimeRangeKind(
+        time_column="event_date",
+        batch_size=30,
+    ),
+    cron="@daily",
+    columns={
+        "id": "int",
+        "order_id": "int",
+        "item_id": "int",
+        "quantity": "int",
+        "event_date": "date",
+    },
+    audits=[
+        (
+            "NOT_NULL",
+            {"columns": [to_column(c) for c in ("id", "order_id", "item_id", "quantity")]},
+        ),
+        ("assert_order_items_quantity_exceeds_threshold", {"quantity": 0}),
+    ],
+)
+def execute(
+    context: ExecutionContext,
+    start: datetime,
+    end: datetime,
+    execution_time: datetime,
+    **kwargs: t.Any,
+) -> t.Generator[pd.DataFrame, None, None]:
+    orders_table = context.table("sushi.orders")
+    items_table = get_items_table(context)
+
+    rng = random.Random()
+
+    for dt in iter_dates(start, end):
+        rng.seed(int(dt.timestamp()))
+
+        # Generate query with sqlglot dialect/quoting
+        orders = context.fetchdf(
+            exp.select("*").from_(orders_table).where(f"event_date = CAST('{to_ds(dt)}' AS DATE)"),
+            quote_identifiers=True,
+        )
+
+        # Normalize column names to support Snowflake.
+        orders = orders.rename(columns={col: col.lower() for col in orders.columns})  # type: ignore
+
+        # Generate query with sqlglot dialect/quoting
+        items = context.fetchdf(
+            exp.select("*").from_(items_table).where(f"event_date = CAST('{to_ds(dt)}' AS DATE)"),
+            quote_identifiers=True,
+        )
+
+        # Normalize column names to support Snowflake.
+        items = items.rename(columns={col: col.lower() for col in items.columns})  # type: ignore
+
+        dfs = []
+
+        for order_id in orders["id"]:
+            n = rng.randint(1, 5)
+
+            dfs.append(
+                pd.DataFrame(
+                    {
+                        "order_id": order_id,
+                        "item_id": items.sample(n=n)["id"],
+                        "quantity": np.random.randint(1, 10, n),
+                        "event_date": to_date(dt),
+                    }
+                )
+                .reset_index()
+                .rename(columns={"index": "id"})
+            )
+
+        yield pd.concat(dfs)

package/recce-source/integration_tests/sqlmesh/models/orders.py

@@ -0,0 +1,70 @@
+import random
+import typing as t
+from datetime import datetime, timedelta
+
+import pandas as pd
+from helper import iter_dates  # type: ignore
+
+from sqlmesh import ExecutionContext, model
+from sqlmesh.core.model import IncrementalByTimeRangeKind
+from sqlmesh.utils.date import to_date
+
+CUSTOMERS = list(range(0, 100))
+WAITERS = list(range(0, 10))
+
+
+@model(
+    "sushi.orders",
+    description="Table of sushi orders.",
+    kind=IncrementalByTimeRangeKind(time_column="event_date", batch_size=30),
+    start="1 week ago",
+    cron="@daily",
+    grains=[
+        "id AS order_id",
+    ],
+    references=[
+        "customer_id",
+        "waiter_id",
+    ],
+    columns={
+        "id": "int",
+        "customer_id": "int",
+        "waiter_id": "int",
+        "start_ts": "int",
+        "end_ts": "int",
+        "event_date": "date",
+    },
+)
+def execute(
+    context: ExecutionContext,
+    start: datetime,
+    end: datetime,
+    execution_time: datetime,
+    **kwargs: t.Any,
+) -> pd.DataFrame:
+    dfs = []
+    for dt in iter_dates(start, end):
+        num_orders = random.randint(10, 30)
+
+        start_ts = [
+            int((dt + timedelta(seconds=random.randint(0, 80000))).timestamp())
+            for _ in range(num_orders)
+        ]
+
+        end_ts = [int(s + random.randint(0, 60 * 60)) for s in start_ts]
+
+        dfs.append(
+            pd.DataFrame(
+                {
+                    "customer_id": random.choices(CUSTOMERS, k=num_orders),
+                    "waiter_id": random.choices(WAITERS, k=num_orders),
+                    "start_ts": start_ts,
+                    "end_ts": end_ts,
+                    "event_date": to_date(dt),
+                }
+            )
+            .reset_index()
+            .rename(columns={"index": "id"})
+        )
+
+    return pd.concat(dfs)