From 426a7149a401ca1dac16b3ba6e272d7c436bb2b8 Mon Sep 17 00:00:00 2001 From: skshetry <18718008+skshetry@users.noreply.github.com> Date: Fri, 5 Jan 2024 19:16:12 +0545 Subject: [PATCH] database: drop dbt implementation (#10222) * remove dbt implementation in the database * lazy load client * fixup * fixup * fix server side cursors * use exec_driver_sql for testing connection --- dvc/commands/imp_db.py | 137 ++------------- dvc/config_schema.py | 2 - dvc/database.py | 124 ++++++++++++++ dvc/database/__init__.py | 22 --- dvc/database/dbt_models.py | 40 ----- dvc/database/dbt_query.py | 102 ----------- dvc/database/dbt_runner.py | 61 ------- dvc/database/dbt_utils.py | 100 ----------- dvc/database/serializer.py | 54 ------ dvc/database/sqla.py | 85 --------- dvc/dependency/__init__.py | 11 +- dvc/dependency/db.py | 285 ++++--------------------------- dvc/repo/imp_db.py | 47 ++--- dvc/stage/__init__.py | 10 +- pyproject.toml | 3 +- tests/func/test_import_db.py | 215 +++++++++-------------- tests/func/test_import_db_e2e.py | 147 ---------------- 17 files changed, 285 insertions(+), 1160 deletions(-) create mode 100644 dvc/database.py delete mode 100644 dvc/database/__init__.py delete mode 100644 dvc/database/dbt_models.py delete mode 100644 dvc/database/dbt_query.py delete mode 100644 dvc/database/dbt_runner.py delete mode 100644 dvc/database/dbt_utils.py delete mode 100644 dvc/database/serializer.py delete mode 100644 dvc/database/sqla.py delete mode 100644 tests/func/test_import_db_e2e.py diff --git a/dvc/commands/imp_db.py b/dvc/commands/imp_db.py index d24e03680..d4f6b3d5a 100644 --- a/dvc/commands/imp_db.py +++ b/dvc/commands/imp_db.py @@ -1,9 +1,7 @@ import argparse -from funcy import compact, merge - from dvc.cli import completion -from dvc.cli.command import CmdBase +from dvc.cli.command import CmdBase, CmdBaseNoRepo from dvc.cli.utils import append_doc_link from dvc.log import logger from dvc.ui import ui @@ -11,94 +9,35 @@ from dvc.ui import ui logger = logger.getChild(__name__) -class CmdTestDb(CmdBase): +class CmdTestDb(CmdBaseNoRepo): def run(self): - from dvc.database import get_client - from dvc.database.dbt_utils import DBT_PROJECT_FILE, is_dbt_project - from dvc.dependency.db import _get_dbt_config + from dvc.config import Config + from dvc.database import client from dvc.exceptions import DvcException connection = self.args.conn - - db_config = self.repo.config.get("db", {}) - config = db_config.get(connection, {}) - if connection and not config: + db_config = Config.from_cwd().get("db", {}) + if connection not in db_config: raise DvcException(f"connection {connection} not found in config") - cli_config = compact( - { - "url": self.args.url, - "username": self.args.username, - "password": self.args.password, - } - ) - conn_config = merge(config, cli_config) - - cli_dbt_config = compact( - {"profile": self.args.profile, "target": self.args.target} - ) - dbt_config = merge(_get_dbt_config(self.repo.config), cli_dbt_config) - - project_dir = self.repo.root_dir - if not (conn_config or dbt_config): - if not self.args.dbt_conn: - raise DvcException( - "no config set; provide arguments or set a configuration" - ) - - if is_dbt_project(project_dir): - ui.write("Using", DBT_PROJECT_FILE, "for testing", styled=True) - else: - raise DvcException( - f"no config set and {DBT_PROJECT_FILE} is missing; " - "provide arguments or set a configuration" - ) - - adapter = get_client(conn_config, project_dir=project_dir, **dbt_config) - with adapter as db: + config = db_config.get(connection, {}) + if self.args.url: + config["url"] = self.args.url + if self.args.username: + config["username"] = self.args.username + if self.args.password: + config["password"] = self.args.password + with client(config) as db: ui.write(f"Testing with {db}", styled=True) - - creds = getattr(db, "creds", {}) - for k, v in creds.items(): - ui.write("\t", f"{k}:", v, styled=True) - - if creds: - ui.write() - db.test_connection() - ui.write("Connection successful", styled=True) class CmdImportDb(CmdBase): def run(self): - from dvc.exceptions import InvalidArgumentError - - if self.args.table or self.args.sql: - arg = "--table" if self.args.table else "--sql" - options = { - "url": self.args.url, - "rev": self.args.rev, - "project_dir": self.args.project_dir, - } - opt = next((o for o, v in options.items() if v), None) - if opt: - raise InvalidArgumentError(f"argument {opt}: not allowed with {arg}") - - if not self.args.conn and not self.args.dbt_conn: - raise InvalidArgumentError(f"{arg} requires --conn") - if self.args.model and self.args.conn: - raise InvalidArgumentError("argument --model: not allowed with --conn") - self.repo.imp_db( - url=self.args.url, - rev=self.args.rev, - project_dir=self.args.project_dir, sql=self.args.sql, table=self.args.table, - model=self.args.model, - profile=self.args.profile, - target=self.args.target, output_format=self.args.output_format, out=self.args.out, force=self.args.force, @@ -116,43 +55,9 @@ def add_parser(subparsers, parent_parser): help=IMPORT_HELP, formatter_class=argparse.RawTextHelpFormatter, ) - import_parser.add_argument( - "--url", - help=argparse.SUPPRESS, - # help="Location of dbt repository", - ) - import_parser.add_argument( - "--rev", - nargs="?", - help=argparse.SUPPRESS, - # help="Git revision (e.g. SHA, branch, tag)", - metavar="", - ) - import_parser.add_argument( - "--project-dir", - nargs="?", - help=argparse.SUPPRESS, - # help="Subdirectory to the dbt project location", - ) - group = import_parser.add_mutually_exclusive_group(required=True) group.add_argument("--sql", help="SQL query to snapshot") group.add_argument("--table", help="Table to snapshot") - group.add_argument( - "--model", - help=argparse.SUPPRESS, - # help="Model name to download", - ) - import_parser.add_argument( - "--profile", - help=argparse.SUPPRESS, - # help="Profile to use", - ) - import_parser.add_argument( - "--target", - help=argparse.SUPPRESS, - # help="Target to use", - ) import_parser.add_argument( "--output-format", default="csv", @@ -177,16 +82,9 @@ def add_parser(subparsers, parent_parser): ) import_parser.add_argument( "--conn", - nargs="?", + required=True, help="Database connection to use, needs to be set in config", ) - import_parser.add_argument( - "--dbt-conn", - action="store_true", - default=False, - help=argparse.SUPPRESS, - # help="Use dbt connection", - ) import_parser.set_defaults(func=CmdImportDb) @@ -197,11 +95,8 @@ def add_parser(subparsers, parent_parser): description=append_doc_link(TEST_DB_HELP, "test-db"), add_help=False, ) - test_db_parser.add_argument("--conn") - test_db_parser.add_argument("--dbt-conn", action="store_true", default=False) + test_db_parser.add_argument("--conn", required=True) test_db_parser.add_argument("--url") test_db_parser.add_argument("--password") test_db_parser.add_argument("--username") - test_db_parser.add_argument("--profile") - test_db_parser.add_argument("--target") test_db_parser.set_defaults(func=CmdTestDb) diff --git a/dvc/config_schema.py b/dvc/config_schema.py index 38740b117..7d56192e4 100644 --- a/dvc/config_schema.py +++ b/dvc/config_schema.py @@ -321,8 +321,6 @@ SCHEMA = { "feature": FeatureSchema( { Optional("machine", default=False): Bool, - "dbt_profile": str, - "dbt_target": str, }, ), "plots": { diff --git a/dvc/database.py b/dvc/database.py new file mode 100644 index 000000000..ebbb35341 --- /dev/null +++ b/dvc/database.py @@ -0,0 +1,124 @@ +import os +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union + +from sqlalchemy import create_engine +from sqlalchemy.engine import make_url as _make_url +from sqlalchemy.exc import NoSuchModuleError + +from dvc import env +from dvc.exceptions import DvcException +from dvc.log import logger +from dvc.types import StrOrBytesPath +from dvc.utils import env2bool + +if TYPE_CHECKING: + from sqlalchemy.engine import URL, Connectable, Engine + from sqlalchemy.sql.expression import Selectable + + +logger = logger.getChild(__name__) + + +def noop(_): + pass + + +def make_url(url: Union["URL", str], **kwargs: Any) -> "URL": + return _make_url(url).set(**kwargs) + + +def url_from_config(config: Union[str, "URL", Dict[str, str]]) -> "URL": + if isinstance(config, dict): + return make_url(**config) + return make_url(config) + + +@dataclass +class Serializer: + sql: "Union[str, Selectable]" + con: "Union[str, Connectable]" + chunksize: int = 10_000 + + def to_csv(self, file: StrOrBytesPath, progress=noop): + import pandas as pd + + with open(file, mode="wb") as f: + idfs = pd.read_sql(self.sql, self.con, chunksize=self.chunksize) + for i, df in enumerate(idfs): + df.to_csv(f, header=i == 0, index=False) + progress(len(df)) + + def to_json(self, file: StrOrBytesPath, progress=noop): # noqa: ARG002 + import pandas as pd + + path = os.fsdecode(file) + df = pd.read_sql(self.sql, self.con) + df.to_json(path, orient="records") + + def export(self, file: StrOrBytesPath, format: str = "csv", progress=noop): # noqa: A002 + if format == "json": + return self.to_json(file, progress=progress) + return self.to_csv(file, progress=progress) + + +@dataclass +class Client: + engine: "Engine" + + def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None: + try: + with self.engine.connect() as conn: + conn.exec_driver_sql("select 1") + except Exception as exc: + if callable(onerror): + onerror() + logger.exception( + "Could not connect to the database. " + "Check your database credentials and try again.", + exc_info=False, + ) + raise DvcException("The database returned the following error") from exc + + def export( + self, + sql: "Union[str, Selectable]", + file: StrOrBytesPath, + format: str = "csv", # noqa: A002 + progress=noop, + ) -> None: + con = self.engine.connect() + if format == "csv": + con = con.execution_options(stream_results=True) # use server-side cursors + + with con: + serializer = Serializer(sql, con) + return serializer.export(file, format=format, progress=progress) + + +@contextmanager +def handle_error(url: "URL"): + try: + yield + except (ModuleNotFoundError, NoSuchModuleError) as e: + # TODO: write installation instructions + driver = url.drivername + raise DvcException(f"Could not load database driver for {driver!r}") from e + + +@contextmanager +def client( + url_or_config: Union[str, "URL", Dict[str, str]], **engine_kwargs: Any +) -> Iterator[Client]: + url = url_from_config(url_or_config) + echo = env2bool(env.DVC_SQLALCHEMY_ECHO, False) + engine_kwargs.setdefault("echo", echo) + + with handle_error(url): + engine = create_engine(url, **engine_kwargs) + + try: + yield Client(engine) + finally: + engine.dispose() diff --git a/dvc/database/__init__.py b/dvc/database/__init__.py deleted file mode 100644 index 19ca5749f..000000000 --- a/dvc/database/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -from typing import Any, ContextManager, Optional, Union - -from . import dbt_query, sqla -from .dbt_models import get_model -from .serializer import export - -Client = Union[sqla.SQLAlchemyClient, dbt_query.DbtClient] - - -def get_client( - config, - project_dir: Optional[str] = None, - profile: Optional[str] = None, - target: Optional[str] = None, - **kwargs: Any, -) -> "ContextManager[Client]": - if config: - return sqla.client(config, **kwargs) - return dbt_query.client(project_dir=project_dir, profile=profile, target=target) - - -__all__ = ["export", "get_client", "get_model", "Client"] diff --git a/dvc/database/dbt_models.py b/dvc/database/dbt_models.py deleted file mode 100644 index 347c44009..000000000 --- a/dvc/database/dbt_models.py +++ /dev/null @@ -1,40 +0,0 @@ -from typing import List, Optional - -from dvc.log import logger - -from .dbt_runner import dbt_show -from .dbt_utils import check_dbt -from .serializer import AgateSerializer - -logger = logger.getChild(__name__) - - -def ref( - name: str, - package: Optional[str] = None, - version: Optional[int] = None, -) -> str: - parts: List[str] = [] - if package: - parts.append(repr(package)) - - parts.append(repr(name)) - if version: - parts.append(f"{version=}") - - inner = ",".join(parts) - return "{{ ref(" + inner + ") }}" - - -@check_dbt("model") -def get_model( - name: str, - package: Optional[str] = None, - version: Optional[int] = None, - profile: Optional[str] = None, - target: Optional[str] = None, -) -> AgateSerializer: - model = ref(name, package, version=version) - q = f"select * from {model}" # noqa: S608 - table = dbt_show(inline=q, profile=profile, target=target) - return AgateSerializer(table) diff --git a/dvc/database/dbt_query.py b/dvc/database/dbt_query.py deleted file mode 100644 index a7f3d6479..000000000 --- a/dvc/database/dbt_query.py +++ /dev/null @@ -1,102 +0,0 @@ -import os -from contextlib import contextmanager -from typing import ( - TYPE_CHECKING, - Any, - Callable, - ContextManager, - Dict, - Iterator, - Optional, - Union, -) - -from attrs import define, field - -from dvc.exceptions import DvcException -from dvc.log import logger - -from .dbt_utils import check_dbt -from .serializer import AgateSerializer, PandasSQLSerializer - -if TYPE_CHECKING: - from dbt.adapters.base.impl import BaseAdapter - -logger = logger.getChild(__name__) - - -@define -class DbtClient: - adapter: "BaseAdapter" = field(repr=lambda o: type(o).__qualname__) - creds: Dict[str, Any] = field(repr=False) - - @contextmanager - def query(self, sql: str) -> Iterator[Union[AgateSerializer, PandasSQLSerializer]]: - from dbt.adapters.sql import SQLAdapter - - with self.adapter.connection_named("execute"): - if isinstance(self.adapter, SQLAdapter): - # try to use connection to stream results - # otherwise fallback to fetch all at once - try: - connection = self.adapter.connections.get_thread_connection().handle - except Exception: # noqa: BLE001 - logger.debug("", exc_info=True) - else: - # pandas raises warning that it does not support - # dbapi2-based raw connection, but it does seem to work. - yield PandasSQLSerializer(sql, connection) - return - _, table = self.adapter.execute(sql, fetch=True) - yield AgateSerializer(table) - - def table( - self, name: str, limit: Optional[int] = None - ) -> ContextManager[Union[AgateSerializer, PandasSQLSerializer]]: - raise NotImplementedError - - def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None: - with self.adapter.connection_named("debug"): - try: - self.adapter.debug_query() - except Exception as exc: - if callable(onerror): - onerror() - - logger.exception( - "Could not connect to the database. " - "Check your database credentials and try again.", - exc_info=False, - ) - raise DvcException("The database returned the following error") from exc - - -@contextmanager -@check_dbt("query") -def client( - project_dir: Optional[str] = None, - profiles_dir: Optional[str] = None, - profile: Optional[str] = None, - target: Optional[str] = None, -) -> Iterator["DbtClient"]: - from dbt.adapters import factory as adapters_factory - from dbt.adapters.base.impl import BaseAdapter - - from .dbt_utils import get_or_build_profile, get_profiles_dir, init_dbt - - profiles_dir = profiles_dir or get_profiles_dir(project_dir) - flags = init_dbt(profiles_dir, os.getcwd(), target=target) - - with flags, adapters_factory.adapter_management(): - # likely invalid connection profile or no adapter - profile_obj = get_or_build_profile(project_dir, profile, target) - - adapters_factory.register_adapter(profile_obj) # type: ignore[arg-type] - adapter = adapters_factory.get_adapter(profile_obj) # type: ignore[arg-type] - assert isinstance(adapter, BaseAdapter) - try: - creds = dict(profile_obj.credentials.connection_info()) - except: # noqa: E722 - creds = {} - - yield DbtClient(adapter, creds) diff --git a/dvc/database/dbt_runner.py b/dvc/database/dbt_runner.py deleted file mode 100644 index 3413d9dc7..000000000 --- a/dvc/database/dbt_runner.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import TYPE_CHECKING, List, Optional, Sequence, Union - -from dvc.exceptions import DvcException -from dvc.log import logger - -if TYPE_CHECKING: - from agate import Table - from dbt.contracts.results import RunResult - -logger = logger.getChild(__name__) - - -class DbtInternalError(DvcException): - pass - - -def _kw_to_cmd_args(**kwargs: Union[None, bool, str, int]) -> List[str]: - args: List[str] = [] - for key, value in kwargs.items(): - key = key.replace("_", "-") - if value is None: - continue # skip creating a flag in this case - if value is True: - args.append(f"--{key}") - elif value is False: - args.append(f"--no-{key}") - else: - args.extend([f"--{key}", str(value)]) - return args - - -def _dbt_invoke(*posargs: str, quiet: bool = True, **kw: Union[None, bool, str, int]): - from dbt.cli.main import dbtRunner - - args = _kw_to_cmd_args(quiet=quiet or None) # global options - args.extend([*posargs, *_kw_to_cmd_args(**kw)]) - - runner = dbtRunner() - result = runner.invoke(args) - if result.success: - return result.result - raise DbtInternalError(f"failed to run dbt {posargs[0]}") from result.exception - - -def dbt_show( - inline: Optional[str] = None, - limit: int = -1, - profile: Optional[str] = None, - target: Optional[str] = None, -) -> "Table": - from dbt.contracts.results import RunExecutionResult - - result = _dbt_invoke( - "show", inline=inline, limit=limit, profile=profile, target=target - ) - assert isinstance(result, RunExecutionResult) - - run_results: Sequence["RunResult"] = result.results - run_result, *_ = run_results - assert run_result.agate_table is not None - return run_result.agate_table diff --git a/dvc/database/dbt_utils.py b/dvc/database/dbt_utils.py deleted file mode 100644 index 87111ddb1..000000000 --- a/dvc/database/dbt_utils.py +++ /dev/null @@ -1,100 +0,0 @@ -import os -from contextlib import contextmanager -from importlib.util import find_spec -from typing import TYPE_CHECKING, Optional - -from funcy import cut_prefix - -from dvc.exceptions import DvcException -from dvc.log import logger -from dvc.utils import packaging - -if TYPE_CHECKING: - from dbt.config.profile import Profile - -logger = logger.getChild(__name__) - -DBT_PROJECT_FILE = "dbt_project.yml" - - -@contextmanager -def check_dbt(action: Optional[str] = None): - if not (find_spec("dbt") and find_spec("dbt.cli")): - action = f" {action}" if action else "" - raise DvcException(f"Could not run{action}. dbt-core is not installed") - yield packaging.check_required_version(pkg="dbt-core") - - -@contextmanager -def init_dbt( - profiles_dir: str, - project_dir: str, - target: Optional[str] = None, -): - from argparse import Namespace - - from dbt import flags - - prev = flags.get_flags() - try: - args = Namespace( - use_colors=True, - project_dir=project_dir, - profiles_dir=profiles_dir, - target=target, - ) - flags.set_from_args(args, None) - yield flags.get_flags() - finally: - flags.set_flags(prev) - - -def get_profiles_dir(project_dir: Optional[str] = None) -> str: - from dbt.cli.resolvers import default_profiles_dir - - if profiles_dir := os.getenv("DBT_PROFILES_DIR"): - return profiles_dir - if project_dir and os.path.isfile(os.path.join(project_dir, "profiles.yml")): - return project_dir - return os.fspath(default_profiles_dir()) - - -@contextmanager -def _handle_profile_parsing_error(): - from dbt.exceptions import DbtRuntimeError - - try: - yield - except DbtRuntimeError as e: - cause = e.__cause__ is not None and e.__cause__.__context__ - if isinstance(cause, ModuleNotFoundError) and ( - adapter := cut_prefix(cause.name, "dbt.adapters.") - ): - # DbtRuntimeError is very noisy, so send it to debug - logger.debug("", exc_info=True) - raise DvcException(f"dbt-{adapter} dependency is missing") from cause - raise DvcException("failed to read connection profiles") from e - - -def is_dbt_project(project_dir: str): - return os.path.isfile(os.path.join(project_dir, DBT_PROJECT_FILE)) - - -@_handle_profile_parsing_error() -def get_or_build_profile( - project_dir: Optional[str], profile: Optional[str], target: Optional[str] -) -> "Profile": - from dbt.config.profile import Profile - from dbt.config.renderer import ProfileRenderer - from dbt.config.runtime import load_profile - - if project_dir and is_dbt_project(project_dir): - return load_profile( - project_dir, {}, profile_name_override=profile, target_override=target - ) - - if not profile: - raise DvcException("No profile specified to query from.") - - renderer = ProfileRenderer({}) - return Profile.render(renderer, profile, target_override=target) diff --git a/dvc/database/serializer.py b/dvc/database/serializer.py deleted file mode 100644 index 3858489a4..000000000 --- a/dvc/database/serializer.py +++ /dev/null @@ -1,54 +0,0 @@ -from dataclasses import dataclass -from typing import TYPE_CHECKING, Union - -if TYPE_CHECKING: - import sqlalchemy as sa - from agate import Table - - -def noop(_): - pass - - -@dataclass -class PandasSQLSerializer: - sql: "Union[str, sa.sql.Selectable]" - con: "sa.engine.Connection" - chunksize: int = 10_000 - - def to_csv(self, file: str, progress=noop) -> None: - import pandas as pd - - with open(file, mode="wb") as f: - idfs = pd.read_sql_query(self.sql, self.con, chunksize=self.chunksize) - for i, df in enumerate(idfs): - df.to_csv(f, header=i == 0, index=False) - progress(len(df)) - - def to_json(self, file: str, progress=noop) -> None: # noqa: ARG002 - import pandas as pd - - df = pd.read_sql_query(self.sql, self.con) - df.to_json(file, orient="records") - - -@dataclass -class AgateSerializer: - table: "Table" - - def to_csv(self, file: str, progress=noop) -> None: # noqa: ARG002 - return self.table.to_csv(file) - - def to_json(self, file: str, progress=noop) -> None: # noqa: ARG002 - return self.table.to_json(file) - - -def export( - serializer: Union[PandasSQLSerializer, AgateSerializer], - file: str, - format: str = "csv", # noqa: A002 - progress=noop, -) -> None: - if format == "csv": - return serializer.to_csv(file, progress=progress) - return serializer.to_json(file, progress=progress) diff --git a/dvc/database/sqla.py b/dvc/database/sqla.py deleted file mode 100644 index 6487deb8f..000000000 --- a/dvc/database/sqla.py +++ /dev/null @@ -1,85 +0,0 @@ -from contextlib import contextmanager -from dataclasses import dataclass -from typing import Any, Callable, ContextManager, Dict, Iterator, Optional, Union - -from sqlalchemy import create_engine, select, text -from sqlalchemy.engine import URL, Engine -from sqlalchemy.engine import make_url as _make_url -from sqlalchemy.exc import NoSuchModuleError -from sqlalchemy.sql import table -from sqlalchemy.sql.expression import Selectable - -from dvc import env -from dvc.exceptions import DvcException -from dvc.log import logger -from dvc.utils import env2bool - -from .serializer import PandasSQLSerializer - -logger = logger.getChild(__name__) - - -def make_url(url: Union[URL, str], **kwargs: Any) -> URL: - return _make_url(url).set(**kwargs) - - -def url_from_config(config: Union[str, URL, Dict[str, str]]) -> URL: - if isinstance(config, (str, URL)): - return make_url(config) - return make_url(**config) - - -@dataclass -class SQLAlchemyClient: - engine: Engine - - @contextmanager - def query(self, sql: Union[str, Selectable]) -> Iterator[PandasSQLSerializer]: - with self.engine.connect().execution_options(stream_results=True) as conn: - yield PandasSQLSerializer(sql, conn) - - def table( - self, name: str, limit: Optional[int] = None - ) -> ContextManager[PandasSQLSerializer]: - stmt = select("*").select_from(table(name)).limit(limit) - return self.query(stmt) - - def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None: - try: - with self.engine.connect() as conn: - conn.execute(text("select 1")) - except Exception as exc: - if callable(onerror): - onerror() - logger.exception( - "Could not connect to the database. " - "Check your database credentials and try again.", - exc_info=False, - ) - raise DvcException("The database returned the following error") from exc - - -@contextmanager -def handle_error(url: URL): - try: - yield - except (ModuleNotFoundError, NoSuchModuleError) as e: - # TODO: write installation instructions - raise DvcException( - f"Could not load database driver for {url.drivername!r}" - ) from e - - -@contextmanager -def client( - url_or_config: Union[str, URL, Dict[str, str]], **engine_kwargs: Any -) -> Iterator[SQLAlchemyClient]: - url = url_from_config(url_or_config) - with handle_error(url): - echo = env2bool(env.DVC_SQLALCHEMY_ECHO, False) - engine_kwargs.setdefault("echo", echo) - engine = create_engine(url, **engine_kwargs) - try: - yield SQLAlchemyClient(engine) - finally: - engine.dispose() diff --git a/dvc/dependency/__init__.py b/dvc/dependency/__init__.py index b555d1079..e3def3c9b 100644 --- a/dvc/dependency/__init__.py +++ b/dvc/dependency/__init__.py @@ -4,7 +4,7 @@ from typing import Any, Dict, List, Mapping, Set from dvc.output import ARTIFACT_SCHEMA, DIR_FILES_SCHEMA, Output from .base import Dependency -from .db import DB_SCHEMA, PARAM_DB, DbDependency, DbtDependency +from .db import DbDependency from .param import ParamsDependency from .repo import RepoDependency @@ -15,7 +15,7 @@ from .repo import RepoDependency SCHEMA: Mapping[str, Any] = { **ARTIFACT_SCHEMA, **RepoDependency.REPO_SCHEMA, - **DB_SCHEMA, + **DbDependency.DB_SCHEMA, Output.PARAM_FILES: [DIR_FILES_SCHEMA], Output.PARAM_FS_CONFIG: dict, } @@ -23,16 +23,13 @@ SCHEMA: Mapping[str, Any] = { def _get(stage, p, info, **kwargs): d = info or {} - db = d.get(PARAM_DB, {}) params = d.pop(ParamsDependency.PARAM_PARAMS, None) repo = d.pop(RepoDependency.PARAM_REPO, None) if params: return ParamsDependency(stage, p, params) - if DbtDependency.PARAM_MODEL in db: - return DbtDependency(repo, stage, info) - if db: - return DbDependency(stage, info) + if DbDependency.PARAM_DB in d: + return DbDependency(stage, d) assert p if repo: diff --git a/dvc/dependency/db.py b/dvc/dependency/db.py index dfeee21ad..10a5dbf29 100644 --- a/dvc/dependency/db.py +++ b/dvc/dependency/db.py @@ -1,58 +1,20 @@ -import os -from contextlib import contextmanager, nullcontext -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional from funcy import compact, log_durations from dvc.exceptions import DvcException from dvc.log import logger -from dvc.scm import SCM from .base import Dependency if TYPE_CHECKING: - from rich.status import Status - from dvc.output import Output from dvc.stage import Stage logger = logger.getChild(__name__) -PARAM_DB = "db" -PARAM_PROFILE = "profile" -PARAM_FILE_FORMAT = "file_format" - - -def _get_dbt_config(config: Dict) -> Dict: - conf = config.get("feature", {}) - pref = "dbt_" - return {k.lstrip(pref): v for k, v in conf.items() if k.startswith(pref)} - - -@contextmanager -def log_status( - msg, status: Optional["Status"] = None, log=logger.debug -) -> Iterator["Status"]: - from funcy import log_durations - - from dvc.ui import ui - - with log_durations(log, msg), status or ui.status(msg) as st: - st.update(msg) - yield st - - -@contextmanager -def chdir(path): - wdir = os.getcwd() - os.chdir(path) - try: - yield - finally: - os.chdir(wdir) - - @contextmanager def download_progress(to: "Output") -> Iterator[Callable[[int], Any]]: from dvc.ui import ui @@ -80,14 +42,27 @@ class AbstractDependency(Dependency): class DbDependency(AbstractDependency): PARAM_CONNECTION = "connection" - PARAM_DBT_CONNECTION = "dbt_connection" + PARAM_DB = "db" PARAM_QUERY = "query" PARAM_TABLE = "table" - QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str, PARAM_TABLE: str} + PARAM_FILE_FORMAT = "file_format" + DB_SCHEMA = { + PARAM_DB: { + PARAM_QUERY: str, + PARAM_CONNECTION: str, + PARAM_FILE_FORMAT: str, + PARAM_TABLE: str, + } + } def __init__(self, stage: "Stage", info, *args, **kwargs): super().__init__(stage, info, *args, **kwargs) - self.target = None + self.db_info: Dict[str, str] = self.info.get(self.PARAM_DB, {}) + self.connection = self.db_info.get(self.PARAM_CONNECTION) + + @property + def sql(self) -> Optional[str]: + return self.db_info.get(self.PARAM_QUERY) or self.db_info.get(self.PARAM_TABLE) def __repr__(self): return "{}:{}".format( @@ -97,10 +72,7 @@ class DbDependency(AbstractDependency): def __str__(self): from dvc.utils.humanize import truncate_text - db_info = self.info.get(PARAM_DB, {}) - query = db_info.get(self.PARAM_QUERY) - table = db_info.get(self.PARAM_TABLE) - return truncate_text(table or query or "[no table or query]", 50) + return truncate_text(self.sql or "", 50) def workspace_status(self): return False # no workspace to check @@ -112,8 +84,8 @@ class DbDependency(AbstractDependency): """nothing to save.""" def dumpd(self, **kwargs): - db_info = compact(self.info.get(PARAM_DB, {})) - return {PARAM_DB: db_info} if db_info else {} + db_info = compact(self.db_info) + return {self.PARAM_DB: db_info} if db_info else {} def update(self, rev=None): """nothing to update.""" @@ -125,209 +97,22 @@ class DbDependency(AbstractDependency): file_format: Optional[str] = None, **kwargs: Any, ) -> None: - from dvc.database import export, get_client + from dvc.database import client + from dvc.ui import ui - db_info = self.info.get(PARAM_DB, {}) - query = db_info.get(self.PARAM_QUERY) - table = db_info.get(self.PARAM_TABLE) - if not (query or table): + sql = self.sql + if not sql: raise DvcException("Cannot download: no query or table specified") - if query and table: - raise DvcException( - "Cannot download: query and table are mutually exclusive" - ) - dbt_config = _get_dbt_config(self.repo.config) - profile = db_info.get(PARAM_PROFILE) or dbt_config.get(PARAM_PROFILE) - target = self.target or dbt_config.get("target") + db_config = self.repo.config.get(self.PARAM_DB, {}) + config = db_config.get(self.connection) + if not config: + raise DvcException(f"connection {self.connection} not found in config") - connection = db_info.get(self.PARAM_CONNECTION) - - db_config = self.repo.config.get("db", {}) - config = db_config.get(connection) - if connection and not config: - raise DvcException(f"connection {connection} not found in config") - - project_dir = self.repo.root_dir - file = to.fs_path - file_format = file_format or db_info.get(PARAM_FILE_FORMAT) or "csv" - assert file_format - with get_client( - config, project_dir=project_dir, profile=profile, target=target - ) as db: - logger.debug("using adapter: %s", db) - with log_status("Testing connection") as status: + file_format = file_format or self.db_info.get(self.PARAM_FILE_FORMAT, "csv") + with client(config) as db: + msg = "Testing connection" + with log_durations(logger.debug, msg), ui.status(msg) as status: db.test_connection(onerror=status.stop) - - query = db.query(query) if query else db.table(table) - with log_status("Executing query") as status, query as serializer: - logger.debug("using serializer: %s", serializer) - status.stop() - with download_progress(to) as prog: - return export(serializer, file, format=file_format, progress=prog) - - -class DbtDependency(AbstractDependency): - PARAM_MODEL = "model" - PARAM_VERSION = "version" - PARAM_PROJECT_DIR = "project_dir" - DBT_SCHEMA = { - PARAM_MODEL: str, - PARAM_VERSION: str, - PARAM_PROJECT_DIR: str, - } - - def __init__(self, def_repo: Dict[str, Any], stage: "Stage", info, *args, **kwargs): - self.def_repo = def_repo or {} - self.target = None - super().__init__(stage, info, *args, **kwargs) - - def __repr__(self): - return "{}:{}".format( - self.__class__.__name__, - "".join(f"{k}=={v}" for k, v in {**self.def_repo, **self.info}.items()), - ) - - def __str__(self): - from .repo import RepoDependency - - repo = self.def_repo.get(RepoDependency.PARAM_URL) - rev = self.def_repo.get(RepoDependency.PARAM_REV) - - db_info = self.info.get(PARAM_DB, {}) - db = db_info.get(self.PARAM_MODEL, "") - project_dir = db_info.get(self.PARAM_PROJECT_DIR, "") - repo_info = "" - if repo: - repo_info += repo - if rev: - repo_info += f"@{rev}" - if project_dir: - repo_info += f":/{project_dir}" - return db + (f"({repo_info})" if repo_info else "") - - @property - def locked_rev(self): - from .repo import RepoDependency - - return self.def_repo.get(RepoDependency.PARAM_REV_LOCK) - - @property - def rev(self): - from .repo import RepoDependency - - return self.def_repo.get(RepoDependency.PARAM_REV) - - def workspace_status(self): - if not self.def_repo: - return - - current = self._get_clone(self.locked_rev or self.rev).get_rev() - updated = self._get_clone(self.rev).get_rev() - if current != updated: - return {str(self): "update available"} - return {} - - def status(self): - return self.workspace_status() - - def save(self): - from .repo import RepoDependency - - if not self.def_repo: - return - - rev = self._get_clone(self.locked_rev or self.rev).get_rev() - if self.def_repo.get(RepoDependency.PARAM_REV_LOCK) is None: - self.def_repo[RepoDependency.PARAM_REV_LOCK] = rev - - def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]: - from .repo import RepoDependency - - def_repo = {} - if self.def_repo: - def_repo = RepoDependency._dump_def_repo(self.def_repo) - - db_info = compact(self.info.get(PARAM_DB, {})) - return compact({RepoDependency.PARAM_REPO: def_repo, PARAM_DB: db_info}) - - def _get_clone(self, rev): - from dvc.repo.open_repo import _cached_clone - - from .repo import RepoDependency - - url = self.def_repo.get(RepoDependency.PARAM_URL) - repo_root = self.repo.root_dir if self.repo else os.getcwd() - return SCM(_cached_clone(url, rev) if url else repo_root) - - def update(self, rev=None): - from .repo import RepoDependency - - if not self.def_repo: - return - - if rev: - self.def_repo[RepoDependency.PARAM_REV] = rev - else: - rev = self.rev - self.def_repo[RepoDependency.PARAM_REV_LOCK] = self._get_clone(rev).get_rev() - - def download( - self, - to: "Output", - jobs: Optional[int] = None, # noqa: ARG002 - file_format: Optional[str] = None, - ) -> None: - from .repo import RepoDependency - - project_dir = self.info.get(PARAM_DB, {}).get(self.PARAM_PROJECT_DIR, "") - if self.def_repo: - ctx = repo = self._get_clone(self.locked_rev or self.rev) - self.def_repo[RepoDependency.PARAM_REV_LOCK] = repo.get_rev() - root = wdir = repo.root_dir - else: - ctx = nullcontext() - root = self.repo.root_dir - wdir = self.stage.wdir - - project_path = os.path.join(wdir, project_dir) if project_dir else root - with ctx, chdir(project_path): - self._download_db(to, file_format=file_format) - - def _download_db( - self, - to: "Output", - version: Optional[int] = None, - file_format: Optional[str] = None, - ) -> None: - from dvc.database import export, get_model - - db_info = self.info.get(PARAM_DB, {}) - model = db_info.get(self.PARAM_MODEL) - version = version or db_info.get(self.PARAM_VERSION) - file_format = file_format or db_info.get(PARAM_FILE_FORMAT) or "csv" - assert file_format - if not model: - raise DvcException("Cannot download, no model specified") - - dbt_config = _get_dbt_config(self.repo.config) - profile = db_info.get(PARAM_PROFILE) or dbt_config.get(PARAM_PROFILE) - target = self.target or db_info.get("target") or dbt_config.get("target") - - with log_status("Downloading model"): - serializer = get_model( - model, version=version, profile=profile, target=target - ) - # NOTE: we keep everything in memory, and then export it out later. - with download_progress(to) as progress: - export(serializer, to.fs_path, format=file_format, progress=progress) - - -DB_SCHEMA = { - PARAM_DB: { - PARAM_PROFILE: str, - PARAM_FILE_FORMAT: str, - **DbDependency.QUERY_SCHEMA, - **DbtDependency.DBT_SCHEMA, - }, -} + with download_progress(to) as progress: + db.export(sql, to.fs_path, format=file_format, progress=progress) diff --git a/dvc/repo/imp_db.py b/dvc/repo/imp_db.py index 9b71f996d..0366dd941 100644 --- a/dvc/repo/imp_db.py +++ b/dvc/repo/imp_db.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Dict, Optional + +from funcy import compact from dvc.exceptions import OutputDuplicationError from dvc.repo.scm_context import scm_context @@ -13,18 +15,11 @@ from . import locked @locked @scm_context -def imp_db( # noqa: PLR0913 +def imp_db( self: "Repo", - url: Optional[str] = None, - rev: Optional[str] = None, - project_dir: Optional[str] = None, sql: Optional[str] = None, table: Optional[str] = None, - model: Optional[str] = None, - version: Optional[int] = None, frozen: bool = True, - profile: Optional[str] = None, - target: Optional[str] = None, output_format: str = "csv", out: Optional[str] = None, force: bool = False, @@ -35,41 +30,33 @@ def imp_db( # noqa: PLR0913 "The functionality may change or break without notice, " "which could lead to unexpected behavior." ) - assert model or sql or table - erepo = None - if model and url: - erepo = {"url": url} - if rev: - erepo["rev"] = rev - + assert sql or table assert output_format in ("csv", "json") - db: Dict[str, Any] = {"file_format": output_format} - if profile: - db["profile"] = profile + db: Dict[str, str] = compact( + { + "connection": connection, + "file_format": output_format, + "query": sql, + "table": table, + } + ) - if model: - db.update({"model": model, "version": version, "project_dir": project_dir}) - else: - db.update({"connection": connection, "query": sql, "table": table}) - - file_name = model or table or "results" + file_name = table or "results" out = out or f"{file_name}.{output_format}" - out = resolve_output(url or ".", out, force=force) + out = resolve_output(".", out, force=force) + path, wdir, out = resolve_paths(self, out, always_local=True) stage = self.stage.create( single_stage=True, validate=False, fname=path, + deps=[None], wdir=wdir, - deps=[url], outs=[out], - erepo=erepo, - fs_config=None, db=db, ) - stage.deps[0].target = target try: self.check_graph(stages={stage}) except OutputDuplicationError as exc: diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 6047df30c..3e75af283 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -286,17 +286,17 @@ class Stage(params.StageParams): if not self.is_import: return False - from dvc.dependency import DbDependency, DbtDependency + from dvc.dependency import DbDependency - return isinstance(self.deps[0], (DbDependency, DbtDependency)) + return isinstance(self.deps[0], DbDependency) @property def is_versioned_import(self) -> bool: - from dvc.dependency import DbDependency, DbtDependency + from dvc.dependency import DbDependency return ( self.is_import - and not isinstance(self.deps[0], (DbDependency, DbtDependency)) + and not isinstance(self.deps[0], DbDependency) and self.deps[0].fs.version_aware ) @@ -462,7 +462,7 @@ class Stage(params.StageParams): if not (self.is_repo_import or self.is_import): raise StageUpdateError(self.relpath) - # always force update DbDep/DbtDep since we don't know if it's changed + # always force update DbDep since we don't know if it's changed force = self.is_db_import update_import( self, diff --git a/pyproject.toml b/pyproject.toml index 1eb843141..86c1a10de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -112,9 +112,8 @@ tests = [ "pytest-xdist>=3.2", 'pywin32>=225; sys_platform == "win32"', # optional test dependency 'tzdata; sys_platform == "win32"', # for testing with celery - "dbt-sqlite", # optional dependency for testing `import-db` "sqlalchemy>=1,<3", # optional dependency for `import-db` - "pandas>=1", + "pandas>=1", # optional dependency for `import-db` ] webdav = ["dvc-webdav>=3,<4"] webhdfs = ["dvc-webhdfs>=3.1,<4"] diff --git a/tests/func/test_import_db.py b/tests/func/test_import_db.py index 039bc89d5..9f5dd3c41 100644 --- a/tests/func/test_import_db.py +++ b/tests/func/test_import_db.py @@ -1,165 +1,116 @@ -"""import_db tests without involving any databases.""" +import os +import sqlite3 +from contextlib import closing + +import pandas as pd import pytest -from agate import Table - -from dvc.database.serializer import AgateSerializer - -data1 = Table([(i, i) for i in range(1, 10)], ["id", "value"]) -data2 = Table([(i, i) for i in range(10, 20)], ["id", "value"]) - -serializers = AgateSerializer(data1), AgateSerializer(data2) +from funcy import compact @pytest.fixture -def client(mocker): - m = mocker.patch("dvc.database.get_client") - client = mocker.MagicMock() - client.query.return_value.__enter__.side_effect = serializers - client.table.return_value.__enter__.side_effect = serializers - m.return_value.__enter__.return_value = client - return m +def db_path(tmp_dir): + return tmp_dir / "main.db" -def test_sql_with_dbt(client, tmp_dir, dvc): - stage = dvc.imp_db( - sql="select * from model", profile="profile", output_format="json" - ) +@pytest.fixture +def seed_db(db_path): + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE model (id INTEGER PRIMARY KEY, value INTEGER)") - db = {"file_format": "json", "profile": "profile", "query": "select * from model"} + def inner(values): + conn.executemany("INSERT INTO model(value) VALUES(?)", [(i,) for i in values]) + conn.commit() - assert (tmp_dir / stage.relpath).parse() == { - "md5": "dae475048115315c8c26a10538e4b2ec", - "frozen": True, - "deps": [{"db": db}], - "outs": [ - { - "hash": "md5", - "md5": "56226d443067b195c90b427db557e1f2", - "path": "results.json", - "size": 243, - } - ], - } - - dvc.update(stage.addressing) - assert (tmp_dir / stage.relpath).parse() == { - "md5": "79b42f9ddb2a3ed999eb494ccab470c6", - "frozen": True, - "deps": [{"db": db}], - "outs": [ - { - "hash": "md5", - "md5": "08d806ee7e1f309c4561750805f02276", - "path": "results.json", - "size": 290, - } - ], - } + with closing(conn): + yield inner -def test_sql_with_conn_string(client, tmp_dir, dvc): +@pytest.fixture +def db_connection(dvc, db_path): with dvc.config.edit(level="local") as conf: - conf["db"] = {"conn": {"url": "conn"}} + conf["db"] = {"conn": {"url": f"sqlite:///{db_path.fs_path}"}} + return "conn" - stage = dvc.imp_db( - sql="select * from model", connection="conn", output_format="json" + +def load_data(file, output_format): + if output_format == "json": + return pd.read_json(file, orient="records") + return pd.read_csv(file) + + +@pytest.mark.parametrize("output_format", ("csv", "json")) +@pytest.mark.parametrize( + "args,file_name", + [ + ({"sql": "select * from model"}, "results"), + ({"table": "model"}, "model"), + ], +) +def test( + M, + tmp_dir, + scm, + dvc, + db_connection, + seed_db, + output_format, + args, + file_name, +): + seed_db(values=range(5)) + if output_format == "json": + file_size = 96, 192 + md5 = "6039fe7565d212b339aaa446ca234e5d", "e1b8adf4d9eb9ab2b64d3ab0bb5f65ac" + elif os.name == "nt": + file_size = 35, 61 + md5 = "14c34db5ddd184345c06f74718539f04", "3bb836e6d43c9afa43a9d73b36bbbab4" + else: + file_size = 29, 50 + md5 = "6f7fc0d701d1ac13eec83d79fffaf427", "c04f712f8167496a2fb43f289f2b7e28" + + db = compact( + { + "file_format": output_format, + "connection": db_connection, + "table": args.get("table"), + "query": args.get("sql"), + } ) - db = {"file_format": "json", "connection": "conn", "query": "select * from model"} + stage = dvc.imp_db(**args, connection=db_connection, output_format=output_format) + + output_file = f"{file_name}.{output_format}" + df = load_data(output_file, output_format) + assert df.values.tolist() == [[i + 1, i] for i in range(5)] assert (tmp_dir / stage.relpath).parse() == { - "md5": "f0dcca8e8411907adb25aa547d7c432c", + "md5": M.instance_of(str), "frozen": True, "deps": [{"db": db}], "outs": [ { + "md5": md5[0], + "size": file_size[0], "hash": "md5", - "md5": "56226d443067b195c90b427db557e1f2", - "path": "results.json", - "size": 243, + "path": output_file, } ], } + seed_db(values=range(5, 10)) + dvc.update(stage.addressing) + + df = load_data(output_file, output_format) + assert df.values.tolist() == [[i + 1, i] for i in range(10)] assert (tmp_dir / stage.relpath).parse() == { - "md5": "8c0441b206d8ee7c7b9e3d54871c3cbb", + "md5": M.instance_of(str), "frozen": True, "deps": [{"db": db}], "outs": [ { + "md5": md5[1], + "size": file_size[1], "hash": "md5", - "md5": "08d806ee7e1f309c4561750805f02276", - "path": "results.json", - "size": 290, - } - ], - } - - -def test_table_with_conn_string(client, tmp_dir, dvc): - with dvc.config.edit(level="local") as conf: - conf["db"] = {"conn": {"url": "conn"}} - - stage = dvc.imp_db(table="model", connection="conn", output_format="json") - db = {"file_format": "json", "connection": "conn", "table": "model"} - assert (tmp_dir / stage.relpath).parse() == { - "md5": "39dfcd61935e039a7bc3f158d7fcf62c", - "frozen": True, - "deps": [{"db": db}], - "outs": [ - { - "hash": "md5", - "md5": "56226d443067b195c90b427db557e1f2", - "path": "model.json", - "size": 243, - } - ], - } - - dvc.update(stage.addressing) - assert (tmp_dir / stage.relpath).parse() == { - "md5": "20d77a700710bad7d094c460f5939b7c", - "frozen": True, - "deps": [{"db": db}], - "outs": [ - { - "hash": "md5", - "md5": "08d806ee7e1f309c4561750805f02276", - "path": "model.json", - "size": 290, - } - ], - } - - -def test_model(mocker, tmp_dir, dvc): - mocker.patch("dvc.database.get_model", side_effect=serializers) - stage = dvc.imp_db(model="model", output_format="json") - - assert (tmp_dir / stage.relpath).parse() == { - "md5": "a8157fb5b28457e1910506e9d4d4fca3", - "frozen": True, - "deps": [{"db": {"file_format": "json", "model": "model"}}], - "outs": [ - { - "hash": "md5", - "md5": "56226d443067b195c90b427db557e1f2", - "path": "model.json", - "size": 243, - } - ], - } - - dvc.update(stage.addressing) - assert (tmp_dir / stage.relpath).parse() == { - "md5": "2f84287f109fc2b8755a9e64ad910b5d", - "frozen": True, - "deps": [{"db": {"file_format": "json", "model": "model"}}], - "outs": [ - { - "hash": "md5", - "md5": "08d806ee7e1f309c4561750805f02276", - "path": "model.json", - "size": 290, + "path": output_file, } ], } diff --git a/tests/func/test_import_db_e2e.py b/tests/func/test_import_db_e2e.py deleted file mode 100644 index 3abbe9ddc..000000000 --- a/tests/func/test_import_db_e2e.py +++ /dev/null @@ -1,147 +0,0 @@ -"""E2E tests involving dvc/dbt and a ~~real~~ sqlite db. - -Any other tests should likely be mocked tests in `test_import_db` rather than here. -""" - -import os -import sqlite3 -from contextlib import closing - -import pytest -from agate import Table - -from dvc.types import StrPath - -pytest.importorskip("dbt.adapters.sqlite", reason="dbt-sqlite not installed") - - -@pytest.fixture(autouse=True) -def disable_tracking(monkeypatch): - monkeypatch.setenv("DO_NOT_TRACK", "1") - - -@pytest.fixture(autouse=True) -def mock_sqlite_adapter(mocker): - # https://github.com/codeforkjeff/dbt-sqlite/issues/47 - mocker.patch( - "dbt.adapters.sqlite.connections.SQLiteCredentials.unique_field", "sqlite_host" - ) - - -@pytest.fixture -def db_path(tmp_dir): - d = tmp_dir / "db" - d.mkdir() - return d / "main.db" - - -@pytest.fixture -def seed_db(db_path): - conn = sqlite3.connect(db_path) - conn.execute("CREATE TABLE model (id INTEGER PRIMARY KEY, value INTEGER)") - - def inner(values): - conn.executemany("INSERT INTO model(value) VALUES(?)", [(i,) for i in values]) - conn.commit() - - with closing(conn): - yield inner - - -@pytest.fixture -def db_config(dvc, db_path): - # only needed for `sql_conn` tests - with dvc.config.edit(level="local") as conf: - conf["db"] = {"conn": {"url": f"sqlite:///{db_path.fs_path}"}} - return "conn" - - -@pytest.fixture -def dbt_profile(tmp_dir, scm, db_path): - (tmp_dir / "profiles.yml").dump( - { - "sqlite_profile": { - "target": "dev", - "outputs": { - "dev": { - "type": "sqlite", - "threads": 1, - "database": "database", - "schema": "main", - "schemas_and_paths": {"main": os.fspath(db_path)}, - "schema_directory": os.fspath(db_path.parent), - } - }, - } - } - ) - scm.add_commit(["profiles.yml"], message="add profiles.yml") - return "sqlite_profile" - - -@pytest.fixture -def dbt_project(tmp_dir, scm, dbt_profile): - (tmp_dir / "dbt_project.yml").dump({"name": "project", "profile": dbt_profile}) - scm.add_commit(["dbt_project.yml"], message="add dbt_project") - return tmp_dir - - -@pytest.fixture -def dbt_model(dbt_project): - dbt_project.scm_gen( - {"models": {"model.sql": "select * from models"}}, commit="add models" - ) - return "model" - - -@pytest.fixture(params=("sql", "model", "external_model", "sql_conn", "table")) -def import_db_parameters(request: pytest.FixtureRequest): - if request.param == "sql": - profile = request.getfixturevalue("dbt_profile") - return {"sql": "select * from model", "profile": profile} - if request.param == "sql_conn": - conn = request.getfixturevalue("db_config") - return {"sql": "select * from model", "connection": conn} - if request.param == "table": - conn = request.getfixturevalue("db_config") - return {"table": "model", "connection": conn} - - dbt_project = request.getfixturevalue("dbt_project") - return { - "model": request.getfixturevalue("dbt_model"), - "url": None if request.param == "model" else os.fspath(dbt_project), - } - - -@pytest.fixture -def file_name(import_db_parameters): - if model := import_db_parameters.get("model"): - return model - if table := import_db_parameters.get("table"): - return table - return "results" - - -def load_table(file: StrPath, typ: str) -> "Table": - return getattr(Table, f"from_{typ}")(file) - - -@pytest.mark.filterwarnings("ignore::ResourceWarning") # dbt leaks fileobj from logger -@pytest.mark.parametrize("output_format", ("csv", "json")) -def test_e2e( - tmp_dir, scm, dvc, seed_db, import_db_parameters, file_name, output_format -): - seed_db(values=range(5)) - - stage = dvc.imp_db(**import_db_parameters, output_format=output_format) - - output_file = tmp_dir / f"{file_name}.{output_format}" - output = load_table(output_file, output_format) - assert output.rows == [(i + 1, i) for i in range(5)] - - seed_db(values=range(5, 10)) - - dvc.update(stage.addressing) - - output = load_table(output_file, output_format) - assert output.rows == [(i + 1, i) for i in range(10)]