Feat!: Introduce multiple gateway virtual layer

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

@@ -379,6 +379,8 @@ 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
+        self.catalogs: t.Dict[str, str] = {}
 
         gw_model_defaults = self.config.gateways[self.selected_gateway].model_defaults
         if gw_model_defaults:
@@ -584,6 +586,15 @@ def load(self, update_schemas: bool = True) -> GenericContext[C]:
         """Load all files in the context's path."""
         load_start_ts = time.perf_counter()
 
+        # In a multi virtual layer setup we need the catalog of each engine
+        # since there is a possibility of not a shared catalog between them
+        if self.gateway_managed_virtual_layer:
+            self.catalogs = {
+                name: adapter.default_catalog
+                for name, adapter in self.engine_adapters.items()
+                if adapter.default_catalog
+            }
+
         loaded_projects = [loader.load() for loader in self._loaders]
 
         self.dag = DAG()
@@ -2207,16 +2218,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."""
@@ -2287,14 +2288,18 @@ 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:
         self._cleanup_environments()
-        expired_snapshots = self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl)
+        expired_snapshots, snapshot_gateways = self.state_sync.delete_expired_snapshots(
+            ignore_ttl=ignore_ttl
+        )
+
         self.snapshot_evaluator.cleanup(
             expired_snapshots,
-            self._snapshot_gateways,
+            snapshot_gateways,
             on_complete=self.console.update_cleanup_progress,
         )
 

sqlmesh/core/context_diff.py

@@ -53,6 +53,8 @@ class ContextDiff(PydanticModel):
     """Whether the currently stored environment record is in unfinalized state."""
     normalize_environment_name: bool
     """Whether the environment name should be normalized."""
+    gateway_managed_virtual_layer: bool = False
+    """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 +98,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.
 
@@ -118,7 +121,11 @@ def create(
         env = state_reader.get_environment(environment)
 
         create_from_env_exists = False
-        if env is None or env.expired:
+        if (
+            env is None
+            or env.expired
+            or env.gateway_managed_virtual_layer != gateway_managed_virtual_layer
+        ):
             env = state_reader.get_environment(create_from.lower())
 
             if not env and create_from != c.PROD:
@@ -226,6 +233,7 @@ def create(
             diff_rendered=diff_rendered,
             previous_environment_statements=previous_environment_statements,
             environment_statements=environment_statements,
+            gateway_managed_virtual_layer=gateway_managed_virtual_layer,
         )
 
     @classmethod

sqlmesh/core/environment.py

@@ -32,12 +32,15 @@ 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_virtual_layer: 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_virtual_layer: bool = False
 
     @field_validator("name", mode="before")
     @classmethod
@@ -49,6 +52,11 @@ def _sanitize_name(cls, v: str) -> str:
     def _validate_normalize_name(cls, v: t.Any) -> bool:
         return True if v is None else bool(v)
 
+    @field_validator("gateway_managed_virtual_layer", mode="before")
+    @classmethod
+    def _validate_gateway_managed_virtual_layer(cls, v: t.Any) -> bool:
+        return False if v is None else bool(v)
+
     @t.overload
     @classmethod
     def sanitize_name(cls, v: str) -> str: ...
@@ -194,6 +202,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_virtual_layer=self.gateway_managed_virtual_layer,
         )
 
     @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,
+                        catalogs=self.context.catalogs,
                     )
                 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,
+                            catalogs=self.context.catalogs,
                         ):
                             if model.enabled:
                                 models[model.fqn] = model

sqlmesh/core/model/definition.py

@@ -1907,6 +1907,13 @@ def create_models_from_blueprints(
         else:
             gateway_name = None
 
+        if (
+            (catalogs := loader_kwargs.pop("catalogs", None))
+            and gateway_name
+            and (catalog := catalogs.get(gateway_name))
+        ):
+            loader_kwargs["default_catalog"] = catalog
+
         model_blueprints.append(
             loader(
                 path=path,

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_virtual_layer=self._context_diff.gateway_managed_virtual_layer,
         )
 
         self._latest_plan: t.Optional[Plan] = None

sqlmesh/core/plan/evaluator.py

@@ -423,8 +423,16 @@ def _demote_snapshots(
         environment_naming_info: EnvironmentNamingInfo,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
+        # In a multi virtual layer setup we need the gateway info from the snapshots for demotion
+        snapshots_to_demote: t.List[Snapshot] = []
+        if environment_naming_info.gateway_managed_virtual_layer:
+            removed_snapshots = self.state_sync.get_snapshots(target_snapshots)
+            snapshots_to_demote = [removed_snapshots[s.snapshot_id] for s in target_snapshots]
+
         self.snapshot_evaluator.demote(
-            target_snapshots, environment_naming_info, on_complete=on_complete
+            snapshots_to_demote or target_snapshots,
+            environment_naming_info,
+            on_complete=on_complete,
         )
 
     def _restate(self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapshot]) -> None:

sqlmesh/core/snapshot/evaluator.py

@@ -226,15 +226,23 @@ def promote(
             deployability_index: Determines snapshots that are deployable in the context of this promotion.
             on_complete: A callback to call on each successfully promoted snapshot.
         """
-        self._create_schemas(
-            [
-                s.qualified_view_name.table_for_environment(
-                    environment_naming_info, dialect=self.adapter.dialect
+
+        gateway_by_schema: t.Dict[t.Any, str] = {}
+        tables: t.List[t.Any] = []
+        for snapshot in target_snapshots:
+            if snapshot.is_model and not snapshot.is_symbolic:
+                table = snapshot.qualified_view_name.table_for_environment(
+                    environment_naming_info,
+                    dialect=self._get_adapter(snapshot.model_gateway).dialect
+                    if environment_naming_info.gateway_managed_virtual_layer
+                    else self.adapter.dialect,
                 )
-                for s in target_snapshots
-                if s.is_model and not s.is_symbolic
-            ]
-        )
+                tables.append(table)
+                if environment_naming_info.gateway_managed_virtual_layer:
+                    table_schema = d.schema_(table.db, catalog=table.catalog)
+                    gateway_by_schema[table_schema] = snapshot.model_gateway or ""
+        self._create_schemas(tables=tables, gateways=gateway_by_schema)
+
         deployability_index = deployability_index or DeployabilityIndex.all_deployable()
         with self.concurrent_context():
             concurrent_apply_to_snapshots(
@@ -920,7 +928,11 @@ def _promote_snapshot(
         table_mapping: t.Optional[t.Dict[str, str]] = None,
     ) -> None:
         if snapshot.is_model:
-            adapter = self.adapter
+            adapter = (
+                self._get_adapter(snapshot.model_gateway)
+                if environment_naming_info.gateway_managed_virtual_layer
+                else self.adapter
+            )
             table_name = snapshot.table_name(deployability_index.is_representative(snapshot))
             view_name = snapshot.qualified_view_name.for_environment(
                 environment_naming_info, dialect=adapter.dialect
@@ -953,7 +965,12 @@ def _demote_snapshot(
         environment_naming_info: EnvironmentNamingInfo,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]],
     ) -> None:
-        adapter = self.adapter
+        adapter = (
+            self._get_adapter(snapshot.model_gateway)
+            if environment_naming_info.gateway_managed_virtual_layer
+            and isinstance(snapshot, Snapshot)
+            else self.adapter
+        )
         view_name = snapshot.qualified_view_name.for_environment(
             environment_naming_info, dialect=adapter.dialect
         )