Adding schema ownership

Signed-off-by: Gabe Pesco <PescoG@medinsight.milliman.com>

Author: Gabe Pesco

sqlmesh/core/config/__init__.py

@@ -32,6 +32,7 @@
     load_configs as load_configs,
 )
 from sqlmesh.core.config.migration import MigrationConfig as MigrationConfig
+from sqlmesh.core.config.ownership import OwnershipConfig as OwnershipConfig
 from sqlmesh.core.config.model import ModelDefaultsConfig as ModelDefaultsConfig
 from sqlmesh.core.config.naming import NameInferenceConfig as NameInferenceConfig
 from sqlmesh.core.config.linter import LinterConfig as LinterConfig

sqlmesh/core/config/ownership.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import re
+import typing as t
+
+from pydantic.functional_validators import BeforeValidator
+
+from sqlmesh.core.config.base import BaseConfig
+from sqlmesh.core.config.common import compile_regex_mapping
+
+if t.TYPE_CHECKING:
+    OwnershipMapping = t.Dict[re.Pattern, str]
+else:
+    OwnershipMapping = t.Annotated[
+        t.Dict[re.Pattern, str], BeforeValidator(compile_regex_mapping)
+    ]
+
+
+class OwnershipConfig(BaseConfig):
+    """Configuration for object ownership rules applied at creation time.
+
+    Maps environment name regex patterns to owner principals. The first
+    matching pattern wins. Ownership is applied immediately when schemas and
+    views are created, so even a partially-completed run leaves objects in a
+    manageable state.
+
+    Example::
+
+        ownership:
+          environment_owner_mapping:
+            "^prod$": "svc_prod_spn"
+            ".*": "group:shared-developers"
+    """
+
+    environment_owner_mapping: OwnershipMapping = {}
+
+    def resolve_owner(self, environment_name: str) -> t.Optional[str]:
+        """Return the configured owner for the given environment name, or None."""
+        for pattern, owner in self.environment_owner_mapping.items():
+            if pattern.fullmatch(environment_name):
+                return owner
+        return None

sqlmesh/core/config/root.py

@@ -30,6 +30,7 @@
 from sqlmesh.core.config.format import FormatConfig
 from sqlmesh.core.config.gateway import GatewayConfig
 from sqlmesh.core.config.janitor import JanitorConfig
+from sqlmesh.core.config.ownership import OwnershipConfig
 from sqlmesh.core.config.migration import MigrationConfig
 from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.config.naming import NameInferenceConfig as NameInferenceConfig
@@ -118,6 +119,7 @@ class Config(BaseConfig):
         gateway_managed_virtual_layer: Whether the models' views in the virtual layer are created by the model-specific gateway rather than the default gateway.
         infer_python_dependencies: Whether to statically analyze Python code to automatically infer Python package requirements.
         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.
+        ownership: Ownership rules applied at schema/view creation time. Maps environment name patterns to owner principals so objects are correctly owned even after a partial run.
         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.
         format: The formatting options for SQL code.
@@ -175,6 +177,7 @@ class Config(BaseConfig):
     janitor: JanitorConfig = JanitorConfig()
     cache_dir: t.Optional[str] = None
     dbt: t.Optional[DbtConfig] = None
+    ownership: OwnershipConfig = Field(default_factory=OwnershipConfig)
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
         "gateways": UpdateStrategy.NESTED_UPDATE,
@@ -194,6 +197,7 @@ class Config(BaseConfig):
         "after_all": UpdateStrategy.EXTEND,
         "linter": UpdateStrategy.NESTED_UPDATE,
         "dbt": UpdateStrategy.NESTED_UPDATE,
+        "ownership": UpdateStrategy.NESTED_UPDATE,
     }
 
     _connection_config_validator = connection_config_validator

sqlmesh/core/config/scheduler.py

@@ -128,12 +128,18 @@ class BuiltInSchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig)
     type_: t.Literal["builtin"] = Field(alias="type", default="builtin")
 
     def create_plan_evaluator(self, context: GenericContext) -> PlanEvaluator:
