Feat: Add 'state_schema_naming_pattern' to infer the state schema per dbt target

Author: erindru

sqlmesh/cli/project_init.py

@@ -116,7 +116,13 @@ def _gen_config(
     - invalidselectstarexpansion
     - noambiguousprojections
 """,
-        ProjectTemplate.DBT: f"""# --- Virtual Data Environment Mode ---
+        ProjectTemplate.DBT: f"""# --- State ---
+# This default configuration ensures that each dbt target gets its own isolated state.
+# If this is undesirable, you may configure the state connection manually.
+# https://sqlmesh.readthedocs.io/en/stable/integrations/dbt/?h=dbt#selecting-a-different-state-connection
+state_schema_naming_pattern: sqlmesh_state_@{{dbt_profile_name}}_@{{dbt_target_name}}
+
+# --- Virtual Data Environment Mode ---
 # Enable Virtual Data Environments (VDE) for *development* environments.
 # Note that the production environment in dbt projects is not virtual by default to maintain compatibility with existing tooling.
 # https://sqlmesh.readthedocs.io/en/stable/guides/configuration/#virtual-data-environment-modes

sqlmesh/core/config/__init__.py

@@ -36,6 +36,6 @@
 from sqlmesh.core.config.naming import NameInferenceConfig as NameInferenceConfig
 from sqlmesh.core.config.linter import LinterConfig as LinterConfig
 from sqlmesh.core.config.plan import PlanConfig as PlanConfig
-from sqlmesh.core.config.root import Config as Config
+from sqlmesh.core.config.root import Config as Config, DbtConfigInfo as DbtConfigInfo
 from sqlmesh.core.config.run import RunConfig as RunConfig
 from sqlmesh.core.config.scheduler import BuiltInSchedulerConfig as BuiltInSchedulerConfig

sqlmesh/core/config/base.py

@@ -140,3 +140,17 @@ def update_with(self: T, other: t.Union[t.Dict[str, t.Any], T]) -> T:
             setattr(updated, field, value)
 
         return updated
+
+
+class DbtConfigInfo(PydanticModel):
+    """
+    This is like DbtNodeInfo except it applies to config instead of DAG nodes.
+
+    It's intended to capture information from a dbt project loaded by the DbtLoader so that it can be used for things like
+    variable substitutions in regular project config.
+    """
+
+    profile_name: str
+    """Which profile in the dbt project is being used"""
+    target_name: str
+    """Which target of the specified profile is being used"""

sqlmesh/core/config/gateway.py

@@ -2,7 +2,6 @@
 
 import typing as t
 
-from sqlmesh.core import constants as c
 from sqlmesh.core.config.base import BaseConfig
 from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.config.common import variables_validator
@@ -33,7 +32,7 @@ class GatewayConfig(BaseConfig):
     state_connection: t.Optional[SerializableConnectionConfig] = None
     test_connection: t.Optional[SerializableConnectionConfig] = None
     scheduler: t.Optional[SchedulerConfig] = None
-    state_schema: t.Optional[str] = c.SQLMESH
+    state_schema: t.Optional[str] = None
     variables: t.Dict[str, t.Any] = {}
     model_defaults: t.Optional[ModelDefaultsConfig] = None
 

sqlmesh/core/config/root.py

@@ -4,6 +4,7 @@
 import re
 import typing as t
 import zlib
+import logging
 
 from pydantic import Field
 from pydantic.functional_validators import BeforeValidator
@@ -19,7 +20,7 @@
     TableNamingConvention,
     VirtualEnvironmentMode,
 )
-from sqlmesh.core.config.base import BaseConfig, UpdateStrategy
+from sqlmesh.core.config.base import BaseConfig, UpdateStrategy, DbtConfigInfo
 from sqlmesh.core.config.common import variables_validator, compile_regex_mapping
 from sqlmesh.core.config.connection import (
     ConnectionConfig,
@@ -49,6 +50,8 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import model_validator
 
+logger = logging.getLogger(__name__)
+
 
 def validate_no_past_ttl(v: str) -> str:
     current_time = now()
@@ -96,6 +99,8 @@ class Config(BaseConfig):
         default_test_connection: The default connection to use for tests if one is not specified in a gateway.
         default_scheduler: The default scheduler configuration to use if one is not specified in a gateway.
         default_gateway: The default gateway.
+        state_schema_naming_pattern: A pattern supporting variable substitutions to determine the state schema name, rather than just using 'sqlmesh'.
+            Only applies when the state schema is not explicitly set in the gateway config
         notification_targets: The notification targets to use.
         project: The project name of this config. Used for multi-repo setups.
         snapshot_ttl: The period of time that a model snapshot that is not a part of any environment should exist before being deleted.
@@ -128,6 +133,7 @@ class Config(BaseConfig):
         before_all: SQL statements or macros to be executed at the start of the `sqlmesh plan` and `sqlmesh run` commands.
         after_all: SQL statements or macros to be executed at the end of the `sqlmesh plan` and `sqlmesh run` commands.
         cache_dir: The directory to store the SQLMesh cache. Defaults to .cache in the project folder.
+        dbt_config_info: Dbt-specific properties (such as profile and target) for dbt projects loaded by the dbt loader
     """
 
     gateways: GatewayDict = {"": GatewayConfig()}
@@ -137,6 +143,7 @@ class Config(BaseConfig):
     )
     default_scheduler: SchedulerConfig = BuiltInSchedulerConfig()
     default_gateway: str = ""
