Feat!: Add support for multiple virtual layers controlled by gateway (#4101)

Co-authored-by: Iaroslav Zeigerman <zeigerman.ia@gmail.com>

Author: themisvaltinos

docs/reference/configuration.md

@@ -35,6 +35,7 @@ Configuration options for SQLMesh environment creation and promotion.
 | `physical_schema_override`    | (Deprecated) Use `physical_schema_mapping` instead. A mapping from model schema names to names of schemas in which physical tables for the corresponding models will be placed.                                                                                                                    | dict[string, string] | N        |
 | `physical_schema_mapping`     | A mapping from regular expressions to names of schemas in which physical tables for the corresponding models [will be placed](../guides/configuration.md#physical-table-schemas). (Default physical schema name: `sqlmesh__[model schema]`)                                                        | dict[string, string] | N        |
 | `environment_suffix_target`   | Whether SQLMesh views should append their environment name to the `schema` or `table` - [additional details](../guides/configuration.md#view-schema-override). (Default: `schema`)                                                                                                                 | string               | N        |
+| `gateway_managed_virtual_layer`   | Whether SQLMesh views of the virtual layer will be created by the default gateway or model specified gateways - [additional details](../guides/configuration.md#view-schema-override). (Default: False)                                                                                                                 | boolean               | N        |
 | `environment_catalog_mapping` | A mapping from regular expressions to catalog names. The catalog name is used to determine the target catalog for a given environment.                                                                                                                                                             | dict[string, string] | N        |
 | `log_limit`                   | The default number of logs to keep (Default: `20`)                                                                                                                                                                                                                                                 | int                  | N        |
 

sqlmesh/core/config/root.py

@@ -72,6 +72,7 @@ class Config(BaseConfig):
         model_defaults: Default values for model definitions.
         physical_schema_mapping: A mapping from regular expressions to names of schemas in which physical tables for corresponding models will be placed.
         environment_suffix_target: Indicates whether to append the environment name to the schema or table name.
+        gateway_managed_virtual_layer: Whether the models' views in the virtual layer are created by the model-specific gateway rather than the default gateway.
         environment_catalog_mapping: A mapping from regular expressions to catalog names. The catalog name is used to determine the target catalog for a given environment.
         default_target_environment: The name of the environment that will be the default target for the `sqlmesh plan` and `sqlmesh run` commands.
         log_limit: The default number of logs to keep.
@@ -110,6 +111,7 @@ class Config(BaseConfig):
     environment_suffix_target: EnvironmentSuffixTarget = Field(
         default=EnvironmentSuffixTarget.default
     )
+    gateway_managed_virtual_layer: bool = False
     environment_catalog_mapping: t.Dict[re.Pattern, str] = {}
     default_target_environment: str = c.PROD
     log_limit: int = c.DEFAULT_LOG_LIMIT

sqlmesh/core/context.py

@@ -364,6 +364,7 @@ def __init__(
         self._environment_statements: t.List[EnvironmentStatements] = []
         self._excluded_requirements: t.Set[str] = set()
         self._default_catalog: t.Optional[str] = None
+        self._default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
 
@@ -380,6 +381,7 @@ def __init__(
         self.pinned_environments = Environment.sanitize_names(self.config.pinned_environments)
         self.auto_categorize_changes = self.config.plan.auto_categorize_changes
         self.selected_gateway = gateway or self.config.default_gateway_name
+        self.gateway_managed_virtual_layer = self.config.gateway_managed_virtual_layer
 
         gw_model_defaults = self.config.gateways[self.selected_gateway].model_defaults
         if gw_model_defaults:
@@ -2212,16 +2214,6 @@ def _model_tables(self) -> t.Dict[str, str]:
             for fqn, snapshot in self.snapshots.items()
         }
 
-    @property
-    def _snapshot_gateways(self) -> t.Dict[str, str]:
-        """Mapping of snapshot name to the gateway if specified in the model."""
-
-        return {
-            fqn: snapshot.model.gateway
-            for fqn, snapshot in self.snapshots.items()
-            if snapshot.is_model and snapshot.model.gateway
-        }
-
     @cached_property
     def engine_adapters(self) -> t.Dict[str, EngineAdapter]:
         """Returns all the engine adapters for the gateways defined in the configuration."""
@@ -2232,6 +2224,17 @@ def engine_adapters(self) -> t.Dict[str, EngineAdapter]:
                 self._engine_adapters[gateway_name] = adapter
         return self._engine_adapters
 
+    @cached_property
+    def default_catalog_per_gateway(self) -> t.Dict[str, str]:
+        """Returns the default catalogs for each engine adapter."""
+        if self._default_catalog_per_gateway is None:
+            self._default_catalog_per_gateway = {
+                name: adapter.default_catalog
+                for name, adapter in self.engine_adapters.items()
+                if adapter.default_catalog
+            }
+        return self._default_catalog_per_gateway
+
     def _get_engine_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
         if gateway:
             if adapter := self.engine_adapters.get(gateway):
@@ -2292,22 +2295,33 @@ def _context_diff(
             ensure_finalized_snapshots=ensure_finalized_snapshots,
             diff_rendered=diff_rendered,
             environment_statements=self._environment_statements,
+            gateway_managed_virtual_layer=self.gateway_managed_virtual_layer,
         )
 
     def _run_janitor(self, ignore_ttl: bool = False) -> None:
+        # Clean up expired environments by removing their views and schemas
         self._cleanup_environments()
-        expired_snapshots = self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+
+        # Identify and delete expired snapshots
+        cleanup_targets = self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+
+        # Remove the expired snapshots tables
         self.snapshot_evaluator.cleanup(
-            expired_snapshots,
-            self._snapshot_gateways,
+            target_snapshots=cleanup_targets,
             on_complete=self.console.update_cleanup_progress,
         )
 
         self.state_sync.compact_intervals()
 
     def _cleanup_environments(self) -> None:
         expired_environments = self.state_sync.delete_expired_environments()
-        cleanup_expired_views(self.engine_adapter, expired_environments, console=self.console)
+
+        cleanup_expired_views(
+            default_adapter=self.engine_adapter,
+            engine_adapters=self.engine_adapters,
+            environments=expired_environments,
+            console=self.console,
+        )
 
     def _try_connection(self, connection_name: str, validator: t.Callable[[], None]) -> None:
         connection_name = connection_name.capitalize()

sqlmesh/core/context_diff.py

@@ -53,6 +53,10 @@ class ContextDiff(PydanticModel):
     """Whether the currently stored environment record is in unfinalized state."""
     normalize_environment_name: bool
     """Whether the environment name should be normalized."""
+    previous_gateway_managed_virtual_layer: bool
+    """Whether the previous environment's virtual layer's views were created by the model specified gateways."""
+    gateway_managed_virtual_layer: bool
+    """Whether the virtual layer's views will be created by the model specified gateways."""
     create_from: str
     """The name of the environment the target environment will be created from if new."""
     create_from_env_exists: bool
@@ -96,6 +100,7 @@ def create(
         excluded_requirements: t.Optional[t.Set[str]] = None,
         diff_rendered: bool = False,
         environment_statements: t.Optional[t.List[EnvironmentStatements]] = [],
+        gateway_managed_virtual_layer: bool = False,
     ) -> ContextDiff:
         """Create a ContextDiff object.
 
@@ -226,6 +231,8 @@ def create(
             diff_rendered=diff_rendered,
             previous_environment_statements=previous_environment_statements,
             environment_statements=environment_statements,
+            previous_gateway_managed_virtual_layer=env.gateway_managed if env else False,
+            gateway_managed_virtual_layer=gateway_managed_virtual_layer,
         )
 
     @classmethod
@@ -263,6 +270,8 @@ def create_no_diff(cls, environment: str, state_reader: StateReader) -> ContextD
             previous_requirements=env.requirements,
             requirements=env.requirements,
             previous_environment_statements=[],
+            previous_gateway_managed_virtual_layer=env.gateway_managed,
+            gateway_managed_virtual_layer=env.gateway_managed,
         )
 
     @property
@@ -273,6 +282,7 @@ def has_changes(self) -> bool:
             or self.is_unfinalized_environment
             or self.has_requirement_changes
             or self.has_environment_statements_changes
+            or self.previous_gateway_managed_virtual_layer != self.gateway_managed_virtual_layer
         )
 
     @property

sqlmesh/core/environment.py

@@ -16,7 +16,7 @@
 from sqlmesh.utils.date import TimeLike, now_timestamp
 from sqlmesh.utils.jinja import JinjaMacroRegistry
 from sqlmesh.utils.metaprogramming import Executable
-from sqlmesh.utils.pydantic import PydanticModel, field_validator
+from sqlmesh.utils.pydantic import PydanticModel, field_validator, ValidationInfo
 
 T = t.TypeVar("T", bound="EnvironmentNamingInfo")
 PydanticType = t.TypeVar("PydanticType", bound="PydanticModel")
@@ -32,22 +32,27 @@ class EnvironmentNamingInfo(PydanticModel):
         catalog_name_override: The name of the catalog to use for this environment if an override was provided
         normalize_name: Indicates whether the environment's name will be normalized. For example, if it's
             `dev`, then it will become `DEV` when targeting Snowflake.
+        gateway_managed: Determines whether the virtual layer's views are created by the model-specific
+            gateways, otherwise the default gateway is used. Default: False.
     """
 
     name: str = c.PROD
     suffix_target: EnvironmentSuffixTarget = Field(default=EnvironmentSuffixTarget.SCHEMA)
     catalog_name_override: t.Optional[str] = None
     normalize_name: bool = True
+    gateway_managed: bool = False
 
     @field_validator("name", mode="before")
     @classmethod
     def _sanitize_name(cls, v: str) -> str:
         return word_characters_only(v).lower()
 
-    @field_validator("normalize_name", mode="before")
+    @field_validator("normalize_name", "gateway_managed", mode="before")
     @classmethod
-    def _validate_normalize_name(cls, v: t.Any) -> bool:
-        return True if v is None else bool(v)
+    def _validate_boolean_field(cls, v: t.Any, info: ValidationInfo) -> bool:
+        if v is None:
+            return info.field_name == "normalize_name"
+        return bool(v)
 
     @t.overload
     @classmethod
@@ -194,6 +199,7 @@ def naming_info(self) -> EnvironmentNamingInfo:
             suffix_target=self.suffix_target,
             catalog_name_override=self.catalog_name_override,
             normalize_name=self.normalize_name,
+            gateway_managed=self.gateway_managed,
         )
 
     @property

sqlmesh/core/loader.py

@@ -468,6 +468,7 @@ def _load() -> t.List[Model]:
                         default_catalog=self.context.default_catalog,
                         infer_names=self.config.model_naming.infer_names,
                         signal_definitions=signals,
+                        default_catalog_per_gateway=self.context.default_catalog_per_gateway,
                     )
                 except Exception as ex:
                     raise ConfigError(f"Failed to load model definition at '{path}'.\n{ex}")
@@ -525,6 +526,7 @@ def _load_python_models(
                             default_catalog=self.context.default_catalog,
                             infer_names=self.config.model_naming.infer_names,
                             audit_definitions=audits,
+                            default_catalog_per_gateway=self.context.default_catalog_per_gateway,
                         ):
                             if model.enabled:
                                 models[model.fqn] = model

sqlmesh/core/model/decorator.py

@@ -93,6 +93,7 @@ def models(
         path: Path,
         module_path: Path,
         dialect: t.Optional[str] = None,
+        default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
         **loader_kwargs: t.Any,
     ) -> t.List[Model]:
         return create_models_from_blueprints(
@@ -103,6 +104,7 @@ def models(
             path=path,
             module_path=module_path,
             dialect=dialect,
+            default_catalog_per_gateway=default_catalog_per_gateway,
             **loader_kwargs,
         )
 

sqlmesh/core/model/definition.py

@@ -1886,6 +1886,7 @@ def create_models_from_blueprints(
     path: Path = Path(),
     module_path: Path = Path(),
     dialect: DialectType = None,
+    default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
     **loader_kwargs: t.Any,
 ) -> t.List[Model]:
     model_blueprints: t.List[Model] = []
@@ -1907,6 +1908,13 @@ def create_models_from_blueprints(
         else:
             gateway_name = None
 
+        if (
+            default_catalog_per_gateway
+            and gateway_name
+            and (catalog := default_catalog_per_gateway.get(gateway_name)) is not None
+        ):
+            loader_kwargs["default_catalog"] = catalog
+
         model_blueprints.append(
             loader(
                 path=path,
@@ -1927,6 +1935,7 @@ def load_sql_based_models(
     path: Path = Path(),
     module_path: Path = Path(),
     dialect: DialectType = None,
+    default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
     **loader_kwargs: t.Any,
 ) -> t.List[Model]:
     gateway: t.Optional[exp.Expression] = None
@@ -1964,6 +1973,7 @@ def load_sql_based_models(
         path=path,
         module_path=module_path,
         dialect=dialect,
+        default_catalog_per_gateway=default_catalog_per_gateway,
         **loader_kwargs,
     )
 

sqlmesh/core/plan/builder.py

@@ -151,6 +151,7 @@ def __init__(
             name=self._context_diff.environment,
             suffix_target=environment_suffix_target,
             normalize_name=self._context_diff.normalize_environment_name,
+            gateway_managed=self._context_diff.gateway_managed_virtual_layer,
         )
 
         self._latest_plan: t.Optional[Plan] = None

sqlmesh/core/snapshot/definition.py

@@ -478,6 +478,7 @@ class SnapshotTableInfo(PydanticModel, SnapshotInfoMixin, frozen=True):
     base_table_name_override: t.Optional[str] = None
     custom_materialization: t.Optional[str] = None
     dev_table_suffix: str
+    model_gateway: t.Optional[str] = None
 
     def __lt__(self, other: SnapshotTableInfo) -> bool:
         return self.name < other.name
@@ -1179,6 +1180,7 @@ def table_info(self) -> SnapshotTableInfo:
             node_type=self.node_type,
             custom_materialization=custom_materialization,
             dev_table_suffix=self.dev_table_suffix,
+            model_gateway=self.model_gateway,
         )
 
     @property