feat: add callable resolver support to OwnershipConfig

- OwnershipConfig gains environment_owner_resolver and
  physical_owner_resolver callable fields for cases where the owner
  principal must be resolved at plan execution time (e.g. dynamic SPN
  identity via adapter.current_user()).
- Callable resolvers take precedence over the static mapping/string
  fields when both are set.
- resolve_owner() and resolve_physical_owner() now accept the active
  EngineAdapter so callables can query the connection.
- Add current_user() to the base EngineAdapter (SELECT CURRENT_USER()).
- Replace the getattr/isinstance guard in BuiltInSchedulerConfig with
  OwnershipConfig.is_active.
- Update all resolve_owner/resolve_physical_owner call sites in
  BuiltInPlanEvaluator to thread the adapter through.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 2 people

sqlmesh/core/config/ownership.py

@@ -9,34 +9,70 @@
 from sqlmesh.core.config.common import compile_regex_mapping
 
 if t.TYPE_CHECKING:
+    from sqlmesh.core.engine_adapter.base import EngineAdapter
     OwnershipMapping = t.Dict[re.Pattern, str]
+    EnvironmentOwnerResolver = t.Callable[[str, EngineAdapter], t.Optional[str]]
+    PhysicalOwnerResolver = t.Callable[[EngineAdapter], t.Optional[str]]
 else:
     OwnershipMapping = t.Annotated[t.Dict[re.Pattern, str], BeforeValidator(compile_regex_mapping)]
+    EnvironmentOwnerResolver = t.Callable
+    PhysicalOwnerResolver = t.Callable
 
 
 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.
+    For static YAML-based config, use ``environment_owner_mapping`` and
+    ``physical_owner``.  For programmatic config where the principal must be
+    resolved at plan-execution time (e.g. via ``adapter.current_user()`` or a
+    Databricks API call), supply ``environment_owner_resolver`` and/or
+    ``physical_owner_resolver`` instead — callables take precedence over the
+    static fields.
 
-    Example::
+    Example (YAML)::
 
         ownership:
           environment_owner_mapping:
             "^prod$": "svc_prod_spn"
             ".*": "group:shared-developers"
           physical_owner: "group:shared-developers"
+
+    Example (Python)::
+
+        OwnershipConfig(
+            environment_owner_resolver=lambda env, adapter: (
+                adapter.current_user() if env == "prod" else "group:shared-developers"
+            ),
+            physical_owner="group:shared-developers",
+        )
     """
 
     environment_owner_mapping: OwnershipMapping = {}
+    environment_owner_resolver: t.Optional[EnvironmentOwnerResolver] = None
     physical_owner: t.Optional[str] = None
+    physical_owner_resolver: t.Optional[PhysicalOwnerResolver] = None
 
-    def resolve_owner(self, environment_name: str) -> t.Optional[str]:
-        """Return the configured owner for the given environment name, or None."""
+    @property
+    def is_active(self) -> bool:
+        """True when any ownership rule is configured."""
+        return bool(
+            self.environment_owner_resolver is not None
+            or self.environment_owner_mapping
+            or self.physical_owner is not None
+            or self.physical_owner_resolver is not None
+        )
+
+    def resolve_owner(self, environment_name: str, adapter: "EngineAdapter") -> t.Optional[str]:
+        """Return the configured owner for the given environment, or None."""
+        if self.environment_owner_resolver is not None:
+            return self.environment_owner_resolver(environment_name, adapter)
         for pattern, owner in self.environment_owner_mapping.items():
             if pattern.fullmatch(environment_name):
                 return owner
         return None
+
+    def resolve_physical_owner(self, adapter: "EngineAdapter") -> t.Optional[str]:
+        """Return the configured physical-layer owner, or None."""
+        if self.physical_owner_resolver is not None:
+            return self.physical_owner_resolver(adapter)
+        return self.physical_owner

sqlmesh/core/config/scheduler.py

@@ -128,14 +128,8 @@ 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
+        ownership = context.config.ownership
+        ownership_config = ownership if ownership.is_active else None
         return BuiltInPlanEvaluator(
             state_sync=context.state_sync,
             snapshot_evaluator=context.snapshot_evaluator,

sqlmesh/core/engine_adapter/base.py

@@ -1418,6 +1418,17 @@ def _create_schema(
                 raise
             logger.warning("Failed to create %s '%s': %s", kind.lower(), schema_name, e)
 
+    def current_user(self) -> str:
+        """Return the identity of the currently-connected principal.
+
+        Uses SQL ``CURRENT_USER()`` which is supported by Spark/Databricks and
+        DuckDB.  Override in adapters where a different mechanism is required.
+        """
+        row = self.fetchone("SELECT CURRENT_USER()")
+        if not row:
+            raise SQLMeshError("Could not determine current user: CURRENT_USER() returned no rows")
+        return row[0]
+
     def alter_schema_owner(self, schema_name: SchemaName, owner: str) -> None:
         """Set the owner of a schema.
 

