tests: add functional test for db (#10101)

* tests: add functional test for db

* fix test

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)

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"]

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,
+            }
+        ],
+    }

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)