chore(dbt): Instruct users of missing dependency when creating a Snow…

…flake virtual dataset (#290)

* Initial implementation
* chore(dbt): Instruct users of missing dependency when creating a Snowflake virtual dataset
* Fixing tests
* Addressing PR feedback

setup.cfg

@@ -84,9 +84,8 @@ exclude =
     tests
 
 [options.extras_require]
-# Add here additional requirements for extra features, to install with:
-# `pip install preset-cli[PDF]` like:
-# PDF = ReportLab; RXP
+# TODO: Implement additional optional dependencies
+snowflake = snowflake-sqlalchemy==1.4.4
 
 # Add here test requirements (semicolon/line-separated)
 testing =

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

@@ -8,15 +8,16 @@
 import logging
 from typing import Any, Dict, List, Optional, Tuple
 
-from sqlalchemy.engine import create_engine
 from sqlalchemy.engine.url import URL as SQLAlchemyURL
 from sqlalchemy.engine.url import make_url
 from yarl import URL
 
 from preset_cli.api.clients.dbt import ModelSchema
 from preset_cli.api.clients.superset import SupersetClient, SupersetMetricDefinition
 from preset_cli.api.operators import OneToMany
+from preset_cli.cli.superset.sync.dbt.lib import create_engine_with_check
 from preset_cli.exceptions import CLIError, SupersetError
+from preset_cli.lib import raise_cli_errors
 
 DEFAULT_CERTIFICATION = {"details": "This table is produced by dbt"}
 
@@ -56,6 +57,7 @@ def clean_metadata(metadata: Dict[str, Any]) -> Dict[str, Any]:
     return metadata
 
 
+@raise_cli_errors
 def create_dataset(
     client: SupersetClient,
     database: Dict[str, Any],
@@ -75,7 +77,7 @@ def create_dataset(
             "table_name": model.get("alias") or model["name"],
         }
     else:
-        engine = create_engine(url)
+        engine = create_engine_with_check(url)
         quote = engine.dialect.identifier_preparer.quote
         source = ".".join(quote(model[key]) for key in ("database", "schema", "name"))
         kwargs = {

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

@@ -12,9 +12,12 @@
 
 import yaml
 from jinja2 import Environment
+from sqlalchemy.engine import Engine, create_engine
 from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import NoSuchModuleError
 
 from preset_cli.api.clients.dbt import ModelSchema
+from preset_cli.exceptions import CLIError
 
 _logger = logging.getLogger(__name__)
 
@@ -198,6 +201,31 @@ def build_snowflake_sqlalchemy_params(target: Dict[str, Any]) -> Dict[str, Any]:
     return parameters
 
 
+def create_engine_with_check(url: URL) -> Engine:
+    """
+    Returns a SQLAlchemy engine or raises an error if missing required dependency.
+    """
+    try:
+        return create_engine(url)
+    except NoSuchModuleError as exc:
+        string_url = str(url)
+        dialect = string_url.split("://", maxsplit=1)[0]
+        # TODO: Handle more DB engines that require an additional package
+        if dialect == "snowflake":
+            raise CLIError(
+                (
+                    "Missing required package. "
+                    'Please run ``pip install "preset-cli[snowflake]"`` to install it.'
+                ),
+                1,
+            ) from exc
+        raise NotImplementedError(
+            f"Unable to build a SQLAlchemy Engine for the {dialect} connection. Please file an "
+            "issue at https://github.com/preset-io/backend-sdk/issues/new?labels=enhancement&"
+            f"title=Missing+package+for+{dialect}.",
+        ) from exc
+
+
 def env_var(var: str, default: Optional[str] = None) -> str:
     """
     Simplified version of dbt's ``env_var``.

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

@@ -769,7 +769,7 @@ def test_create_dataset_virtual(mocker: MockerFixture) -> None:
     Test ``create_dataset`` for virtual datasets.
     """
     create_engine = mocker.patch(
-        "preset_cli.cli.superset.sync.dbt.datasets.create_engine",
+        "preset_cli.cli.superset.sync.dbt.lib.create_engine",
     )
     create_engine().dialect.identifier_preparer.quote = lambda token: token
     client = mocker.MagicMock()
@@ -792,6 +792,45 @@ def test_create_dataset_virtual(mocker: MockerFixture) -> None:
     )
 
 
+def test_create_dataset_virtual_missing_dependency(
+    capsys: pytest.CaptureFixture[str],
+    mocker: MockerFixture,
+) -> None:
+    """
+    Test ``create_dataset`` for virtual datasets when the DB connection requires
+    an additional package.
+    """
+    client = mocker.MagicMock()
+
+    with pytest.raises(NotImplementedError):
+        create_dataset(
+            client,
+            {
+                "id": 1,
+                "schema": "public",
+                "name": "Database",
+                "sqlalchemy_uri": "blah://user@host/examples",
+            },
+            models[0],
+        )
+
+    with pytest.raises(SystemExit) as excinfo:
+        create_dataset(
+            client,
+            {
+                "id": 1,
+                "schema": "public",
+                "name": "other_db",
+                "sqlalchemy_uri": "snowflake://user@host/examples",
+            },
+            models[0],
+        )
+    output_content = capsys.readouterr()
+    assert excinfo.type == SystemExit
+    assert excinfo.value.code == 1
+    assert "preset-cli[snowflake]" in output_content.out
+
+
 def test_get_or_create_dataset_existing_dataset(mocker: MockerFixture) -> None:
     """
     Test the ``get_or_create_dataset`` helper when it finds a dataset.

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

@@ -11,17 +11,20 @@
 import pytest
 from pyfakefs.fake_filesystem import FakeFilesystem
 from pytest_mock import MockerFixture
+from sqlalchemy.engine.url import URL
 
 from preset_cli.api.clients.dbt import ModelSchema
 from preset_cli.cli.superset.sync.dbt.lib import (
     apply_select,
     as_number,
     build_sqlalchemy_params,
+    create_engine_with_check,
     env_var,
     filter_models,
     list_failed_models,
     load_profiles,
 )
+from preset_cli.exceptions import CLIError
 
 
 def test_build_sqlalchemy_params_postgres(mocker: MockerFixture) -> None:
@@ -225,6 +228,37 @@ def test_build_sqlalchemy_params_unsupported() -> None:
     )
 
 
+def test_create_engine_with_check(mocker: MockerFixture) -> None:
+    """
+    Test the ``create_engine_with_check`` method.
+    """
+    mock_engine = mocker.patch("preset_cli.cli.superset.sync.dbt.lib.create_engine")
+    test = create_engine_with_check(URL("blah://blah"))
+    assert test == mock_engine.return_value
+
+
+def test_create_engine_with_check_missing_snowflake() -> None:
+    """
+    Test the ``create_engine_with_check`` method when the Snowflake driver is
+    not installed.
+    """
+    with pytest.raises(CLIError) as excinfo:
+        create_engine_with_check(URL("snowflake://blah"))
+    assert 'run ``pip install "preset-cli[snowflake]"``' in str(excinfo.value)
+
+
+def test_create_engine_with_check_missing_unknown_driver() -> None:
+    """
+    Test the ``create_engine_with_check`` method when a SQLAlchemy driver is
+    not installed.
+    """
+    with pytest.raises(NotImplementedError) as excinfo:
+        create_engine_with_check(URL("mssql+odbc://blah"))
+    assert "Unable to build a SQLAlchemy Engine for the mssql+odbc connection" in str(
+        excinfo.value,
+    )
+
+
 def test_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
     """
     Test the ``env_var`` implementation.