Feat!: Make Environment Suffix Target Configurable (#1220)

* wip - add environment suffix target

* remove bad import

* properly cleanup table suffix views

* fix return type

* add test covering basic promotion functionality

* test cleanup expired views

* add context janitor test

* only demote if there are removed values

* fix migration script

* add integration test

* update Airflow tests

* document environment suffix target

Author: eakmanrq

docs/reference/configuration.md

@@ -223,8 +223,21 @@ adhere to at your organization and therefore need more control over the schema n
 
 ### View Schema Override
 
-Coming soon. Please message us on [slack](https://tobikodata.com/slack) if you are interested in this feature so we can better understand your 
-use case and make sure the new feature satisfies your needs.
+By default SQLMesh appends the environment name to the schema name when creating new environments. This can be changed
+to instead append a suffix at the end of table instead. This means that new environment views will be created in the 
+same schema as production but be differentiated having their names end with `__<env>`.
+
+Config example:
+    
+```yaml linenums="1"
+environment_suffix_target: table
+```
+
+If you had a model name of `db.users`, and you were creating a `dev` environment, then the view would be created as `db.users__dev` instead of the default behavior of `db__dev.users`.
+
+The default behavior of appending the suffix to schemas is recommended because it leaves production with a single clean
+interface for accessing the views. However if you are deploying SQLMesh in an environment with tight restrictions on 
+schema creation then this can be a useful way of reducing the number of schemas that need to be created.
 
 ## Additional details
 

examples/sushi/config.py

@@ -6,6 +6,7 @@
     CategorizerConfig,
     Config,
     DuckDBConnectionConfig,
+    EnvironmentSuffixTarget,
     GatewayConfig,
     ModelDefaultsConfig,
     SparkConnectionConfig,
@@ -104,3 +105,10 @@
     ],
     model_defaults=ModelDefaultsConfig(dialect="duckdb"),
 )
+
+
+environment_suffix_config = Config(
+    default_connection=DuckDBConnectionConfig(),
+    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    environment_suffix_target=EnvironmentSuffixTarget.TABLE,
+)

sqlmesh/core/config/__init__.py

@@ -1,4 +1,5 @@
 from sqlmesh.core.config.categorizer import AutoCategorizationMode, CategorizerConfig