+    state_schema_naming_pattern: t.Optional[str] = None
     notification_targets: t.List[NotificationTarget] = []
     project: str = ""
     snapshot_ttl: NoPastTTLString = c.DEFAULT_SNAPSHOT_TTL
@@ -173,6 +180,7 @@ class Config(BaseConfig):
     linter: LinterConfig = LinterConfig()
     janitor: JanitorConfig = JanitorConfig()
     cache_dir: t.Optional[str] = None
+    dbt_config_info: t.Optional[DbtConfigInfo] = None
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
         "gateways": UpdateStrategy.NESTED_UPDATE,
@@ -344,8 +352,27 @@ def get_test_connection(
     def get_scheduler(self, gateway_name: t.Optional[str] = None) -> SchedulerConfig:
         return self.get_gateway(gateway_name).scheduler or self.default_scheduler
 
-    def get_state_schema(self, gateway_name: t.Optional[str] = None) -> t.Optional[str]:
-        return self.get_gateway(gateway_name).state_schema
+    def get_state_schema(self, gateway_name: t.Optional[str] = None) -> str:
+        state_schema = self.get_gateway(gateway_name).state_schema
+
+        if state_schema is None and self.state_schema_naming_pattern:
+            substitutions = {}
+            if dbt := self.dbt_config_info:
+                # TODO: keeping this simple for now rather than trying to set up a Jinja or SQLMesh Macro rendering context
+                substitutions.update(
+                    {
+                        "@{dbt_profile_name}": dbt.profile_name,
+                        # TODO @iaroslav: what was the problem with using target name instead of the default schema name again?
+                        "@{dbt_target_name}": dbt.target_name,
+                    }
+                )
+            state_schema = self.state_schema_naming_pattern
+            for pattern, value in substitutions.items():
+                state_schema = state_schema.replace(pattern, value)
+
+            logger.info("Inferring state schema: %s", state_schema)
+
+        return state_schema or c.SQLMESH
 
     @property
     def default_gateway_name(self) -> str:

sqlmesh/dbt/loader.py

@@ -11,6 +11,7 @@
     ConnectionConfig,
     GatewayConfig,
     ModelDefaultsConfig,
+    DbtConfigInfo,
 )
 from sqlmesh.core.environment import EnvironmentStatements
 from sqlmesh.core.loader import CacheBase, LoadedProject, Loader