sqlmesh/core/plan/evaluator.py

@@ -175,7 +175,11 @@ def visit_physical_layer_update_stage(
             self.console.log_success(skip_message)
             return
 
-        physical_owner = self.ownership_config.physical_owner if self.ownership_config else None
+        physical_owner = (
+            self.ownership_config.resolve_physical_owner(self.snapshot_evaluator.adapter)
+            if self.ownership_config
+            else None
+        )
         completion_status = None
         progress_stopped = False
         try:
@@ -214,7 +218,11 @@ def visit_physical_layer_update_stage(
     def visit_physical_layer_schema_creation_stage(
         self, stage: stages.PhysicalLayerSchemaCreationStage, plan: EvaluatablePlan
     ) -> None:
-        physical_owner = self.ownership_config.physical_owner if self.ownership_config else None
+        physical_owner = (
+            self.ownership_config.resolve_physical_owner(self.snapshot_evaluator.adapter)
+            if self.ownership_config
+            else None
+        )
         try:
             self.snapshot_evaluator.create_physical_schemas(
                 stage.snapshots, stage.deployability_index, owner=physical_owner
@@ -442,7 +450,9 @@ def _promote_snapshots(
     ) -> None:
         owner: t.Optional[str] = None
         if self.ownership_config:
-            owner = self.ownership_config.resolve_owner(environment_naming_info.name)
+            owner = self.ownership_config.resolve_owner(
+                environment_naming_info.name, self.snapshot_evaluator.adapter
+            )
         self.snapshot_evaluator.promote(
             target_snapshots,
             start=plan.start,

tests/core/engine_adapter/test_spark.py

@@ -1154,6 +1154,24 @@ def test_alter_schema_owner_base_noop(make_mocked_engine_adapter: t.Callable):
     assert alter_calls == []
 
 
+def test_current_user(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    adapter.cursor.fetchone.return_value = ("spn-abc-123",)
+    result = adapter.current_user()
+    assert result == "spn-abc-123"
+    sql_calls = to_sql_calls(adapter)
+    assert any("CURRENT_USER" in s.upper() for s in sql_calls)
+
+
+def test_current_user_base_noop(make_mocked_engine_adapter: t.Callable):
+    from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
+
+    adapter = make_mocked_engine_adapter(DuckDBEngineAdapter)
+    adapter.cursor.fetchone.return_value = ("duckdb-user",)
+    result = adapter.current_user()
+    assert result == "duckdb-user"
+
+
 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/test_config.py

@@ -614,52 +614,58 @@ def test_load_duckdb_attach_config(tmp_path):
 
 
 def test_ownership_config_resolve_owner():
+    mock_adapter = mock.MagicMock()
     config = OwnershipConfig(
         environment_owner_mapping={
             "^prod$": "svc_prod_spn",
             ".*": "group:shared-developers",
         }
     )
-    assert config.resolve_owner("prod") == "svc_prod_spn"
-    assert config.resolve_owner("dev_alice") == "group:shared-developers"
-    assert config.resolve_owner("staging") == "group:shared-developers"
+    assert config.resolve_owner("prod", mock_adapter) == "svc_prod_spn"
+    assert config.resolve_owner("dev_alice", mock_adapter) == "group:shared-developers"
+    assert config.resolve_owner("staging", mock_adapter) == "group:shared-developers"
     # "production" does not match ^prod$ so falls through to .*
-    assert config.resolve_owner("production") == "group:shared-developers"
+    assert config.resolve_owner("production", mock_adapter) == "group:shared-developers"
 
 
 def test_ownership_config_empty_returns_none():
-    assert OwnershipConfig().resolve_owner("prod") is None
-    assert OwnershipConfig().resolve_owner("dev_env") is None
+    mock_adapter = mock.MagicMock()
+    assert OwnershipConfig().resolve_owner("prod", mock_adapter) is None
+    assert OwnershipConfig().resolve_owner("dev_env", mock_adapter) is None
 
 
 def test_ownership_config_first_match_wins():
     # The catch-all .* comes before a more specific pattern — it always wins.
     # This documents the ordering contract: users must put specific patterns first.
+    mock_adapter = mock.MagicMock()
     config = OwnershipConfig(
         environment_owner_mapping={
             ".*": "catch_all_owner",
             "^prod$": "prod_owner",
         }
     )
-    assert config.resolve_owner("prod") == "catch_all_owner"
+    assert config.resolve_owner("prod", mock_adapter) == "catch_all_owner"
 
 
 def test_ownership_config_case_sensitive():
     # Patterns are compiled without re.IGNORECASE, so matching is case-sensitive.
+    mock_adapter = mock.MagicMock()
     config = OwnershipConfig(environment_owner_mapping={"^prod$": "svc_prod"})
-    assert config.resolve_owner("prod") == "svc_prod"
-    assert config.resolve_owner("PROD") is None
-    assert config.resolve_owner("Prod") is None
+    assert config.resolve_owner("prod", mock_adapter) == "svc_prod"
+    assert config.resolve_owner("PROD", mock_adapter) is None
+    assert config.resolve_owner("Prod", mock_adapter) is None
 
 
 def test_ownership_config_no_match_returns_none():
+    mock_adapter = mock.MagicMock()
     config = OwnershipConfig(environment_owner_mapping={"^prod$": "svc_prod"})
-    assert config.resolve_owner("staging") is None
-    assert config.resolve_owner("dev_bob") is None
+    assert config.resolve_owner("staging", mock_adapter) is None
+    assert config.resolve_owner("dev_bob", mock_adapter) is None
 
 
 def test_ownership_config_deserialization_from_dict():
     # Simulates YAML/dict-based config loading (as produced by load_config_from_yaml).
+    mock_adapter = mock.MagicMock()
     config = Config(
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
         ownership={
@@ -669,15 +675,16 @@ def test_ownership_config_deserialization_from_dict():
             }
         },
     )
-    assert config.ownership.resolve_owner("prod") == "svc_prod_spn"
-    assert config.ownership.resolve_owner("dev") == "group:shared-developers"
+    assert config.ownership.resolve_owner("prod", mock_adapter) == "svc_prod_spn"
+    assert config.ownership.resolve_owner("dev", mock_adapter) == "group:shared-developers"
 
 
 def test_ownership_config_nested_update():
     # Config.ownership uses UpdateStrategy.NESTED_UPDATE.
     # When two Configs are merged, the second one's environment_owner_mapping
     # replaces the first's (REPLACE semantics within OwnershipConfig since
     # environment_owner_mapping has no explicit strategy).
+    mock_adapter = mock.MagicMock()
     c1 = Config(
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
         ownership=OwnershipConfig(environment_owner_mapping={"^prod$": "spn_prod"}),
@@ -688,15 +695,16 @@ def test_ownership_config_nested_update():
     )
     merged = c1.update_with(c2)
     # c2's mapping fully replaces c1's — the ^prod$ pattern is gone
-    assert merged.ownership.resolve_owner("prod") == "grp_devs"
-    assert merged.ownership.resolve_owner("dev_alice") == "grp_devs"
+    assert merged.ownership.resolve_owner("prod", mock_adapter) == "grp_devs"
+    assert merged.ownership.resolve_owner("dev_alice", mock_adapter) == "grp_devs"
 
 
 def test_config_ownership_defaults_to_empty():
     # Configs without an explicit ownership block have a no-op OwnershipConfig.
+    mock_adapter = mock.MagicMock()
     config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
     assert config.ownership.environment_owner_mapping == {}
-    assert config.ownership.resolve_owner("prod") is None
+    assert config.ownership.resolve_owner("prod", mock_adapter) is None
 
 
 def test_ownership_config_physical_owner():
@@ -718,7 +726,70 @@ def test_ownership_config_physical_owner_deserialization():
         },
     )
     assert config.ownership.physical_owner == "group:data-platform"
-    assert config.ownership.resolve_owner("prod") == "svc_prod"
+    assert config.ownership.resolve_owner("prod", mock.MagicMock()) == "svc_prod"
+
+
+def test_ownership_config_resolve_owner_callable():
+    # A callable resolver takes precedence over environment_owner_mapping and
+    # receives (env_name, adapter) so it can call adapter.current_user() etc.
+    mock_adapter = mock.MagicMock()
+    mock_adapter.current_user.return_value = "spn-dynamic-uuid"
+
+    config = OwnershipConfig(
+        environment_owner_mapping={".*": "group:fallback"},
+        environment_owner_resolver=lambda env, adapter: (
+            adapter.current_user() if env == "prod" else "group:shared-developers"
+        ),
+    )
+
+    assert config.resolve_owner("prod", mock_adapter) == "spn-dynamic-uuid"
+    assert config.resolve_owner("dev_alice", mock_adapter) == "group:shared-developers"
+    mock_adapter.current_user.assert_called_once()
+
+
+def test_ownership_config_resolver_overrides_mapping():
+    # Resolver always wins when set, even if the mapping would also match.
+    mock_adapter = mock.MagicMock()
+    config = OwnershipConfig(
+        environment_owner_mapping={"^prod$": "static-owner"},
+        environment_owner_resolver=lambda env, adapter: "dynamic-owner",
+    )
+    assert config.resolve_owner("prod", mock_adapter) == "dynamic-owner"
+
+
+def test_ownership_config_resolve_physical_owner_callable():
+    mock_adapter = mock.MagicMock()
+    mock_adapter.current_user.return_value = "spn-uuid-123"
+
+    config = OwnershipConfig(
+        physical_owner_resolver=lambda adapter: adapter.current_user(),
+    )
+    assert config.resolve_physical_owner(mock_adapter) == "spn-uuid-123"
+    mock_adapter.current_user.assert_called_once()
+
+
+def test_ownership_config_resolve_physical_owner_static():
+    mock_adapter = mock.MagicMock()
+    config = OwnershipConfig(physical_owner="group:data-platform")
+    assert config.resolve_physical_owner(mock_adapter) == "group:data-platform"
+    mock_adapter.current_user.assert_not_called()
+
+
+def test_ownership_config_physical_owner_resolver_overrides_static():
+    mock_adapter = mock.MagicMock()
+    config = OwnershipConfig(
+        physical_owner="static-owner",
+        physical_owner_resolver=lambda adapter: "dynamic-owner",
+    )
+    assert config.resolve_physical_owner(mock_adapter) == "dynamic-owner"
+
+
+def test_ownership_config_is_active():
+    assert not OwnershipConfig().is_active
+    assert OwnershipConfig(environment_owner_mapping={".*": "grp"}).is_active
+    assert OwnershipConfig(environment_owner_resolver=lambda e, a: None).is_active
+    assert OwnershipConfig(physical_owner="grp").is_active
+    assert OwnershipConfig(physical_owner_resolver=lambda a: "grp").is_active
 
 
 def test_load_model_defaults_audits(tmp_path):