diff --git a/dvc/commands/imp_db.py b/dvc/commands/imp_db.py index b87b7bd96..d24e03680 100644 --- a/dvc/commands/imp_db.py +++ b/dvc/commands/imp_db.py @@ -18,7 +18,7 @@ class CmdTestDb(CmdBase): from dvc.dependency.db import _get_dbt_config from dvc.exceptions import DvcException - connection = self.args.connection + connection = self.args.conn db_config = self.repo.config.get("db", {}) config = db_config.get(connection, {}) @@ -34,11 +34,18 @@ class CmdTestDb(CmdBase): ) conn_config = merge(config, cli_config) - cli_dbt_config = {"profile": self.args.profile, "target": self.args.target} + 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: @@ -65,61 +72,87 @@ class CmdTestDb(CmdBase): class CmdImportDb(CmdBase): def run(self): - if not (self.args.sql or self.args.model): - raise argparse.ArgumentTypeError("Either of --sql or --model is required.") + 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, - connection=self.args.connection, + connection=self.args.conn, ) return 0 def add_parser(subparsers, parent_parser): - IMPORT_HELP = ( - "Download file or directory tracked by DVC or by Git " - "into the workspace, and track it." - ) - + IMPORT_HELP = "Snapshot a table or a SQL query result to a CSV/JSON format" import_parser = subparsers.add_parser( "import-db", parents=[parent_parser], - description=append_doc_link(IMPORT_HELP, "import"), - add_help=False, + description=append_doc_link(IMPORT_HELP, "import-db"), + help=IMPORT_HELP, + formatter_class=argparse.RawTextHelpFormatter, ) import_parser.add_argument( - "--url", help="Location of DVC or Git repository to download from" + "--url", + help=argparse.SUPPRESS, + # help="Location of dbt repository", ) import_parser.add_argument( "--rev", nargs="?", - help="Git revision (e.g. SHA, branch, tag)", + help=argparse.SUPPRESS, + # help="Git revision (e.g. SHA, branch, tag)", metavar="", ) import_parser.add_argument( - "--project-dir", nargs="?", help="Subdirectory to the dbt project location" + "--project-dir", + nargs="?", + help=argparse.SUPPRESS, + # help="Subdirectory to the dbt project location", ) - group = import_parser.add_mutually_exclusive_group() - group.add_argument( - "--sql", - help="SQL query", - ) + 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="Model name to download", + 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("--profile", help="Profile to use") - import_parser.add_argument("--target", help="Target to use") import_parser.add_argument( "--output-format", default="csv", @@ -144,10 +177,16 @@ def add_parser(subparsers, parent_parser): ) import_parser.add_argument( "--conn", - dest="connection", nargs="?", 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) @@ -158,7 +197,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", dest="connection") + test_db_parser.add_argument("--conn") + test_db_parser.add_argument("--dbt-conn", action="store_true", default=False) test_db_parser.add_argument("--url") test_db_parser.add_argument("--password") test_db_parser.add_argument("--username") diff --git a/dvc/database/dbt_query.py b/dvc/database/dbt_query.py index c37cf6014..a7f3d6479 100644 --- a/dvc/database/dbt_query.py +++ b/dvc/database/dbt_query.py @@ -1,6 +1,15 @@ import os from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ContextManager, + Dict, + Iterator, + Optional, + Union, +) from attrs import define, field @@ -41,6 +50,11 @@ class DbtClient: _, 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: @@ -50,8 +64,8 @@ class DbtClient: onerror() logger.exception( - "dvc was unable to connect to the specified database. " - "Please check your database credentials and try again.", + "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 diff --git a/dvc/database/serializer.py b/dvc/database/serializer.py index 13a7d0447..3858489a4 100644 --- a/dvc/database/serializer.py +++ b/dvc/database/serializer.py @@ -12,8 +12,8 @@ def noop(_): @dataclass class PandasSQLSerializer: - sql: "Union[sa.TextClause, str]" - con: "sa.Connection" + sql: "Union[str, sa.sql.Selectable]" + con: "sa.engine.Connection" chunksize: int = 10_000 def to_csv(self, file: str, progress=noop) -> None: diff --git a/dvc/database/sqla.py b/dvc/database/sqla.py index 6a2486672..6487deb8f 100644 --- a/dvc/database/sqla.py +++ b/dvc/database/sqla.py @@ -1,14 +1,18 @@ from contextlib import contextmanager from dataclasses import dataclass -from typing import Any, Callable, Dict, Iterator, Optional, Union +from typing import Any, Callable, ContextManager, Dict, Iterator, Optional, Union -from sqlalchemy import create_engine, text +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 @@ -30,9 +34,15 @@ class SQLAlchemyClient: engine: Engine @contextmanager - def query(self, sql: str) -> Iterator[PandasSQLSerializer]: + def query(self, sql: Union[str, Selectable]) -> Iterator[PandasSQLSerializer]: with self.engine.connect().execution_options(stream_results=True) as conn: - yield PandasSQLSerializer(text(sql), 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: @@ -66,6 +76,8 @@ def client( ) -> 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) diff --git a/dvc/dependency/__init__.py b/dvc/dependency/__init__.py index c96ef5f93..b555d1079 100644 --- a/dvc/dependency/__init__.py +++ b/dvc/dependency/__init__.py @@ -29,10 +29,10 @@ def _get(stage, p, info, **kwargs): if params: return ParamsDependency(stage, p, params) - if DbDependency.PARAM_QUERY in db: - return DbDependency(stage, info) - if db: + if DbtDependency.PARAM_MODEL in db: return DbtDependency(repo, stage, info) + if db: + return DbDependency(stage, info) assert p if repo: diff --git a/dvc/dependency/db.py b/dvc/dependency/db.py index 2df04d8d3..dfeee21ad 100644 --- a/dvc/dependency/db.py +++ b/dvc/dependency/db.py @@ -79,9 +79,11 @@ class AbstractDependency(Dependency): class DbDependency(AbstractDependency): - PARAM_QUERY = "query" PARAM_CONNECTION = "connection" - QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str} + PARAM_DBT_CONNECTION = "dbt_connection" + PARAM_QUERY = "query" + PARAM_TABLE = "table" + QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str, PARAM_TABLE: str} def __init__(self, stage: "Stage", info, *args, **kwargs): super().__init__(stage, info, *args, **kwargs) @@ -96,8 +98,9 @@ class DbDependency(AbstractDependency): from dvc.utils.humanize import truncate_text db_info = self.info.get(PARAM_DB, {}) - query = db_info.get(self.PARAM_QUERY, "[query]") - return truncate_text(query, 50) + 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) def workspace_status(self): return False # no workspace to check @@ -126,20 +129,29 @@ class DbDependency(AbstractDependency): db_info = self.info.get(PARAM_DB, {}) query = db_info.get(self.PARAM_QUERY) - if not query: - raise DvcException("Cannot download: no query specified") + table = db_info.get(self.PARAM_TABLE) + if not (query or table): + 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") 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: @@ -147,15 +159,12 @@ class DbDependency(AbstractDependency): with log_status("Testing connection") as status: db.test_connection(onerror=status.stop) - file_format = file_format or db_info.get(PARAM_FILE_FORMAT, "csv") - assert file_format - with log_status("Executing query") as status, db.query(query) as serializer: - 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) - with download_progress(to) as progress: - return export( - serializer, to.fs_path, format=file_format, progress=progress - ) + status.stop() + with download_progress(to) as prog: + return export(serializer, file, format=file_format, progress=prog) class DbtDependency(AbstractDependency): diff --git a/dvc/env.py b/dvc/env.py index 6f481f622..72c5b8900 100644 --- a/dvc/env.py +++ b/dvc/env.py @@ -16,5 +16,6 @@ DVC_STUDIO_OFFLINE = "DVC_STUDIO_OFFLINE" DVC_STUDIO_REPO_URL = "DVC_STUDIO_REPO_URL" DVC_STUDIO_TOKEN = "DVC_STUDIO_TOKEN" # noqa: S105 DVC_STUDIO_URL = "DVC_STUDIO_URL" +DVC_SQLALCHEMY_ECHO = "DVC_SQLALCHEMY_ECHO" DVC_SYSTEM_CONFIG_DIR = "DVC_SYSTEM_CONFIG_DIR" DVC_UPDATER_ENDPOINT = "DVC_UPDATER_ENDPOINT" diff --git a/dvc/repo/imp_db.py b/dvc/repo/imp_db.py index 00f494b99..9b71f996d 100644 --- a/dvc/repo/imp_db.py +++ b/dvc/repo/imp_db.py @@ -19,6 +19,7 @@ def imp_db( # noqa: PLR0913 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, @@ -34,6 +35,7 @@ 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} @@ -47,12 +49,12 @@ def imp_db( # noqa: PLR0913 db["profile"] = profile if model: - out = out or f"{model}.{output_format}" db.update({"model": model, "version": version, "project_dir": project_dir}) else: - out = out or f"results.{output_format}" - db.update({"query": sql, "connection": connection}) + db.update({"connection": connection, "query": sql, "table": table}) + file_name = model or table or "results" + out = out or f"{file_name}.{output_format}" out = resolve_output(url or ".", out, force=force) path, wdir, out = resolve_paths(self, out, always_local=True) stage = self.stage.create( diff --git a/tests/func/test_import_db.py b/tests/func/test_import_db.py index d4f82ec68..039bc89d5 100644 --- a/tests/func/test_import_db.py +++ b/tests/func/test_import_db.py @@ -15,11 +15,12 @@ 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 test_sql(client, tmp_dir, dvc): +def test_sql_with_dbt(client, tmp_dir, dvc): stage = dvc.imp_db( sql="select * from model", profile="profile", output_format="json" ) @@ -56,7 +57,7 @@ def test_sql(client, tmp_dir, dvc): } -def test_sql_conn_string(client, tmp_dir, dvc): +def test_sql_with_conn_string(client, tmp_dir, dvc): with dvc.config.edit(level="local") as conf: conf["db"] = {"conn": {"url": "conn"}} @@ -94,6 +95,42 @@ def test_sql_conn_string(client, tmp_dir, dvc): } +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") diff --git a/tests/func/test_import_db_e2e.py b/tests/func/test_import_db_e2e.py index f0934d5d6..3abbe9ddc 100644 --- a/tests/func/test_import_db_e2e.py +++ b/tests/func/test_import_db_e2e.py @@ -94,7 +94,7 @@ def dbt_model(dbt_project): return "model" -@pytest.fixture(params=("sql", "model", "external_model", "sql_conn")) +@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") @@ -102,6 +102,9 @@ def import_db_parameters(request: pytest.FixtureRequest): 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 { @@ -114,6 +117,8 @@ def import_db_parameters(request: pytest.FixtureRequest): 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" diff --git a/tests/unit/command/test_help.py b/tests/unit/command/test_help.py index 0db6796ea..d5be4c58c 100644 --- a/tests/unit/command/test_help.py +++ b/tests/unit/command/test_help.py @@ -28,7 +28,7 @@ def command_tuples(): # the no. of commands will usually go up, # but if we ever remove commands and drop below, adjust the magic number accordingly - assert len(commands) >= 96 + assert len(commands) >= 112 return sorted(commands)