PR feedback

Author: erindru

sqlmesh/core/context.py

@@ -93,7 +93,7 @@
 from sqlmesh.core.reference import ReferenceGraph
 from sqlmesh.core.scheduler import Scheduler, CompletionStatus
 from sqlmesh.core.schema_loader import create_external_models_file
-from sqlmesh.core.selector import Selector
+from sqlmesh.core.selector import Selector, NativeSelector
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
     Snapshot,
@@ -348,8 +348,6 @@ class GenericContext(BaseContext, t.Generic[C]):
         load: Whether or not to automatically load all models and macros (default True).
         console: The rich instance used for printing out CLI command results.
         users: A list of users to make known to SQLMesh.
-        dbt_mode: A flag to indicate we are running in 'dbt mode' which means that things like
-            model selections should use the dbt names and not the native SQLMesh names
     """
 
     CONFIG_TYPE: t.Type[C]
@@ -370,7 +368,7 @@ def __init__(
         load: bool = True,
         users: t.Optional[t.List[User]] = None,
         config_loader_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
-        dbt_mode: bool = False,
+        selector: t.Optional[t.Type[Selector]] = None,
     ):
         self.configs = (
             config
@@ -393,7 +391,7 @@ def __init__(
         self._engine_adapter: t.Optional[EngineAdapter] = None
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
-        self._dbt_mode = dbt_mode
+        self._selector_cls = selector or NativeSelector
 
         self.path, self.config = t.cast(t.Tuple[Path, C], next(iter(self.configs.items())))
 
@@ -2897,15 +2895,14 @@ def _new_state_sync(self) -> StateSync:
     def _new_selector(
         self, models: t.Optional[UniqueKeyDict[str, Model]] = None, dag: t.Optional[DAG[str]] = None
     ) -> Selector:
-        return Selector(
+        return self._selector_cls(
             self.state_reader,
             models=models or self._models,
             context_path=self.path,
             dag=dag,
             default_catalog=self.default_catalog,
             dialect=self.default_dialect,
             cache_dir=self.cache_dir,
-            dbt_mode=self._dbt_mode,
         )
 
     def _register_notification_targets(self) -> None:

sqlmesh/core/selector.py

@@ -4,6 +4,7 @@
 import typing as t
 from pathlib import Path
 from itertools import zip_longest
+import abc
 
 from sqlglot import exp
 from sqlglot.errors import ParseError
@@ -27,7 +28,7 @@
     from sqlmesh.core.state_sync import StateReader
 
 
-class Selector:
+class Selector(abc.ABC):
     def __init__(
         self,
         state_reader: StateReader,
@@ -37,7 +38,6 @@ def __init__(
         default_catalog: t.Optional[str] = None,
         dialect: t.Optional[str] = None,
         cache_dir: t.Optional[Path] = None,
-        dbt_mode: bool = False,
     ):
         self._state_reader = state_reader
         self._models = models
@@ -46,7 +46,6 @@ def __init__(
         self._default_catalog = default_catalog
         self._dialect = dialect
         self._git_client = GitClient(context_path)
-        self._dbt_mode = dbt_mode
 
         if dag is None:
             self._dag: DAG[str] = DAG()
@@ -243,26 +242,37 @@ def evaluate(node: exp.Expression) -> t.Set[str]:
 
         return evaluate(node)
 
-    def _model_fqn(self, model: Model) -> str:
-        if self._dbt_mode:
-            dbt_fqn = model.dbt_fqn
-            if dbt_fqn is None:
-                raise SQLMeshError("Expecting dbt node information to be populated; it wasnt")
-            return dbt_fqn
-        return model.fqn
+    @abc.abstractmethod
+    def _model_name(self, model: Model) -> str:
+        """Given a model, return the name that a selector pattern contining wildcards should be fnmatch'd on"""
+        pass
+
+    @abc.abstractmethod
+    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
+        """Given a pattern, return the keys of the matching models from :all_models"""
+        pass
+
+
+class NativeSelector(Selector):
+    """Implementation of selectors that matches objects based on SQLMesh native names"""
 
     def _model_name(self, model: Model) -> str:
-        if self._dbt_mode:
-            # dbt always matches on the fqn, not the name
-            return self._model_fqn(model)
         return model.name
 
     def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
-        # note: all_models should be keyed by sqlmesh fqn, not dbt fqn
-        if not self._dbt_mode:
-            fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
-            return {fqn} if fqn in all_models else set()
+        fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
+        return {fqn} if fqn in all_models else set()
 
