Skip to content

Commit ac090d6

Browse files
izeigermanizeigerman
committed
add an integration test
1 parent ee96e5b commit ac090d6

File tree

5 files changed

+153
-3
lines changed

5 files changed

+153
-3
lines changed

‎examples/sushi/config.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import os
22

3+
from sqlmesh.core.config.common import VirtualEnvironmentMode
34
from sqlmesh.core.config import (
45
AutoCategorizationMode,
56
BigQueryConnectionConfig,
@@ -76,6 +77,16 @@
7677
model_defaults=model_defaults,
7778
)
7879

80+
# A configuration used for SQLMesh tests with virtual environment mode set to DEV_ONLY.
81+
test_config_virtual_environment_mode_dev_only = test_config.copy(
82+
update={
83+
"virtual_environment_mode": VirtualEnvironmentMode.DEV_ONLY,
84+
"plan": PlanConfig(
85+
auto_categorize_changes=CategorizerConfig.all_full(),
86+
),
87+
}
88+
)
89+
7990
# A DuckDB config with a physical schema map.
8091
map_config = Config(
8192
default_connection=DuckDBConnectionConfig(),

‎sqlmesh/dbt/basemodel.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from sqlmesh.utils.pydantic import field_validator
3131

3232
if t.TYPE_CHECKING:
33+
from sqlmesh.core.config.common import VirtualEnvironmentMode
3334
from sqlmesh.core.audit.definition import ModelAudit
3435
from sqlmesh.dbt.context import DbtContext
3536

@@ -345,6 +346,9 @@ def sqlmesh_model_kwargs(
345346

346347
@abstractmethod
347348
def to_sqlmesh(
348-
self, context: DbtContext, audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None
349+
self,
350+
context: DbtContext,
351+
audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
352+
virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
349353
) -> Model:
350354
"""Convert DBT model into sqlmesh Model"""

‎sqlmesh/dbt/model.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from sqlmesh.utils.pydantic import field_validator
3030

3131
if t.TYPE_CHECKING:
32+
from sqlmesh.core.config.common import VirtualEnvironmentMode
3233
from sqlmesh.core.audit.definition import ModelAudit
3334
from sqlmesh.dbt.context import DbtContext
3435

@@ -421,7 +422,10 @@ def sqlmesh_config_fields(self) -> t.Set[str]:
421422
}
422423

423424
def to_sqlmesh(
424-
self, context: DbtContext, audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None
425+
self,
426+
context: DbtContext,
427+
audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
428+
virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
425429
) -> Model:
426430
"""Converts the dbt model into a SQLMesh model."""
427431
model_dialect = self.dialect(context)
@@ -573,6 +577,7 @@ def to_sqlmesh(
573577
# Note: any table dependencies that are not referenced using the `ref` macro will not be included.
574578
extract_dependencies_from_query=False,
575579
allow_partials=allow_partials,
580+
virtual_environment_mode=virtual_environment_mode,
576581
**optional_kwargs,
577582
**model_kwargs,
578583
)

‎sqlmesh/dbt/seed.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from sqlmesh.dbt.column import ColumnConfig
2121

2222
if t.TYPE_CHECKING:
23+
from sqlmesh.core.config.common import VirtualEnvironmentMode
2324
from sqlmesh.core.audit.definition import ModelAudit
2425
from sqlmesh.dbt.context import DbtContext
2526

@@ -38,7 +39,10 @@ class SeedConfig(BaseModelConfig):
3839
quote_columns: t.Optional[bool] = False
3940

4041
def to_sqlmesh(
41-
self, context: DbtContext, audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None
42+
self,
43+
context: DbtContext,
44+
audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
45+
virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
4246
) -> Model:
4347
"""Converts the dbt seed into a SQLMesh model."""
4448
seed_path = self.path.absolute().as_posix()
@@ -83,6 +87,7 @@ def to_sqlmesh(
8387
SeedKind(path=seed_path),
8488
dialect=self.dialect(context),
8589
audit_definitions=audit_definitions,
90+
virtual_environment_mode=virtual_environment_mode,
8691
**kwargs,
8792
)
8893

‎tests/core/test_integration.py

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2448,6 +2448,131 @@ def test_unaligned_start_snapshot_with_non_deployable_downstream(init_and_plan_c
24482448
assert snapshot_interval.intervals[0][0] == to_timestamp("2023-01-07")
24492449

24502450

2451+
@time_machine.travel("2023-01-08 15:00:00 UTC")
2452+
def test_virtual_environment_mode_dev_only(init_and_plan_context: t.Callable):
2453+
context, _ = init_and_plan_context(
2454+
"examples/sushi", config="test_config_virtual_environment_mode_dev_only"
2455+
)
2456+
2457+
assert all(
2458+
s.virtual_environment_mode.is_dev_only or not s.is_model or s.is_symbolic
2459+
for s in context.snapshots.values()
2460+
)
2461+
2462+
# Init prod
2463+
context.plan("prod", auto_apply=True, no_prompts=True)
2464+
2465+
# Make a change in dev
2466+
original_model = context.get_model("sushi.waiter_revenue_by_day")
2467+
original_fingerprint = context.get_snapshot(original_model.name).fingerprint
2468+
model = original_model.copy(update={"query": original_model.query.order_by("waiter_id")})
2469+
model = add_projection_to_model(t.cast(SqlModel, model))
2470+
context.upsert_model(model)
2471+
2472+
plan_dev = context.plan_builder("dev").build()
2473+
assert to_timestamp(plan_dev.start) == to_timestamp("2023-01-07")
2474+
assert plan_dev.requires_backfill
2475+
assert plan_dev.missing_intervals == [
2476+
SnapshotIntervals(
2477+
snapshot_id=context.get_snapshot("sushi.top_waiters").snapshot_id,
2478+
intervals=[(to_timestamp("2023-01-07"), to_timestamp("2023-01-08"))],
2479+
),
2480+
SnapshotIntervals(
2481+
snapshot_id=context.get_snapshot("sushi.waiter_revenue_by_day").snapshot_id,
2482+
intervals=[(to_timestamp("2023-01-07"), to_timestamp("2023-01-08"))],
2483+
),
2484+
]
2485+
context.apply(plan_dev)
2486+
2487+
# Make sure the waiter_revenue_by_day model is a table in prod and a view in dev
2488+
table_types_df = context.engine_adapter.fetchdf(
2489+
"SELECT table_schema, table_type FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'waiter_revenue_by_day'"
2490+
)
2491+
assert table_types_df.to_dict("records") == [
2492+
{"table_schema": "sushi", "table_type": "BASE TABLE"},
2493+
{"table_schema": "sushi__dev", "table_type": "VIEW"},
2494+
]
2495+
2496+
# Check that the specified dates were backfilled
2497+
min_event_date = context.engine_adapter.fetchone(
2498+
"SELECT MIN(event_date) FROM sushi__dev.waiter_revenue_by_day"
2499+
)[0]
2500+
assert min_event_date == to_date("2023-01-07")
2501+
2502+
# Make sure the changed models are fully rebuilt when deploying to prod
2503+
plan_prod = context.plan_builder("prod").build()
2504+
assert plan_prod.requires_backfill
2505+
assert plan_prod.missing_intervals == [
2506+
SnapshotIntervals(
2507+
snapshot_id=context.get_snapshot("sushi.top_waiters").snapshot_id,
2508+
intervals=[
2509+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
2510+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
2511+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
2512+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
2513+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
2514+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
2515+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
2516+
],
2517+
),
2518+
SnapshotIntervals(
2519+
snapshot_id=context.get_snapshot("sushi.waiter_revenue_by_day").snapshot_id,
2520+
intervals=[
2521+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
2522+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
2523+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
2524+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
2525+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
2526+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
2527+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
2528+
],
2529+
),
2530+
]
2531+
context.apply(plan_prod)
2532+
assert "one" in context.engine_adapter.columns("sushi.waiter_revenue_by_day")
2533+
assert (
2534+
context.engine_adapter.fetchone(
2535+
"SELECT COUNT(*) FROM sushi.waiter_revenue_by_day WHERE one is NULL"
2536+
)[0]
2537+
== 0
2538+
)
2539+
2540+
# Make sure the revert of a breaking changes results in a full rebuild
2541+
context.upsert_model(original_model)
2542+
assert context.get_snapshot(original_model.name).fingerprint == original_fingerprint
2543+
2544+
plan_prod = context.plan_builder("prod").build()
2545+
assert plan_prod.requires_backfill
2546+
assert plan_prod.missing_intervals == [
2547+
SnapshotIntervals(
2548+
snapshot_id=context.get_snapshot("sushi.top_waiters").snapshot_id,
2549+
intervals=[
2550+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
2551+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
2552+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
2553+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
2554+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
2555+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
2556+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
2557+
],
2558+
),
2559+
SnapshotIntervals(
2560+
snapshot_id=context.get_snapshot("sushi.waiter_revenue_by_day").snapshot_id,
2561+
intervals=[
2562+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
2563+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
2564+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
2565+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
2566+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
2567+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
2568+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
2569+
],
2570+
),
2571+
]
2572+
context.apply(plan_prod)
2573+
assert "one" not in context.engine_adapter.columns("sushi.waiter_revenue_by_day")
2574+
2575+
24512576
@time_machine.travel("2023-01-08 15:00:00 UTC")
24522577
def test_restatement_plan_ignores_changes(init_and_plan_context: t.Callable):
24532578
context, plan = init_and_plan_context("examples/sushi")

0 commit comments

Comments
 (0)