Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore(dbt): Instruct users of missing dependency when creating a Snowflake virtual dataset #290

Merged
merged 4 commits into from
May 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
6 changes: 4 additions & 2 deletions src/preset_cli/cli/superset/sync/dbt/datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"}

Expand Down Expand Up @@ -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],
Expand All @@ -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 = {
Expand Down
28 changes: 28 additions & 0 deletions src/preset_cli/cli/superset/sync/dbt/lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -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
Comment on lines +214 to +221
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the future this could be something like:

if dialect in {"snowflake", "redshift", "bigquery"}:
    raise CLIError(
        (
            "Missing required package. "
            f'Please run ``pip install "preset-cli[{dialect}]"`` 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``.
Expand Down
41 changes: 40 additions & 1 deletion tests/cli/superset/sync/dbt/datasets_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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.
Expand Down
34 changes: 34 additions & 0 deletions tests/cli/superset/sync/dbt/lib_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down
Loading