diff --git a/dvc/repo/imp_db.py b/dvc/repo/imp_db.py index 2de37b60d..eb7cd235f 100644 --- a/dvc/repo/imp_db.py +++ b/dvc/repo/imp_db.py @@ -49,7 +49,7 @@ def imp_db( # noqa: PLR0913 out = out or f"{model}.{output_format}" db.update({"model": model, "version": version, "project_dir": project_dir}) else: - out = out or "results.csv" + out = out or f"results.{output_format}" db["query"] = sql out = resolve_output(url or ".", out, force=force) diff --git a/pyproject.toml b/pyproject.toml index ef8f9f34d..2b9293da0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,8 +71,7 @@ dependencies = [ [project.optional-dependencies] all = ["dvc[azure,gdrive,gs,hdfs,oss,s3,ssh,webdav,webhdfs]"] azure = ["dvc-azure>=2.23.0"] -db = ["dbt-core>=1.5"] -dev = ["dvc[azure,gdrive,gs,hdfs,lint,oss,s3,ssh,tests,webdav,webhdfs,db]"] +dev = ["dvc[azure,gdrive,gs,hdfs,lint,oss,s3,ssh,tests,webdav,webhdfs]"] gdrive = ["dvc-gdrive==2.20"] gs = ["dvc-gs==2.22.1"] hdfs = ["dvc-hdfs==2.19"] @@ -111,6 +110,7 @@ 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` ] webdav = ["dvc-webdav==2.19.1"] webhdfs = ["dvc-webhdfs==2.19"] diff --git a/tests/func/test_import_db.py b/tests/func/test_import_db.py new file mode 100644 index 000000000..bfe35401c --- /dev/null +++ b/tests/func/test_import_db.py @@ -0,0 +1,78 @@ +"""import_db tests without involving any databases.""" + +from agate import Table + +data1 = Table([(i, i) for i in range(1, 10)], ["id", "value"]) +data2 = Table([(i, i) for i in range(10, 20)], ["id", "value"]) + + +def test_sql(mocker, tmp_dir, dvc): + mocker.patch("dvc.utils.db.execute_sql", side_effect=[data1, data2]) + stage = dvc.imp_db( + sql="select * from model", profile="profile", output_format="json" + ) + + db = {"file_format": "json", "profile": "profile", "query": "select * from model"} + + 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, + } + ], + } + + +def test_model(mocker, tmp_dir, dvc): + mocker.patch("dvc.utils.db.get_model", side_effect=[data1, data2]) + 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, + } + ], + } diff --git a/tests/func/test_import_db_e2e.py b/tests/func/test_import_db_e2e.py new file mode 100644 index 000000000..b13ab4c3c --- /dev/null +++ b/tests/func/test_import_db_e2e.py @@ -0,0 +1,128 @@ +"""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 io +import os +import sqlite3 +from contextlib import closing +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from agate import Table + + +dbt_connections = pytest.importorskip( + "dbt.adapters.sql.connections", reason="dbt-core not installed" +) +SQLConnectionManager = dbt_connections.SQLConnectionManager +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.execute("SELECT * FROM model")) as cursor: + return SQLConnectionManager.get_result_from_cursor(cursor, None) + + with closing(conn): + yield inner + + +@pytest.fixture +def dbt_profile(tmp_dir, 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), + } + }, + } + } + ) + return "sqlite_profile" + + +@pytest.fixture +def dbt_project(tmp_dir, dbt_profile): + (tmp_dir / "dbt_project.yml").dump({"name": "project", "profile": dbt_profile}) + return tmp_dir + + +@pytest.fixture +def dbt_model(dbt_project): + dbt_project.gen({"models": {"model.sql": "select * from models"}}) + return "model" + + +@pytest.fixture(params=("sql", "model")) +def import_db_parameters(request): + if request.param == "sql": + profile = request.getfixturevalue("dbt_profile") + return {"sql": "select * from model", "profile": profile} + return {"model": request.getfixturevalue("dbt_model")} + + +@pytest.fixture +def file_name(import_db_parameters): + if model := import_db_parameters.get("model"): + return model + return "results" + + +def format_output(table: "Table", output_format: str): + output = io.StringIO() + getattr(table, f"to_{output_format}")(output) + return output.getvalue() + + +@pytest.mark.filterwarnings("ignore::ResourceWarning") # dbt leaks fileobj from logger +@pytest.mark.parametrize("output_format", ("csv", "json")) +def test_e2e(tmp_dir, dvc, seed_db, import_db_parameters, file_name, output_format): + results: "Table" = 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}" + assert output_file.read_text() == format_output(results, output_format) + + results: "Table" = seed_db(values=range(6, 10)) + + dvc.update(stage.addressing) + assert output_file.read_text() == format_output(results, output_format)