Merge pull request #259 from preset-io/get_models_from_sql

chore(dbt): refactor get_model_from_sql

Author: betodealmeida

src/preset_cli/cli/superset/sync/dbt/command.py

@@ -26,8 +26,7 @@
 from preset_cli.cli.superset.sync.dbt.exposures import ModelKey, sync_exposures
 from preset_cli.cli.superset.sync.dbt.lib import apply_select
 from preset_cli.cli.superset.sync.dbt.metrics import (
-    MultipleModelsError,
-    get_model_from_sql,
+    get_models_from_sql,
     get_superset_metrics_per_model,
 )
 from preset_cli.exceptions import DatabaseNotFoundError
@@ -351,10 +350,10 @@ def process_sl_metrics(
         if sql is None:
             continue
 
-        try:
-            model = get_model_from_sql(sql, dialect, model_map)
-        except MultipleModelsError:
+        models = get_models_from_sql(sql, dialect, model_map)
+        if len(models) > 1:
             continue
+        model = models[0]
 
         sl_metrics.append(
             mf_metric_schema.load(

src/preset_cli/cli/superset/sync/dbt/metrics.py

@@ -325,28 +325,19 @@ def convert_metric_flow_to_superset(
     }
 
 
-class MultipleModelsError(Exception):
-    """
-    Raised when a metric depends on multiple models.
-    """
-
-
-def get_model_from_sql(
+def get_models_from_sql(
     sql: str,
     dialect: MFSQLEngine,
     model_map: Dict[ModelKey, ModelSchema],
-) -> ModelSchema:
+) -> List[ModelSchema]:
     """
     Return the model associated with a SQL query.
     """
     parsed_query = parse_one(sql, dialect=DIALECT_MAP.get(dialect))
     sources = list(parsed_query.find_all(Table))
-    if len(sources) > 1:
-        raise MultipleModelsError(
-            f"Unable to convert metrics with multiple sources: {sql}",
-        )
 
-    table = sources[0]
-    key = ModelKey(table.db, table.name)
+    for table in sources:
+        if ModelKey(table.db, table.name) not in model_map:
+            raise ValueError(f"Unable to find model for SQL source {table}")
 
-    return model_map[key]
+    return [model_map[ModelKey(table.db, table.name)] for table in sources]

tests/cli/superset/sync/dbt/command_test.py

@@ -149,6 +149,14 @@
         "schema": "public",
         "unique_id": "model.superset_examples.messages_channels",
     },
+    {
+        "database": "some_other_table",
+        "description": "",
+        "meta": {},
+        "name": "some_other_table",
+        "schema": "public",
+        "unique_id": "model.superset_examples.some_other_table",
+    },
 ]
 
 dbt_cloud_metrics = [
@@ -975,7 +983,7 @@ def test_dbt_cloud(mocker: MockerFixture) -> None:
     dbt_client.get_sl_metrics.return_value = dbt_metricflow_metrics
     dbt_client.get_sl_metric_sql.side_effect = [
         "SELECT COUNT(*) FROM public.messages_channels",
-        "SELECT COUNT(*) FROM public.messages_channels JOIN some_other_table",
+        "SELECT COUNT(*) FROM public.messages_channels JOIN public.some_other_table",
         None,
     ]
     database = mocker.MagicMock()
@@ -1624,6 +1632,10 @@ def test_dbt_cloud_exposures_only(mocker: MockerFixture, fs: FakeFilesystem) ->
         exposures,
         [
             {"schema": "public", "table_name": "messages_channels"},
+            {"schema": "public", "table_name": "some_other_table"},
         ],
-        {("public", "messages_channels"): dbt_cloud_models[0]},
+        {
+            ("public", "messages_channels"): dbt_cloud_models[0],
+            ("public", "some_other_table"): dbt_cloud_models[1],
+        },
     )

tests/cli/superset/sync/dbt/metrics_test.py

@@ -12,13 +12,12 @@
 from preset_cli.api.clients.dbt import MetricSchema, MFMetricWithSQLSchema, MFSQLEngine
 from preset_cli.cli.superset.sync.dbt.exposures import ModelKey
 from preset_cli.cli.superset.sync.dbt.metrics import (
-    MultipleModelsError,
     convert_metric_flow_to_superset,
     convert_query_to_projection,
     get_metric_expression,
     get_metric_models,
     get_metrics_for_model,
-    get_model_from_sql,
+    get_models_from_sql,
     get_superset_metrics_per_model,
 )
 
@@ -649,30 +648,31 @@ def test_convert_metric_flow_to_superset(mocker: MockerFixture) -> None:
     }
 
 
-def test_get_model_from_sql() -> None:
+def test_get_models_from_sql() -> None:
     """
-    Test the ``get_model_from_sql`` function.
+    Test the ``get_models_from_sql`` function.
     """
     model_map = {
         ModelKey("schema", "table"): {"name": "table"},
+        ModelKey("schema", "a"): {"name": "a"},
+        ModelKey("schema", "b"): {"name": "b"},
     }
 
-    assert get_model_from_sql(
+    assert get_models_from_sql(
         "SELECT 1 FROM project.schema.table",
         MFSQLEngine.BIGQUERY,
         model_map,  # type: ignore
-    ) == {"name": "table"}
+    ) == [{"name": "table"}]
 
-    with pytest.raises(MultipleModelsError) as excinfo:
-        get_model_from_sql(
-            "SELECT 1 FROM schema.a JOIN schema.b",
-            MFSQLEngine.BIGQUERY,
-            {},
-        )
-    assert (
-        str(excinfo.value)
-        == "Unable to convert metrics with multiple sources: SELECT 1 FROM schema.a JOIN schema.b"
-    )
+    assert get_models_from_sql(
+        "SELECT 1 FROM schema.a JOIN schema.b",
+        MFSQLEngine.BIGQUERY,
+        model_map,  # type: ignore
+    ) == [{"name": "a"}, {"name": "b"}]
+
+    with pytest.raises(ValueError) as excinfo:
+        get_models_from_sql("SELECT 1 FROM schema.c", MFSQLEngine.BIGQUERY, {})
+    assert str(excinfo.value) == "Unable to find model for SQL source schema.c"
 
 
 def test_get_superset_metrics_per_model() -> None: