PR feedback

Author: erindru

sqlmesh/cli/project_init.py

@@ -116,11 +116,13 @@ def _gen_config(
     - invalidselectstarexpansion
     - noambiguousprojections
 """,
-        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}}
+        ProjectTemplate.DBT: f"""# --- DBT-specific options ---
+dbt:
+  # This configuration ensures that each dbt target gets its own isolated state.
+  # The inferred state schemas are named "sqlmesh_state_<project name>_<target schema>", eg "sqlmesh_state_jaffle_shop_dev"
+  # If this is undesirable, you may manually configure the gateway to use a specific state schema name
+  # https://sqlmesh.readthedocs.io/en/stable/integrations/dbt/#selecting-a-different-state-connection
+  infer_state_schema_name: True
 
 # --- Virtual Data Environment Mode ---
 # Enable Virtual Data Environments (VDE) for *development* environments.

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, DbtConfigInfo as DbtConfigInfo
+from sqlmesh.core.config.root import Config as Config, DbtConfig as DbtConfig
 from sqlmesh.core.config.run import RunConfig as RunConfig
 from sqlmesh.core.config.scheduler import BuiltInSchedulerConfig as BuiltInSchedulerConfig

sqlmesh/core/config/base.py

@@ -140,17 +140,3 @@ 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/dbt.py

@@ -0,0 +1,13 @@
+from sqlmesh.core.config.base import BaseConfig
+
+
+class DbtConfig(BaseConfig):
+    """
+    Represents dbt-specific options on the SQLMesh root config.
+
+    These options are only taken into account for dbt projects and are ignored on native projects
+    """
+
+    infer_state_schema_name: bool = False
+    """If set, indicates to the dbt loader that the state schema should be inferred based on the profile/target
+    so that each target gets its own isolated state"""

sqlmesh/core/config/gateway.py

@@ -2,6 +2,7 @@
 
 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
@@ -32,7 +33,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] = None
+    state_schema: t.Optional[str] = c.SQLMESH
     variables: t.Dict[str, t.Any] = {}
     model_defaults: t.Optional[ModelDefaultsConfig] = None
 

sqlmesh/core/config/loader.py

@@ -172,11 +172,16 @@ def load_config_from_paths(
         if dbt_project_file:
             from sqlmesh.dbt.loader import sqlmesh_config
 
+            infer_state_schema_name = False
+            if dbt := non_python_config.dbt:
+                infer_state_schema_name = dbt.infer_state_schema_name
+
             dbt_python_config = sqlmesh_config(
                 project_root=dbt_project_file.parent,
                 dbt_profile_name=kwargs.pop("profile", None),
                 dbt_target_name=kwargs.pop("target", None),
                 variables=variables,
+                infer_state_schema_name=infer_state_schema_name,
             )
             if type(dbt_python_config) != config_type:
                 dbt_python_config = convert_config_type(dbt_python_config, config_type)

sqlmesh/core/config/root.py

@@ -20,7 +20,7 @@
     TableNamingConvention,
     VirtualEnvironmentMode,
 )
-from sqlmesh.core.config.base import BaseConfig, UpdateStrategy, DbtConfigInfo
+from sqlmesh.core.config.base import BaseConfig, UpdateStrategy
 from sqlmesh.core.config.common import variables_validator, compile_regex_mapping
 from sqlmesh.core.config.connection import (
     ConnectionConfig,
@@ -37,6 +37,7 @@
 from sqlmesh.core.config.linter import LinterConfig as LinterConfig
 from sqlmesh.core.config.plan import PlanConfig
 from sqlmesh.core.config.run import RunConfig
+from sqlmesh.core.config.dbt import DbtConfig
 from sqlmesh.core.config.scheduler import (
     BuiltInSchedulerConfig,
     SchedulerConfig,
@@ -99,8 +100,6 @@ 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.
@@ -133,7 +132,6 @@ 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()}
@@ -143,7 +141,6 @@ 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
@@ -180,7 +177,7 @@ class Config(BaseConfig):
     linter: LinterConfig = LinterConfig()
     janitor: JanitorConfig = JanitorConfig()
     cache_dir: t.Optional[str] = None
-    dbt_config_info: t.Optional[DbtConfigInfo] = None
+    dbt: t.Optional[DbtConfig] = None
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
         "gateways": UpdateStrategy.NESTED_UPDATE,
@@ -199,6 +196,7 @@ class Config(BaseConfig):
         "before_all": UpdateStrategy.EXTEND,
         "after_all": UpdateStrategy.EXTEND,
         "linter": UpdateStrategy.NESTED_UPDATE,
+        "dbt": UpdateStrategy.NESTED_UPDATE,
     }
 
     _connection_config_validator = connection_config_validator
@@ -352,27 +350,8 @@ 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) -> 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
+    def get_state_schema(self, gateway_name: t.Optional[str] = None) -> t.Optional[str]:
+        return self.get_gateway(gateway_name).state_schema
 
     @property
     def default_gateway_name(self) -> str:

sqlmesh/dbt/loader.py

