diff --git a/setup.cfg b/setup.cfg index 569c3480..74f37943 100644 --- a/setup.cfg +++ b/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 = diff --git a/src/preset_cli/cli/superset/sync/dbt/datasets.py b/src/preset_cli/cli/superset/sync/dbt/datasets.py index 000ec466..44bd3200 100644 --- a/src/preset_cli/cli/superset/sync/dbt/datasets.py +++ b/src/preset_cli/cli/superset/sync/dbt/datasets.py @@ -8,7 +8,6 @@ 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 @@ -16,7 +15,9 @@ 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"} @@ -53,6 +54,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], @@ -72,7 +74,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 = { diff --git a/src/preset_cli/cli/superset/sync/dbt/lib.py b/src/preset_cli/cli/superset/sync/dbt/lib.py index 5e0e1480..5757f58f 100644 --- a/src/preset_cli/cli/superset/sync/dbt/lib.py +++ b/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``. diff --git a/tests/cli/superset/sync/dbt/datasets_test.py b/tests/cli/superset/sync/dbt/datasets_test.py index cc8825e6..69db74b4 100644 --- a/tests/cli/superset/sync/dbt/datasets_test.py +++ b/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. diff --git a/tests/cli/superset/sync/dbt/lib_test.py b/tests/cli/superset/sync/dbt/lib_test.py index 61d35530..4dde45bb 100644 --- a/tests/cli/superset/sync/dbt/lib_test.py +++ b/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.