+
+class DbtSelector(Selector):
+    """Implementation of selectors that matches objects based on the DBT names instead of the SQLMesh native names"""
+
+    def _model_name(self, model: Model) -> str:
+        if dbt_fqn := model.dbt_fqn:
+            return dbt_fqn
+        raise SQLMeshError("dbt node information must be populated to use dbt selectors")
+
+    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
         # a pattern like "staging.customers" should match a model called "jaffle_shop.staging.customers"
         # but not a model called "jaffle_shop.customers.staging"
         # also a pattern like "aging" should not match "staging" so we need to consider components; not substrings

sqlmesh_dbt/operations.py

@@ -231,6 +231,7 @@ def create(
         from sqlmesh.core.console import set_console
         from sqlmesh_dbt.console import DbtCliConsole
         from sqlmesh.utils.errors import SQLMeshError
+        from sqlmesh.core.selector import DbtSelector
 
         # clear any existing handlers set up by click/rich as defaults so that once SQLMesh logging config is applied,
         # we dont get duplicate messages logged from things like console.log_warning()
@@ -250,8 +251,8 @@ def create(
             paths=[project_dir],
             config_loader_kwargs=dict(profile=profile, target=target, variables=vars),
             load=True,
-            # dbt mode enables selectors to use dbt model fqn's rather than SQLMesh model names
-            dbt_mode=True,
+            # DbtSelector selects based on dbt model fqn's rather than SQLMesh model names
+            selector=DbtSelector,
         )
 
         dbt_loader = sqlmesh_context._loaders[0]

tests/core/test_selector.py

@@ -12,7 +12,7 @@
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.model import Model, SqlModel
 from sqlmesh.core.model.common import ParsableSql
-from sqlmesh.core.selector import Selector
+from sqlmesh.core.selector import NativeSelector
 from sqlmesh.core.snapshot import SnapshotChangeCategory
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.date import now_timestamp
@@ -88,7 +88,7 @@ def test_select_models(mocker: MockerFixture, make_snapshot, default_catalog: t.
     local_models[modified_model_v2.fqn] = modified_model_v2.copy(
         update={"mapping_schema": added_model_schema}
     )
-    selector = Selector(state_reader_mock, local_models, default_catalog=default_catalog)
+    selector = NativeSelector(state_reader_mock, local_models, default_catalog=default_catalog)
 
     _assert_models_equal(
         selector.select_models(["db.added_model"], env_name),
@@ -243,7 +243,7 @@ def test_select_models_expired_environment(mocker: MockerFixture, make_snapshot)
 
     local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
     local_models[modified_model_v2.fqn] = modified_model_v2
-    selector = Selector(state_reader_mock, local_models)
+    selector = NativeSelector(state_reader_mock, local_models)
 
     _assert_models_equal(
         selector.select_models(["*.modified_model"], env_name, fallback_env_name="prod"),
@@ -305,7 +305,7 @@ def test_select_change_schema(mocker: MockerFixture, make_snapshot):
     local_child = child.copy(update={"mapping_schema": {'"db"': {'"parent"': {"b": "INT"}}}})
     local_models[local_child.fqn] = local_child
 
-    selector = Selector(state_reader_mock, local_models)
+    selector = NativeSelector(state_reader_mock, local_models)
 
     selected = selector.select_models(["db.parent"], env_name)
     assert selected[local_child.fqn].render_query() != child.render_query()
@@ -339,7 +339,7 @@ def test_select_models_missing_env(mocker: MockerFixture, make_snapshot):
     local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
     local_models[model.fqn] = model
 
-    selector = Selector(state_reader_mock, local_models)
+    selector = NativeSelector(state_reader_mock, local_models)
 
     assert selector.select_models([model.name], "missing_env").keys() == {model.fqn}
     assert not selector.select_models(["missing"], "missing_env")
@@ -563,7 +563,7 @@ def test_expand_model_selections(
         )
         models[model.fqn] = model
 
-    selector = Selector(mocker.Mock(), models)
+    selector = NativeSelector(mocker.Mock(), models)
     assert selector.expand_model_selections(selections) == output
 
 
@@ -576,7 +576,7 @@ def test_model_selection_normalized(mocker: MockerFixture, make_snapshot):
         dialect="bigquery",
     )
     models[model.fqn] = model
-    selector = Selector(mocker.Mock(), models, dialect="bigquery")
+    selector = NativeSelector(mocker.Mock(), models, dialect="bigquery")
     assert selector.expand_model_selections(["db.test_Model"]) == {'"db"."test_Model"'}
 
 
@@ -624,7 +624,7 @@ def test_expand_git_selection(
     git_client_mock.list_uncommitted_changed_files.return_value = []
     git_client_mock.list_committed_changed_files.return_value = [model_a._path, model_c._path]
 
-    selector = Selector(mocker.Mock(), models)
+    selector = NativeSelector(mocker.Mock(), models)
     selector._git_client = git_client_mock
 
     assert selector.expand_model_selections(expressions) == expected_fqns
@@ -658,7 +658,7 @@ def test_select_models_with_external_parent(mocker: MockerFixture):
     local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
     local_models[added_model.fqn] = added_model
 
-    selector = Selector(state_reader_mock, local_models, default_catalog=default_catalog)
+    selector = NativeSelector(state_reader_mock, local_models, default_catalog=default_catalog)
 
     expanded_selections = selector.expand_model_selections(["+*added_model*"])
     assert expanded_selections == {added_model.fqn}
@@ -699,7 +699,7 @@ def test_select_models_local_tags_take_precedence_over_remote(
     local_models[local_existing.fqn] = local_existing
     local_models[local_new.fqn] = local_new
 
-    selector = Selector(state_reader_mock, local_models)
+    selector = NativeSelector(state_reader_mock, local_models)
 
     selected = selector.select_models(["tag:a"], env_name)
 

tests/dbt/cli/test_selectors.py

@@ -1,6 +1,7 @@
 import typing as t
 import pytest
 from sqlmesh_dbt import selectors
+from sqlmesh.core.selector import DbtSelector
 from sqlmesh.core.context import Context
 from pathlib import Path
 
@@ -112,6 +113,7 @@ def test_split_unions_and_intersections(
             ["customers+", "stg_orders"],
             {'"jaffle_shop"."main"."customers"', '"jaffle_shop"."main"."stg_orders"'},
         ),
+        (["*.staging.stg_c*"], {'"jaffle_shop"."main"."stg_customers"'}),
         (["tag:agg"], {'"jaffle_shop"."main"."agg_orders"'}),
         (
             ["staging.stg_customers", "tag:agg"],
@@ -137,6 +139,16 @@ def test_split_unions_and_intersections(
                 '"jaffle_shop"."main"."agg_orders"',
             },
         ),
+        (
+            ["tag:b*"],
+            set(),
+        ),
+        (
+            ["tag:a*"],
+            {
+                '"jaffle_shop"."main"."agg_orders"',
+            },
+        ),
     ],
 )
 def test_select_by_dbt_names(
@@ -155,7 +167,7 @@ def test_select_by_dbt_names(
     assert '"jaffle_shop"."main"."agg_orders"' in ctx.models
 
     selector = ctx._new_selector()
-    assert selector._dbt_mode
+    assert isinstance(selector, DbtSelector)
 
     sqlmesh_selector = selectors.to_sqlmesh(dbt_select=dbt_select, dbt_exclude=[])
     assert sqlmesh_selector
@@ -205,7 +217,7 @@ def test_exclude_by_dbt_names(
     assert '"jaffle_shop"."main"."agg_orders"' in ctx.models
 
     selector = ctx._new_selector()
-    assert selector._dbt_mode
+    assert isinstance(selector, DbtSelector)
 
     sqlmesh_selector = selectors.to_sqlmesh(dbt_select=[], dbt_exclude=dbt_exclude)
     assert sqlmesh_selector
@@ -251,7 +263,7 @@ def test_selection_and_exclusion_by_dbt_names(
     assert '"jaffle_shop"."main"."agg_orders"' in ctx.models
 
     selector = ctx._new_selector()
-    assert selector._dbt_mode
+    assert isinstance(selector, DbtSelector)
 
     sqlmesh_selector = selectors.to_sqlmesh(dbt_select=dbt_select, dbt_exclude=dbt_exclude)
     assert sqlmesh_selector

tests/dbt/conftest.py

@@ -7,6 +7,7 @@
 import pytest
 
 from sqlmesh.core.context import Context
+from sqlmesh.core.selector import DbtSelector
 from sqlmesh.dbt.context import DbtContext
 from sqlmesh.dbt.project import Project
 from sqlmesh.dbt.target import PostgresConfig
@@ -99,7 +100,7 @@ def jaffle_shop_duckdb(copy_to_temp_path: t.Callable[..., t.List[Path]]) -> t.It
 @pytest.fixture
 def jaffle_shop_duckdb_context(jaffle_shop_duckdb: Path) -> Context:
     init_project_if_required(jaffle_shop_duckdb)
-    return Context(paths=[jaffle_shop_duckdb], dbt_mode=True)
+    return Context(paths=[jaffle_shop_duckdb], selector=DbtSelector)
 
 
 @pytest.fixture()