+from sqlmesh.core.config.common import EnvironmentSuffixTarget
 from sqlmesh.core.config.connection import (
     BigQueryConnectionConfig,
     ConnectionConfig,

sqlmesh/core/config/common.py

@@ -1,12 +1,31 @@
 from __future__ import annotations
 
 import typing as t
+from enum import Enum
 
 from pydantic import validator
 
+from sqlmesh.utils import classproperty
 from sqlmesh.utils.errors import ConfigError
 
 
+class EnvironmentSuffixTarget(str, Enum):
+    SCHEMA = "schema"
+    TABLE = "table"
+
+    @property
+    def is_schema(self) -> bool:
+        return self == EnvironmentSuffixTarget.SCHEMA
+
+    @property
+    def is_table(self) -> bool:
+        return self == EnvironmentSuffixTarget.TABLE
+
+    @classproperty
+    def default(cls) -> EnvironmentSuffixTarget:
+        return EnvironmentSuffixTarget.SCHEMA
+
+
 def _concurrent_tasks_validator(v: t.Any) -> int:
     if isinstance(v, str):
         v = int(v)

sqlmesh/core/config/root.py

@@ -2,9 +2,10 @@
 
 import typing as t
 
-from pydantic import root_validator, validator
+from pydantic import Field, root_validator, validator
 
 from sqlmesh.core import constants as c
+from sqlmesh.core.config import EnvironmentSuffixTarget
 from sqlmesh.core.config.base import BaseConfig, UpdateStrategy
 from sqlmesh.core.config.categorizer import CategorizerConfig
 from sqlmesh.core.config.connection import ConnectionConfig, DuckDBConnectionConfig
@@ -41,6 +42,7 @@ class Config(BaseConfig):
         pinned_environments: A list of development environment names that should not be deleted by the janitor task.
         model_defaults: Default values for model definitions.
         include_unmodified: Indicates whether to include unmodified models in the target development environment.
+        environment_suffix_target: Indicates whether to append the environment name to the schema or table name.
     """
 
     gateways: t.Union[t.Dict[str, GatewayConfig], GatewayConfig] = GatewayConfig()
@@ -63,6 +65,9 @@ class Config(BaseConfig):
     username: str = ""
     include_unmodified: bool = False
     physical_schema_override: t.Dict[str, str] = {}
+    environment_suffix_target: EnvironmentSuffixTarget = Field(
+        default=EnvironmentSuffixTarget.default
+    )
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
         "gateways": UpdateStrategy.KEY_UPDATE,

sqlmesh/core/console.py

@@ -23,7 +23,7 @@
 from rich.table import Table
 from rich.tree import Tree
 
-from sqlmesh.core import constants as c
+from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory
 from sqlmesh.core.test import ModelTest
 from sqlmesh.utils import rich as srich
@@ -54,7 +54,11 @@ class Console(abc.ABC):
     with them when their input is needed."""
 
     @abc.abstractmethod
-    def start_evaluation_progress(self, batches: t.Dict[Snapshot, int], environment: str) -> None:
+    def start_evaluation_progress(
+        self,
+        batches: t.Dict[Snapshot, int],
+        environment_naming_info: EnvironmentNamingInfo,
+    ) -> None:
         """Indicates that a new snapshot evaluation progress has begun."""
 
     @abc.abstractmethod
@@ -180,7 +184,7 @@ def __init__(self, console: t.Optional[RichConsole] = None, **kwargs: t.Any) ->
         self.evaluation_model_progress: t.Optional[Progress] = None
         self.evaluation_model_tasks: t.Dict[str, TaskID] = {}
         self.evaluation_model_batches: t.Dict[Snapshot, int] = {}
-        self.evaluation_environment: str = c.PROD
+        self.evaluation_environment_naming_info = EnvironmentNamingInfo()
 
         self.creation_progress: t.Optional[Progress] = None
         self.creation_task: t.Optional[TaskID] = None
@@ -202,7 +206,11 @@ def _prompt(self, message: str, **kwargs: t.Any) -> t.Any:
     def _confirm(self, message: str, **kwargs: t.Any) -> bool:
         return Confirm.ask(message, console=self.console, **kwargs)
 
-    def start_evaluation_progress(self, batches: t.Dict[Snapshot, int], enviornment: str) -> None:
+    def start_evaluation_progress(
+        self,
+        batches: t.Dict[Snapshot, int],
+        environment_naming_info: EnvironmentNamingInfo,
+    ) -> None:
         """Indicates that a new snapshot evaluation progress has begun."""
         if not self.evaluation_progress_live:
             self.evaluation_total_progress = Progress(
@@ -234,11 +242,13 @@ def start_evaluation_progress(self, batches: t.Dict[Snapshot, int], enviornment:
             )
 
             self.evaluation_model_batches = batches
-            self.evaluation_environment = enviornment
+            self.evaluation_environment_naming_info = environment_naming_info
 
     def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         if self.evaluation_model_progress and snapshot.name not in self.evaluation_model_tasks:
-            view_name = snapshot.qualified_view_name.for_environment(self.evaluation_environment)
+            view_name = snapshot.qualified_view_name.for_environment(
+                self.evaluation_environment_naming_info
+            )
             self.evaluation_model_tasks[snapshot.name] = self.evaluation_model_progress.add_task(
                 f"Evaluating {view_name}...",
                 view_name=view_name,
@@ -273,7 +283,7 @@ def stop_evaluation_progress(self, success: bool = True) -> None:
         self.evaluation_model_progress = None
         self.evaluation_model_tasks = {}
         self.evaluation_model_batches = {}
-        self.evaluation_environment = c.PROD
+        self.evaluation_environment_naming_info = EnvironmentNamingInfo()
 
     def start_creation_progress(self, total_tasks: int) -> None:
         """Indicates that a new creation progress has begun."""
@@ -504,7 +514,7 @@ def _show_missing_dates(self, plan: Plan) -> None:
         backfill = Tree("[bold]Models needing backfill (missing dates):")
         for missing in plan.missing_intervals:
             snapshot = plan.context_diff.snapshots[missing.snapshot_name]
-            view_name = snapshot.qualified_view_name.for_environment(plan.environment_name)
+            view_name = snapshot.qualified_view_name.for_environment(plan.environment_naming_info)
             backfill.add(f"{view_name}: {missing.format_intervals(snapshot.model.interval_unit)}")
         self._print(backfill)
 
@@ -1079,7 +1089,7 @@ def _show_missing_dates(self, plan: Plan) -> None:
         self._print("**Models needing backfill (missing dates):**\n\n")
         for missing in plan.missing_intervals:
             snapshot = plan.context_diff.snapshots[missing.snapshot_name]
-            view_name = snapshot.qualified_view_name.for_environment(plan.environment_name)
+            view_name = snapshot.qualified_view_name.for_environment(plan.environment_naming_info)
             self._print(
                 f"* `{view_name}`: {missing.format_intervals(snapshot.model.interval_unit)}\n"
             )
@@ -1149,13 +1159,19 @@ def _confirm(self, message: str, **kwargs: t.Any) -> bool:
         self._print(message)
         return super()._confirm("", **kwargs)
 
-    def start_evaluation_progress(self, batches: t.Dict[Snapshot, int], environment: str) -> None:
+    def start_evaluation_progress(
+        self,
+        batches: t.Dict[Snapshot, int],
+        environment_naming_info: EnvironmentNamingInfo,
+    ) -> None:
         self.evaluation_batches = batches
-        self.evaluation_environment = environment
+        self.evaluation_environment_naming_info = environment_naming_info
 
     def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         if not self.evaluation_batch_progress.get(snapshot.name):
-            view_name = snapshot.qualified_view_name.for_environment(self.evaluation_environment)
+            view_name = snapshot.qualified_view_name.for_environment(
+                self.evaluation_environment_naming_info
+            )
             self.evaluation_batch_progress[snapshot.name] = (view_name, 0)
             print(f"Starting '{view_name}', Total batches: {self.evaluation_batches[snapshot]}")
 

sqlmesh/core/context.py

@@ -57,7 +57,7 @@
     parse,
 )
 from sqlmesh.core.engine_adapter import EngineAdapter
-from sqlmesh.core.environment import Environment
+from sqlmesh.core.environment import Environment, EnvironmentNamingInfo
 from sqlmesh.core.loader import Loader, SqlMeshLoader, update_model_schemas
 from sqlmesh.core.macros import ExecutableOrMacro
 from sqlmesh.core.metric import Metric
@@ -77,7 +77,12 @@
     SnapshotFingerprint,
     to_table_mapping,
 )
-from sqlmesh.core.state_sync import CachingStateSync, StateReader, StateSync
+from sqlmesh.core.state_sync import (
+    CachingStateSync,
+    StateReader,
+    StateSync,
+    cleanup_expired_views,
+)
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import get_all_model_tests, run_model_tests, run_tests
 from sqlmesh.core.user import User
@@ -406,7 +411,11 @@ def run(
         )
         try:
             self.scheduler(environment=environment).run(
-                environment, start, end, execution_time, ignore_cron=ignore_cron
+                environment,
+                start=start,
+                end=end,
+                execution_time=execution_time,
+                ignore_cron=ignore_cron,
             )
         except Exception as e:
             self.notification_target_manager.notify(
@@ -773,6 +782,7 @@ def plan(
             is_dev=environment != c.PROD,
             forward_only=forward_only,
             environment_ttl=environment_ttl,
+            environment_suffix_target=self.config.environment_suffix_target,
             categorizer_config=self.auto_categorize_changes,
             auto_categorization_enabled=not no_auto_categorization,
             effective_from=effective_from,
@@ -1138,7 +1148,11 @@ def _model_tables(self) -> t.Dict[str, str]:
         return {
             name: snapshot.table_name()
             if snapshot.version
-            else snapshot.qualified_view_name.for_environment(c.PROD)
+            else snapshot.qualified_view_name.for_environment(
+                EnvironmentNamingInfo(
+                    name=c.PROD, suffix_target=self.config.environment_suffix_target
+                )
+            )
             for name, snapshot in self.snapshots.items()
         }
 
@@ -1185,14 +1199,7 @@ def _load_configs(
 
     def _run_janitor(self) -> None:
         expired_environments = self.state_sync.delete_expired_environments()
-        expired_schemas = {
-            snapshot.qualified_view_name.schema_for_environment(expired_environment.name)
-            for expired_environment in expired_environments
-            for snapshot in expired_environment.snapshots
-        }
-        for expired_schema in expired_schemas:
-            self.engine_adapter.drop_schema(expired_schema, ignore_if_not_exists=True, cascade=True)
-
+        cleanup_expired_views(self.engine_adapter, expired_environments)
         expired_snapshots = self.state_sync.delete_expired_snapshots()
         self.snapshot_evaluator.cleanup(expired_snapshots)