+        from sqlmesh.core.config.ownership import OwnershipConfig
+
+        ownership_config = getattr(context.config, "ownership", None)
+        if isinstance(ownership_config, OwnershipConfig) and not ownership_config.environment_owner_mapping:
+            ownership_config = None
         return BuiltInPlanEvaluator(
             state_sync=context.state_sync,
             snapshot_evaluator=context.snapshot_evaluator,
             create_scheduler=context.create_scheduler,
             default_catalog=context.default_catalog,
             console=context.console,
+            ownership_config=ownership_config,
         )
 
     def get_default_catalog_per_gateway(self, context: GenericContext) -> t.Dict[str, str]:

sqlmesh/core/engine_adapter/base.py

@@ -1418,6 +1418,20 @@ def _create_schema(
                 raise
             logger.warning("Failed to create %s '%s': %s", kind.lower(), schema_name, e)
 
+    def alter_schema_owner(self, schema_name: SchemaName, owner: str) -> None:
+        """Set the owner of a schema.
+
+        No-op by default. Override in dialect-specific adapters that support ownership control
+        (e.g. Spark/Databricks Unity Catalog: ALTER SCHEMA ... OWNER TO ...).
+        """
+
+    def alter_view_owner(self, view_name: TableName, owner: str) -> None:
+        """Set the owner of a view.
+
+        No-op by default. Override in dialect-specific adapters that support ownership control
+        (e.g. Spark/Databricks Unity Catalog: ALTER VIEW ... OWNER TO ...).
+        """
+
     def drop_schema(
         self,
         schema_name: SchemaName,

sqlmesh/core/engine_adapter/spark.py

@@ -553,6 +553,20 @@ def _build_create_comment_column_exp(
 
         return f"ALTER TABLE {table_sql} ALTER COLUMN {column_sql} COMMENT {comment_sql}"
 
+    def alter_schema_owner(self, schema_name: SchemaName, owner: str) -> None:
+        schema_sql = exp.to_table(schema_name, dialect=self.dialect).sql(
+            dialect=self.dialect, identify=True
+        )
+        owner_sql = exp.to_identifier(owner, quoted=True).sql(dialect=self.dialect)
+        self.execute(f"ALTER SCHEMA {schema_sql} OWNER TO {owner_sql}")
+
+    def alter_view_owner(self, view_name: TableName, owner: str) -> None:
+        view_sql = exp.to_table(view_name, dialect=self.dialect).sql(
+            dialect=self.dialect, identify=True
+        )
+        owner_sql = exp.to_identifier(owner, quoted=True).sql(dialect=self.dialect)
+        self.execute(f"ALTER VIEW {view_sql} OWNER TO {owner_sql}")
+
     @classmethod
     def _wap_branch_name(cls, wap_id: str) -> str:
         return f"{cls.WAP_PREFIX}{wap_id}"

sqlmesh/core/plan/evaluator.py

@@ -40,6 +40,7 @@
 from sqlmesh.core.plan.common import identify_restatement_intervals_across_snapshot_versions
 from sqlmesh.utils import CorrelationId
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
+from sqlmesh.core.config.ownership import OwnershipConfig
 from sqlmesh.utils.errors import PlanError, ConflictingPlanError, SQLMeshError
 from sqlmesh.utils.date import now, to_timestamp
 
@@ -74,12 +75,14 @@ def __init__(
         create_scheduler: t.Callable[[t.Iterable[Snapshot], SnapshotEvaluator], Scheduler],
         default_catalog: t.Optional[str],
         console: t.Optional[Console] = None,
+        ownership_config: t.Optional[OwnershipConfig] = None,
     ):
         self.state_sync = state_sync
         self.snapshot_evaluator = snapshot_evaluator
         self.create_scheduler = create_scheduler
         self.default_catalog = default_catalog
         self.console = console or get_console()
+        self.ownership_config = ownership_config
         self._circuit_breaker: t.Optional[t.Callable[[], bool]] = None
 
     def evaluate(
@@ -434,6 +437,9 @@ def _promote_snapshots(
         deployability_index: t.Optional[DeployabilityIndex] = None,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
     ) -> None:
+        owner: t.Optional[str] = None
+        if self.ownership_config:
+            owner = self.ownership_config.resolve_owner(environment_naming_info.name)
         self.snapshot_evaluator.promote(
             target_snapshots,
             start=plan.start,
@@ -449,6 +455,7 @@ def _promote_snapshots(
             environment_naming_info=environment_naming_info,
             deployability_index=deployability_index,
             on_complete=on_complete,
+            owner=owner,
         )
 
     def _demote_snapshots(

sqlmesh/core/snapshot/evaluator.py

@@ -275,6 +275,7 @@ def promote(
         snapshots: t.Optional[t.Dict[SnapshotId, Snapshot]] = None,
         table_mapping: t.Optional[t.Dict[str, str]] = None,
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]] = None,
+        owner: t.Optional[str] = None,
     ) -> None:
         """Promotes the given collection of snapshots in the target environment by replacing a corresponding
         view with a physical table associated with the given snapshot.
@@ -306,7 +307,7 @@ def promote(
         gateway_table_pairs = [
             (gateway, table) for gateway, tables in tables_by_gateway.items() for table in tables
         ]
-        self._create_schemas(gateway_table_pairs=gateway_table_pairs)
+        self._create_schemas(gateway_table_pairs=gateway_table_pairs, owner=owner)
 
         # Fetch the view data objects for the promoted snapshots to get them cached
         self._get_virtual_data_objects(target_snapshots, environment_naming_info)
@@ -325,6 +326,7 @@ def promote(
                     environment_naming_info=environment_naming_info,
                     deployability_index=deployability_index,  # type: ignore
                     on_complete=on_complete,
+                    owner=owner,
                 ),
                 self.ddl_concurrent_tasks,
             )
@@ -1257,6 +1259,7 @@ def _promote_snapshot(
         execution_time: t.Optional[TimeLike] = None,
         snapshots: t.Optional[t.Dict[SnapshotId, Snapshot]] = None,
         table_mapping: t.Optional[t.Dict[str, str]] = None,
+        owner: t.Optional[str] = None,
     ) -> None:
         if not snapshot.is_model:
             return
@@ -1298,6 +1301,9 @@ def _promote_snapshot(
             render_kwargs["snapshots"] = snapshot_by_name
             adapter.execute(snapshot.model.render_on_virtual_update(**render_kwargs))
 
+        if owner:
+            adapter.alter_view_owner(view_name, owner)
+
         if on_complete is not None:
             on_complete(snapshot)
 
@@ -1449,6 +1455,7 @@ def _create_catalogs(
     def _create_schemas(
         self,
         gateway_table_pairs: t.Iterable[t.Tuple[t.Optional[str], t.Union[exp.Table, str]]],
+        owner: t.Optional[str] = None,
     ) -> None:
         table_exprs = [(gateway, exp.to_table(t)) for gateway, t in gateway_table_pairs]
         unique_schemas = {
@@ -1464,6 +1471,8 @@ def _create_schema(
             logger.info("Creating schema '%s'", schema)
             adapter = self.get_adapter(gateway)
             adapter.create_schema(schema)
+            if owner:
+                adapter.alter_schema_owner(schema, owner)
 
         with self.concurrent_context():
             concurrent_apply_to_values(

tests/core/engine_adapter/test_spark.py

@@ -1092,6 +1092,52 @@ def test_table_format(adapter: SparkEngineAdapter, mocker: MockerFixture):
     ]
 
 
+def test_alter_schema_owner(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    adapter.alter_schema_owner("catalog.my_schema", "svc_prod_spn")
+    assert to_sql_calls(adapter) == [
+        "ALTER SCHEMA `catalog`.`my_schema` OWNER TO `svc_prod_spn`"
+    ]
+
+
+def test_alter_schema_owner_three_part_name(make_mocked_engine_adapter: t.Callable):
+    # Schema references are typically 2-part (catalog.schema), but verify quoting is correct.
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    adapter.alter_schema_owner("my_schema", "svc_prod_spn")
+    assert to_sql_calls(adapter) == ["ALTER SCHEMA `my_schema` OWNER TO `svc_prod_spn`"]
+
+
+def test_alter_view_owner(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    adapter.alter_view_owner("catalog.my_schema.my_view", "svc_prod_spn")
+    assert to_sql_calls(adapter) == [
+        "ALTER VIEW `catalog`.`my_schema`.`my_view` OWNER TO `svc_prod_spn`"
+    ]
+
+
+def test_alter_view_owner_special_chars_in_principal(make_mocked_engine_adapter: t.Callable):
+    # Databricks Unity Catalog principals can contain colons and @ signs.
+    # Verify they are safely backtick-quoted and not interpreted as SQL syntax.
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    adapter.alter_view_owner("catalog.sushi__dev.orders", "group:devs@company.com")
+    assert to_sql_calls(adapter) == [
+        "ALTER VIEW `catalog`.`sushi__dev`.`orders` OWNER TO `group:devs@company.com`"
+    ]
+
+
+def test_alter_schema_owner_base_noop(make_mocked_engine_adapter: t.Callable):
+    # The base EngineAdapter.alter_schema_owner is a no-op: adapters that don't
+    # support ownership control silently skip it without emitting any SQL.
+    from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
+
+    adapter = make_mocked_engine_adapter(DuckDBEngineAdapter)
+    adapter.alter_schema_owner("my_schema", "some_owner")
+    adapter.alter_view_owner("my_schema.my_view", "some_owner")
+    # No ALTER SQL should have been emitted
+    alter_calls = [s for s in to_sql_calls(adapter) if "OWNER" in s.upper()]
+    assert alter_calls == []
+
+
 def test_get_data_object_wap_branch(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(SparkEngineAdapter, patch_get_data_objects=False)
     mocker.patch.object(adapter, "_get_data_objects", return_value=[])

tests/core/integration/test_config.py

@@ -15,6 +15,7 @@
     GatewayConfig,
     ModelDefaultsConfig,
     DuckDBConnectionConfig,
+    OwnershipConfig,
     TableNamingConvention,
     AutoCategorizationMode,
 )
@@ -578,3 +579,47 @@ def test_auto_categorization(sushi_context: Context):
         sushi_context.get_snapshot("sushi.waiter_as_customer_by_day", raise_if_missing=True).version
         == version
     )
+
+
+def test_ownership_config_plan_applies_without_error(
+    tmp_path: Path, monkeypatch: MonkeyPatch
+) -> None:
+    """OwnershipConfig flows through the full plan/apply lifecycle without errors.
+
+    DuckDB's alter_schema_owner/alter_view_owner are no-ops, so we cannot verify
+    that ownership was actually changed — but we confirm the config plumbing
+    doesn't break schema creation, view promotion, or dev environment application.
+    """
+    monkeypatch.chdir(tmp_path)
+
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        default_connection=DuckDBConnectionConfig(),
+        ownership=OwnershipConfig(
+            environment_owner_mapping={
+                "^prod$": "svc_prod_owner",
+                ".*": "group:shared-developers",
+            }
+        ),
+    )
+
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+    (models_dir / "model.sql").write_text(
+        """
+        MODEL (name example_schema.test_model, kind FULL);
+        SELECT '1' AS a
+        """
+    )
+
+    ctx = Context(config=config, paths=tmp_path)
+
+    # Prod plan/apply — exercises virtual layer schema and view creation with ownership config
+    ctx.plan(auto_apply=True)
+    assert ctx.engine_adapter.table_exists("example_schema.test_model")
+
+    # Dev plan/apply — exercises env-suffixed schema and view creation with ownership config
+    ctx.plan(environment="dev", include_unmodified=True, auto_apply=True)
+    metadata = DuckDBMetadata.from_context(ctx)
+    dev_schemas = {s for s in metadata.schemas if "__dev" in s}
+    assert len(dev_schemas) > 0