@@ -67,10 +68,19 @@ def sqlmesh_config(
     if not issubclass(loader, DbtLoader):
         raise ConfigError("The loader must be a DbtLoader.")
 
+    if context.profile_name is None:
+        # Note: Profile.load() mutates `context` and will have already raised an exception if profile_name is not set,
+        # but mypy doesnt know this because the field is defined as t.Optional[str]
+        raise ConfigError(f"profile name must be set")
+
     return Config(
         loader=loader,
         model_defaults=model_defaults,
         variables=variables or {},
+        dbt_config_info=DbtConfigInfo(
+            profile_name=dbt_profile_name or context.profile_name,
+            target_name=dbt_target_name or profile.target_name,
+        ),
         **{
             "default_gateway": profile.target_name if "gateways" not in kwargs else "",
             "gateways": {

tests/dbt/test_config.py

@@ -15,6 +15,7 @@
 from sqlmesh.core.dialect import jinja_query
 from sqlmesh.core.model import SqlModel
 from sqlmesh.core.model.kind import OnDestructiveChange, OnAdditiveChange
+from sqlmesh.core.state_sync import CachingStateSync, EngineAdapterStateSync
 from sqlmesh.dbt.builtin import Api
 from sqlmesh.dbt.column import ColumnConfig
 from sqlmesh.dbt.common import Dependencies
@@ -46,7 +47,8 @@
 )
 from sqlmesh.dbt.test import TestConfig
 from sqlmesh.utils.errors import ConfigError
-from sqlmesh.utils.yaml import load as yaml_load
+from sqlmesh.utils.yaml import load as yaml_load, dump as yaml_dump
+from tests.dbt.conftest import EmptyProjectCreator
 
 pytestmark = pytest.mark.dbt
 
@@ -1211,3 +1213,35 @@ def test_empty_vars_config(tmp_path):
     # Verify the variables are empty (not causing any issues)
     assert project.packages["test_empty_vars"].variables == {}
     assert project.context.variables == {}
+
+
+def test_state_schema_naming_pattern(create_empty_project: EmptyProjectCreator):
+    project_dir, _ = create_empty_project("test_foo", "dev")
+
+    # no state_schema_naming_pattern, creating python config manually doesnt take into account
+    # any config yaml files that may be present, so we get the default state schema
+    config = sqlmesh_config(project_root=project_dir)
+    assert not config.state_schema_naming_pattern
+    assert config.get_state_schema() == "sqlmesh"
+
+    # create_empty_project() uses the default dbt template for sqlmesh yaml config which
+    # sets state_schema_naming_pattern
+    ctx = Context(paths=[project_dir])
+    assert ctx.config.state_schema_naming_pattern
+    assert ctx.config.get_state_schema() == "sqlmesh_state_test_foo_dev"
+    assert isinstance(ctx.state_sync, CachingStateSync)
+    assert isinstance(ctx.state_sync.state_sync, EngineAdapterStateSync)
+    assert ctx.state_sync.state_sync.schema == "sqlmesh_state_test_foo_dev"
+
+    # If the user delberately overrides state_schema then we should respect this choice
+    config_file = project_dir / "sqlmesh.yaml"
+    config_yaml = yaml_load(config_file)
+    config_yaml["gateways"] = {"dev": {"state_schema": "state_override"}}
+    config_file.write_text(yaml_dump(config_yaml))
+
+    ctx = Context(paths=[project_dir])
+    assert ctx.config.state_schema_naming_pattern
+    assert ctx.config.get_state_schema() == "state_override"
+    assert isinstance(ctx.state_sync, CachingStateSync)
+    assert isinstance(ctx.state_sync.state_sync, EngineAdapterStateSync)
+    assert ctx.state_sync.state_sync.schema == "state_override"

tests/dbt/test_integration.py

@@ -19,7 +19,8 @@
 from sqlmesh.core.config.connection import DuckDBConnectionConfig
 from sqlmesh.core.engine_adapter import DuckDBEngineAdapter
 from sqlmesh.utils.pandas import columns_to_types_from_df
-from sqlmesh.utils.yaml import YAML
+from sqlmesh.utils.yaml import YAML, load as yaml_load, dump as yaml_dump
+from sqlmesh_dbt.operations import init_project_if_required
 from tests.utils.pandas import compare_dataframes, create_df
 
 # Some developers had issues with this test freezing locally so we mark it as cicdonly
@@ -604,3 +605,41 @@ def test_dbt_node_info(jaffle_shop_duckdb_context: Context):
         relationship_audit.node.dbt_node_info.name
         == "relationships_orders_customer_id__customer_id__ref_customers_"
     )
+
+
+def test_state_schema_isolation_per_target(jaffle_shop_duckdb: Path):
+    profiles_file = jaffle_shop_duckdb / "profiles.yml"
+
+    profiles_yml = yaml_load(profiles_file)
+
+    # make prod / dev config identical with the exception of a different default schema to simulate using the same warehouse
+    profiles_yml["jaffle_shop"]["outputs"]["prod"] = {
+        **profiles_yml["jaffle_shop"]["outputs"]["dev"]
+    }
+    profiles_yml["jaffle_shop"]["outputs"]["prod"]["schema"] = "prod_schema"
+    profiles_yml["jaffle_shop"]["outputs"]["dev"]["schema"] = "dev_schema"
+
+    profiles_file.write_text(yaml_dump(profiles_yml))
+
+    init_project_if_required(jaffle_shop_duckdb)
+
+    # start off with the prod target
+    prod_ctx = Context(paths=[jaffle_shop_duckdb], config_loader_kwargs={"target": "prod"})
+    assert prod_ctx.config.get_state_schema() == "sqlmesh_state_jaffle_shop_prod"
+    assert all("prod_schema" in fqn for fqn in prod_ctx.models)
+    assert prod_ctx.plan(auto_apply=True).has_changes
+    assert not prod_ctx.plan(auto_apply=True).has_changes
+
+    # dev target should have changes - new state separate from prod
+    dev_ctx = Context(paths=[jaffle_shop_duckdb], config_loader_kwargs={"target": "dev"})
+    assert dev_ctx.config.get_state_schema() == "sqlmesh_state_jaffle_shop_dev"
+    assert all("dev_schema" in fqn for fqn in dev_ctx.models)
+    assert dev_ctx.plan(auto_apply=True).has_changes
+    assert not dev_ctx.plan(auto_apply=True).has_changes
+
+    # no explicitly specified target should use dev because that's what's set for the default in the profiles.yml
+    assert profiles_yml["jaffle_shop"]["target"] == "dev"
+    default_ctx = Context(paths=[jaffle_shop_duckdb])
+    assert default_ctx.config.get_state_schema() == "sqlmesh_state_jaffle_shop_dev"
+    assert all("dev_schema" in fqn for fqn in default_ctx.models)
+    assert not default_ctx.plan(auto_apply=True).has_changes

tests/fixtures/dbt/empty_project/profiles.yml

@@ -3,7 +3,7 @@ empty_project:
   target: __DEFAULT_TARGET__
 
   outputs:
-    duckdb:
+    __DEFAULT_TARGET__:
       type: duckdb
       path: 'empty_project.duckdb'
       threads: 4