@@ -11,7 +11,7 @@
     ConnectionConfig,
     GatewayConfig,
     ModelDefaultsConfig,
-    DbtConfigInfo,
+    DbtConfig as RootDbtConfig,
 )
 from sqlmesh.core.environment import EnvironmentStatements
 from sqlmesh.core.loader import CacheBase, LoadedProject, Loader
@@ -51,6 +51,7 @@ def sqlmesh_config(
     dbt_target_name: t.Optional[str] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
     register_comments: t.Optional[bool] = None,
+    infer_state_schema_name: bool = False,
     **kwargs: t.Any,
 ) -> Config:
     project_root = project_root or Path()
@@ -68,25 +69,28 @@ 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")
+    gateway_kwargs = {}
+    if infer_state_schema_name:
+        profile_name = context.profile_name
+        # Note: we deliberately isolate state based on the target *schema* and not the target name.
+        # It is assumed that the project will define a target, eg 'dev', and then in each users own ~/.dbt/profiles.yml the schema
+        # for the 'dev' target is overriden to something user-specific, rather than making the target name itself user-specific.
+        # This means that the schema name is the indicator of isolated state, not the target name which may be re-used across multiple schemas.
+        target_schema = profile.target.schema_
+        gateway_kwargs["state_schema"] = f"sqlmesh_state_{profile_name}_{target_schema}"
 
     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,
-        ),
+        dbt=RootDbtConfig(infer_state_schema_name=infer_state_schema_name),
         **{
             "default_gateway": profile.target_name if "gateways" not in kwargs else "",
             "gateways": {
                 profile.target_name: GatewayConfig(
                     connection=profile.target.to_sqlmesh(**target_to_sqlmesh_args),
                     state_connection=state_connection,
+                    **gateway_kwargs,
                 )
             },  # type: ignore
             **kwargs,

tests/dbt/test_config.py

@@ -1215,23 +1215,24 @@ def test_empty_vars_config(tmp_path):
     assert project.context.variables == {}
 
 
-def test_state_schema_naming_pattern(create_empty_project: EmptyProjectCreator):
+def test_infer_state_schema_name(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
+    # infer_state_schema_name defaults to False if omitted
     config = sqlmesh_config(project_root=project_dir)
-    assert not config.state_schema_naming_pattern
+    assert config.dbt
+    assert not config.dbt.infer_state_schema_name
     assert config.get_state_schema() == "sqlmesh"
 
     # create_empty_project() uses the default dbt template for sqlmesh yaml config which
-    # sets state_schema_naming_pattern
+    # sets infer_state_schema_name=True
     ctx = Context(paths=[project_dir])
-    assert ctx.config.state_schema_naming_pattern
-    assert ctx.config.get_state_schema() == "sqlmesh_state_test_foo_dev"
+    assert ctx.config.dbt
+    assert ctx.config.dbt.infer_state_schema_name
+    assert ctx.config.get_state_schema() == "sqlmesh_state_test_foo_main"
     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"
+    assert ctx.state_sync.state_sync.schema == "sqlmesh_state_test_foo_main"
 
     # If the user delberately overrides state_schema then we should respect this choice
     config_file = project_dir / "sqlmesh.yaml"
@@ -1240,7 +1241,8 @@ def test_state_schema_naming_pattern(create_empty_project: EmptyProjectCreator):
     config_file.write_text(yaml_dump(config_yaml))
 
     ctx = Context(paths=[project_dir])
-    assert ctx.config.state_schema_naming_pattern
+    assert ctx.config.dbt
+    assert ctx.config.dbt.infer_state_schema_name
     assert ctx.config.get_state_schema() == "state_override"
     assert isinstance(ctx.state_sync, CachingStateSync)
     assert isinstance(ctx.state_sync.state_sync, EngineAdapterStateSync)

tests/dbt/test_integration.py

@@ -625,21 +625,30 @@ def test_state_schema_isolation_per_target(jaffle_shop_duckdb: Path):
 
     # 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 prod_ctx.config.get_state_schema() == "sqlmesh_state_jaffle_shop_prod_schema"
     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 dev_ctx.config.get_state_schema() == "sqlmesh_state_jaffle_shop_dev_schema"
     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 default_ctx.config.get_state_schema() == "sqlmesh_state_jaffle_shop_dev_schema"
     assert all("dev_schema" in fqn for fqn in default_ctx.models)
     assert not default_ctx.plan(auto_apply=True).has_changes
+
+    # an explicit state schema override set in `sqlmesh.yaml` should use that
+    sqlmesh_yaml_file = jaffle_shop_duckdb / "sqlmesh.yaml"
+    sqlmesh_yaml = yaml_load(sqlmesh_yaml_file)
+    sqlmesh_yaml["gateways"] = {"dev": {"state_schema": "sqlmesh_dev_state_override"}}
+    sqlmesh_yaml_file.write_text(yaml_dump(sqlmesh_yaml))
+    default_ctx = Context(paths=[jaffle_shop_duckdb])
+    assert default_ctx.config.get_state_schema() == "sqlmesh_dev_state_override"
+    assert all("dev_schema" in fqn for fqn in default_ctx.models)