testing: bench: use --dataset instead of --size (#10134)

This is no longer just about size (e.g. we have mnist/openorca/used-cars),
so the --dataset makes more sense.

.github/workflows/codeql.yml

@@ -34,6 +34,7 @@ jobs:
       uses: github/codeql-action/init@v2
       with:
         languages: ${{ matrix.language }}
+        queries: security-extended
 
     - name: Perform CodeQL Analysis
       uses: github/codeql-action/analyze@v2

dvc/cli/parser.py

@@ -26,6 +26,7 @@ from dvc.commands import (
     get_url,
     git_hook,
     imp,
+    imp_db,
     imp_url,
     init,
     install,
@@ -65,6 +66,7 @@ COMMANDS = [
     data_sync,
     gc,
     imp,
+    imp_db,
     imp_url,
     config,
     checkout,

dvc/commands/config.py

@@ -1,6 +1,8 @@
 import argparse
 import os
 
+from funcy import set_in
+
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -8,7 +10,7 @@ from dvc.ui import ui
 
 logger = logger.getChild(__name__)
 
-NAME_REGEX = r"^(?P<remote>remote\.)?(?P<section>[^\.]*)\.(?P<option>[^\.]*)$"
+NAME_REGEX = r"^(?P<top>(remote|db)\.)?(?P<section>[^\.]*)\.(?P<option>[^\.]*)$"
 
 
 def _name_type(value):
@@ -17,10 +19,12 @@ def _name_type(value):
     match = re.match(NAME_REGEX, value)
     if not match:
         raise argparse.ArgumentTypeError(
-            "name argument should look like remote.name.option or section.option"
+            "name argument should look like remote.name.option or "
+            "db.name.option or section.option"
         )
+    top = match.group("top")
     return (
-        bool(match.group("remote")),
+        top.strip(".") if top else None,
         match.group("section").lower(),
         match.group("option").lower(),
     )
@@ -49,12 +53,11 @@ class CmdConfig(CmdBaseNoRepo):
             logger.error("name argument is required")
             return 1
 
-        remote, section, opt = self.args.name
+        remote_or_db, section, opt = self.args.name
 
         if self.args.value is None and not self.args.unset:
-            return self._get(remote, section, opt)
-
-        return self._set(remote, section, opt)
+            return self._get(remote_or_db, section, opt)
+        return self._set(remote_or_db, section, opt)
 
     def _list(self):
         if any((self.args.name, self.args.value, self.args.unset)):
@@ -75,18 +78,18 @@ class CmdConfig(CmdBaseNoRepo):
 
         return 0
 
-    def _get(self, remote, section, opt):
+    def _get(self, remote_or_db, section, opt):
         from dvc.config import ConfigError
 
         levels = self._get_appropriate_levels(self.args.level)[::-1]
 
         for level in levels:
             conf = self.config.read(level)
-            if remote:
-                conf = conf["remote"]
+            if remote_or_db:
+                conf = conf[remote_or_db]
 
             try:
-                self._check(conf, remote, section, opt)
+                self._check(conf, remote_or_db, section, opt)
             except ConfigError:
                 if self.args.level:
                     raise
@@ -99,16 +102,15 @@ class CmdConfig(CmdBaseNoRepo):
 
         return 0
 
-    def _set(self, remote, section, opt):
+    def _set(self, remote_or_db, section, opt):
         with self.config.edit(self.args.level) as conf:
-            if remote:
-                conf = conf["remote"]
+            if remote_or_db:
+                conf = conf[remote_or_db]
             if self.args.unset:
-                self._check(conf, remote, section, opt)
+                self._check(conf, remote_or_db, section, opt)
                 del conf[section][opt]
             else:
-                self._check(conf, remote, section)
-                conf[section][opt] = self.args.value
+                conf.update(set_in(conf, [section, opt], self.args.value))
 
         if self.args.name == "cache.type":
             logger.warning(
@@ -119,10 +121,10 @@ class CmdConfig(CmdBaseNoRepo):
 
         return 0
 
-    def _check(self, conf, remote, section, opt=None):
+    def _check(self, conf, remote_or_db, section, opt=None):
         from dvc.config import ConfigError
 
-        name = "remote" if remote else "section"
+        name = remote_or_db or "section"
         if section not in conf:
             raise ConfigError(f"{name} '{section}' doesn't exist")
 

dvc/commands/imp_db.py

@@ -0,0 +1,167 @@
+import argparse
+
+from funcy import compact, merge
+
+from dvc.cli import completion
+from dvc.cli.command import CmdBase
+from dvc.cli.utils import append_doc_link
+from dvc.log import logger
+from dvc.ui import ui
+
+logger = logger.getChild(__name__)
+
+
+class CmdTestDb(CmdBase):
+    def run(self):
+        from dvc.database import get_adapter
+        from dvc.database.dbt_utils import DBT_PROJECT_FILE, is_dbt_project
+        from dvc.dependency.db import _get_dbt_config
+        from dvc.exceptions import DvcException
+
+        connection = self.args.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")
+
+        cli_config = compact(
+            {
+                "url": self.args.url,
+                "username": self.args.username,
+                "password": self.args.password,
+            }
+        )
+        conn_config = merge(config, cli_config)
+
+        cli_dbt_config = {"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 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_adapter(conn_config, project_dir=project_dir, **dbt_config)
+        with adapter 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):
+        if not (self.args.sql or self.args.model):
+            raise argparse.ArgumentTypeError("Either of --sql or --model is required.")
+
+        self.repo.imp_db(
+            url=self.args.url,
+            rev=self.args.rev,
+            project_dir=self.args.project_dir,
+            sql=self.args.sql,
+            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,
+        )
+        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_parser = subparsers.add_parser(
+        "import-db",
+        parents=[parent_parser],
+        description=append_doc_link(IMPORT_HELP, "import"),
+        add_help=False,
+    )
+    import_parser.add_argument(
+        "--url", help="Location of DVC or Git repository to download from"
+    )
+    import_parser.add_argument(
+        "--rev",
+        nargs="?",
+        help="Git revision (e.g. SHA, branch, tag)",
+        metavar="<commit>",
+    )
+    import_parser.add_argument(
+        "--project-dir", nargs="?", help="Subdirectory to the dbt project location"
+    )
+
+    group = import_parser.add_mutually_exclusive_group()
+    group.add_argument(
+        "--sql",
+        help="SQL query",
+    )
+    group.add_argument(
+        "--model",
+        help="Model name to download",
+    )
+    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",
+        const="csv",
+        nargs="?",
+        choices=["csv", "json"],
+        help="Export format",
+    )
+    import_parser.add_argument(
+        "-o",
+        "--out",
+        nargs="?",
+        help="Destination path to download files to",
+        metavar="<path>",
+    ).complete = completion.FILE
+    import_parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        default=False,
+        help="Override destination file or folder if exists.",
+    )
+    import_parser.add_argument(
+        "--conn",
+        dest="connection",
+        nargs="?",
+        help="Database connection to use, needs to be set in config",
+    )
+
+    import_parser.set_defaults(func=CmdImportDb)
+
+    TEST_DB_HELP = "Test the database connection"
+    test_db_parser = subparsers.add_parser(
+        "test-db",
+        parents=[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("--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)

dvc/config.py

@@ -390,10 +390,10 @@ class Config(dict):
 
 
 def _parse_named(conf):
-    result: Dict[str, Dict] = {"remote": {}, "machine": {}}
+    result: Dict[str, Dict] = {"remote": {}, "machine": {}, "db": {}}
 
     for section, val in conf.items():
-        match = re_find(r'^\s*(remote|machine)\s*"(.*)"\s*$', section)
+        match = re_find(r'^\s*(remote|machine|db)\s*"(.*)"\s*$', section)
         if match:
             key, name = match
             result[key][name] = val
@@ -408,7 +408,7 @@ def _pack_named(conf):
     result = compact(conf)
 
     # Transform remote.name -> 'remote "name"'
-    for key in ("remote", "machine"):
+    for key in ("remote", "machine", "db"):
         for name, val in conf[key].items():
             result[f'{key} "{name}"'] = val
         result.pop(key, None)

dvc/config_schema.py

@@ -318,6 +318,8 @@ SCHEMA = {
     "feature": FeatureSchema(
         {
             Optional("machine", default=False): Bool,
+            "dbt_profile": str,
+            "dbt_target": str,
         },
     ),
     "plots": {
@@ -350,4 +352,11 @@ SCHEMA = {
         "repo_url": str,
         Optional("offline", default=False): Bool,
     },
+    "db": {
+        str: {
+            "url": str,
+            "username": str,
+            "password": str,
+        },
+    },
 }

dvc/database/__init__.py

@@ -0,0 +1,22 @@
+from typing import Any, ContextManager, Optional, Union
+
+from . import dbt_query, sqla
+from .dbt_models import get_model
+from .serializer import export
+
+Adapter = Union[sqla.SQLAlchemyAdapter, dbt_query.DbtAdapter]
+
+
+def get_adapter(
+    config,
+    project_dir: Optional[str] = None,
+    profile: Optional[str] = None,
+    target: Optional[str] = None,
+    **kwargs: Any,
+) -> "ContextManager[Adapter]":
+    if config:
+        return sqla.adapter(config, **kwargs)
+    return dbt_query.adapter(project_dir=project_dir, profile=profile, target=target)
+
+
+__all__ = ["export", "get_adapter", "get_model", "Adapter"]

dvc/database/dbt_models.py

@@ -0,0 +1,40 @@
+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)

dvc/database/dbt_query.py

@@ -0,0 +1,88 @@
+import os
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, 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 DbtAdapter:
+    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 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(
+                    "dvc was unable to connect to the specified database. "
+                    "Please 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 adapter(
+    project_dir: Optional[str] = None,
+    profiles_dir: Optional[str] = None,
+    profile: Optional[str] = None,
+    target: Optional[str] = None,
+) -> Iterator["DbtAdapter"]:
+    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 DbtAdapter(adapter, creds)

dvc/database/dbt_runner.py

@@ -0,0 +1,61 @@
+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

dvc/database/dbt_utils.py

@@ -0,0 +1,100 @@
+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)

dvc/database/serializer.py

@@ -0,0 +1,54 @@
+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[sa.TextClause, str]"
+    con: "sa.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)

dvc/database/sqla.py

@@ -0,0 +1,73 @@
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterator, Optional, Union
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import URL, Engine
+from sqlalchemy.engine import make_url as _make_url
+from sqlalchemy.exc import NoSuchModuleError
+
+from dvc.exceptions import DvcException
+from dvc.log import logger
+
+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 SQLAlchemyAdapter:
+    engine: Engine
+
+    @contextmanager
+    def query(self, sql: str) -> Iterator[PandasSQLSerializer]:
+        with self.engine.connect().execution_options(stream_results=True) as conn:
+            yield PandasSQLSerializer(text(sql), conn)
+
+    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(
+                "dvc was unable to connect to the specified database. "
+                "Please 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 adapter(
+    url_or_config: Union[str, URL, Dict[str, str]], **engine_kwargs: Any
+) -> Iterator[SQLAlchemyAdapter]:
+    url = url_from_config(url_or_config)
+    with handle_error(url):
+        engine = create_engine(url, **engine_kwargs)
+    try:
+        yield SQLAlchemyAdapter(engine)
+    finally:
+        engine.dispose()

dvc/dependency/__init__.py

@@ -4,6 +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 .param import ParamsDependency
 from .repo import RepoDependency
 
@@ -14,20 +15,28 @@ from .repo import RepoDependency
 SCHEMA: Mapping[str, Any] = {
     **ARTIFACT_SCHEMA,
     **RepoDependency.REPO_SCHEMA,
+    **DB_SCHEMA,
     Output.PARAM_FILES: [DIR_FILES_SCHEMA],
     Output.PARAM_FS_CONFIG: dict,
 }
 
 
 def _get(stage, p, info, **kwargs):
-    if info and info.get(RepoDependency.PARAM_REPO):
-        repo = info.pop(RepoDependency.PARAM_REPO)
-        return RepoDependency(repo, stage, p, info)
+    d = info or {}
+    db = d.get(PARAM_DB, {})
+    params = d.pop(ParamsDependency.PARAM_PARAMS, None)
+    repo = d.pop(RepoDependency.PARAM_REPO, None)
 
-    if info and info.get(ParamsDependency.PARAM_PARAMS):
-        params = info.pop(ParamsDependency.PARAM_PARAMS)
+    if params:
         return ParamsDependency(stage, p, params)
+    if DbDependency.PARAM_QUERY in db:
+        return DbDependency(stage, info)
+    if db:
+        return DbtDependency(repo, stage, info)
 
+    assert p
+    if repo:
+        return RepoDependency(repo, stage, p, info)
     return Dependency(stage, p, info, **kwargs)
 
 
@@ -44,9 +53,11 @@ def loadd_from(stage, d_list):
     return ret
 
 
-def loads_from(stage, s_list, erepo=None, fs_config=None):
+def loads_from(stage, s_list, erepo=None, fs_config=None, db=None):
     assert isinstance(s_list, list)
     info = {RepoDependency.PARAM_REPO: erepo} if erepo else {}
+    if db:
+        info.update({"db": db})
     return [
         _get(
             stage,

dvc/dependency/db.py

@@ -0,0 +1,324 @@
+import os
+from contextlib import contextmanager, nullcontext
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
+
+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
+    from dvc.ui._rich_progress import DbDownloadProgress
+
+    with log_durations(logger.debug, f"Saving to {to}"), DbDownloadProgress(
+        console=ui.error_console,
+    ) as progress:
+        task = progress.add_task("Saving", total=None, output=to)
+        yield lambda n: progress.advance(task, advance=n)
+        progress.update(task, description="Saved", total=0)
+
+
+class AbstractDependency(Dependency):
+    """Dependency without workspace/fs/fs_path"""
+
+    def __init__(self, stage: "Stage", info: Dict[str, Any], *args, **kwargs):
+        self.repo = stage.repo
+        self.stage = stage
+        self.fs = None
+        self.fs_path = None
+        self.def_path = None  # type: ignore[assignment]
+        self.info = info or {}
+
+
+class DbDependency(AbstractDependency):
+    PARAM_QUERY = "query"
+    PARAM_CONNECTION = "connection"
+    QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str}
+
+    def __init__(self, stage: "Stage", info, *args, **kwargs):
+        super().__init__(stage, info, *args, **kwargs)
+        self.target = None
+
+    def __repr__(self):
+        return "{}:{}".format(
+            self.__class__.__name__, "".join(f"{k}=={v}" for k, v in self.info.items())
+        )
+
+    def __str__(self):
+        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)
+
+    def workspace_status(self):
+        return False  # no workspace to check
+
+    def status(self):
+        return self.workspace_status()
+
+    def save(self):
+        """nothing to save."""
+
+    def dumpd(self, **kwargs):
+        db_info = compact(self.info.get(PARAM_DB, {}))
+        return {PARAM_DB: db_info} if db_info else {}
+
+    def update(self, rev=None):
+        """nothing to update."""
+
+    def download(
+        self,
+        to: "Output",
+        jobs: Optional[int] = None,  # noqa: ARG002
+        file_format: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        from dvc.database import export, get_adapter
+
+        db_info = self.info.get(PARAM_DB, {})
+        query = db_info.get(self.PARAM_QUERY)
+        if not query:
+            raise DvcException("Cannot download: no query 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 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
+        with get_adapter(
+            config, project_dir=project_dir, profile=profile, target=target
+        ) as db:
+            logger.debug("using adapter: %s", db)
+            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()
+                logger.debug("using serializer: %s", serializer)
+                with download_progress(to) as progress:
+                    return export(
+                        serializer, to.fs_path, format=file_format, progress=progress
+                    )
+
+
+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,
+    },
+}

dvc/dependency/repo.py

@@ -64,28 +64,31 @@ class RepoDependency(Dependency):
         if self.def_repo.get(self.PARAM_REV_LOCK) is None:
             self.def_repo[self.PARAM_REV_LOCK] = rev
 
-    def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]:
-        repo = {self.PARAM_URL: self.def_repo[self.PARAM_URL]}
+    @classmethod
+    def _dump_def_repo(cls, def_repo) -> Dict[str, str]:
+        repo = {cls.PARAM_URL: def_repo[cls.PARAM_URL]}
 
-        rev = self.def_repo.get(self.PARAM_REV)
+        rev = def_repo.get(cls.PARAM_REV)
         if rev:
-            repo[self.PARAM_REV] = self.def_repo[self.PARAM_REV]
+            repo[cls.PARAM_REV] = def_repo[cls.PARAM_REV]
 
-        rev_lock = self.def_repo.get(self.PARAM_REV_LOCK)
+        rev_lock = def_repo.get(cls.PARAM_REV_LOCK)
         if rev_lock:
-            repo[self.PARAM_REV_LOCK] = rev_lock
+            repo[cls.PARAM_REV_LOCK] = rev_lock
 
-        config = self.def_repo.get(self.PARAM_CONFIG)
+        config = def_repo.get(cls.PARAM_CONFIG)
         if config:
-            repo[self.PARAM_CONFIG] = config
+            repo[cls.PARAM_CONFIG] = config
 
-        remote = self.def_repo.get(self.PARAM_REMOTE)
+        remote = def_repo.get(cls.PARAM_REMOTE)
         if remote:
-            repo[self.PARAM_REMOTE] = remote
+            repo[cls.PARAM_REMOTE] = remote
+        return repo
 
+    def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]:
         return {
             self.PARAM_PATH: self.def_path,
-            self.PARAM_REPO: repo,
+            self.PARAM_REPO: self._dump_def_repo(self.def_repo),
         }
 
     def update(self, rev: Optional[str] = None):

dvc/dirs.py

@@ -15,7 +15,7 @@ def system_config_dir():
 
 
 def global_config_dir():
-    return os.getenv(env.DVC_SYSTEM_CONFIG_DIR) or platformdirs.user_config_dir(
+    return os.getenv(env.DVC_GLOBAL_CONFIG_DIR) or platformdirs.user_config_dir(
         APPNAME, APPAUTHOR
     )
 

dvc/fs/__init__.py

@@ -5,6 +5,8 @@ from dvc_http import HTTPFileSystem, HTTPSFileSystem  # noqa: F401
 
 from dvc.config import ConfigError as RepoConfigError
 from dvc.config_schema import SCHEMA, Invalid
+
+# pylint: disable=unused-import
 from dvc_objects.fs import (  # noqa: F401
     LocalFileSystem,
     MemoryFileSystem,
@@ -27,7 +29,7 @@ from dvc_objects.fs.path import Path  # noqa: F401
 
 from .callbacks import Callback
 from .data import DataFileSystem  # noqa: F401
-from .dvc import DVCFileSystem  # noqa: F401
+from .dvc import DVCFileSystem
 from .git import GitFileSystem  # noqa: F401
 
 known_implementations.update(
@@ -47,6 +49,8 @@ known_implementations.update(
 def download(
     fs: "FileSystem", fs_path: str, to: str, jobs: Optional[int] = None
 ) -> int:
+    from dvc.scm import lfs_prefetch
+
     with Callback.as_tqdm_callback(
         desc=f"Downloading {fs.path.name(fs_path)}",
         unit="files",
@@ -70,6 +74,8 @@ def download(
             from_infos = [fs_path]
             to_infos = [to]
 
+        if isinstance(fs, DVCFileSystem):
+            lfs_prefetch(fs, from_infos)
         cb.set_size(len(from_infos))
         jobs = jobs or fs.jobs
         generic.copy(

dvc/output.py

@@ -1523,7 +1523,7 @@ CLOUD_SCHEMA = vol.All({str: {**META_SCHEMA, **CHECKSUMS_SCHEMA}}, vol.Length(ma
 ARTIFACT_SCHEMA: Dict[Any, Any] = {
     **CHECKSUMS_SCHEMA,
     **META_SCHEMA,
-    vol.Required(Output.PARAM_PATH): str,
+    Output.PARAM_PATH: str,
     Output.PARAM_PERSIST: bool,
     Output.PARAM_CLOUD: CLOUD_SCHEMA,
     Output.PARAM_HASH: str,

dvc/repo/__init__.py

@@ -76,6 +76,7 @@ class Repo:
     from dvc.repo.get import get as _get  # type: ignore[misc]
     from dvc.repo.get_url import get_url as _get_url  # type: ignore[misc]
     from dvc.repo.imp import imp  # type: ignore[misc]
+    from dvc.repo.imp_db import imp_db  # type: ignore[misc]
     from dvc.repo.imp_url import imp_url  # type: ignore[misc]
     from dvc.repo.install import install  # type: ignore[misc]
     from dvc.repo.ls import ls as _ls  # type: ignore[misc]

dvc/repo/graph.py

@@ -135,6 +135,8 @@ def build_graph(stages, outs_trie=None):
     for stage in stages:
         if stage.is_repo_import:
             continue
+        if stage.is_db_import:
+            continue
 
         for dep in stage.deps:
             dep_key = dep.fs.path.parts(dep.fs_path)
@@ -160,6 +162,8 @@ def build_outs_graph(graph, outs_trie):
     for stage in graph.nodes():
         if stage.is_repo_import:
             continue
+        if stage.is_db_import:
+            continue
         for dep in stage.deps:
             dep_key = dep.fs.path.parts(dep.fs_path)
             overlapping = [n.value for n in outs_trie.prefixes(dep_key)]

dvc/repo/imp_db.py

@@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from dvc.exceptions import OutputDuplicationError
+from dvc.repo.scm_context import scm_context
+from dvc.ui import ui
+from dvc.utils import resolve_output, resolve_paths
+
+if TYPE_CHECKING:
+    from . import Repo
+
+from . import locked
+
+
+@locked
+@scm_context
+def imp_db(  # noqa: PLR0913
+    self: "Repo",
+    url: Optional[str] = None,
+    rev: Optional[str] = None,
+    project_dir: Optional[str] = None,
+    sql: 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,
+    connection: Optional[str] = None,
+):
+    ui.warn("WARNING: import-db is an experimental feature.")
+    ui.warn(
+        "The functionality may change or break without notice, "
+        "which could lead to unexpected behavior."
+    )
+    erepo = None
+    if model and url:
+        erepo = {"url": url}
+        if rev:
+            erepo["rev"] = rev
+
+    assert output_format in ("csv", "json")
+
+    db: Dict[str, Any] = {"file_format": output_format}
+    if profile:
+        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})
+
+    out = resolve_output(url or ".", out, force=force)
+    path, wdir, out = resolve_paths(self, out, always_local=True)
+    stage = self.stage.create(
+        single_stage=True,
+        validate=False,
+        fname=path,
+        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:
+        raise OutputDuplicationError(  # noqa: B904
+            exc.output, set(exc.stages) - {stage}
+        )
+
+    stage.run()
+    stage.frozen = frozen
+    stage.dump()
+    return stage

dvc/repo/index.py

@@ -135,7 +135,11 @@ def _load_data_from_outs(index, prefix, outs):
             hash_info=out.hash_info,
         )
 
-        if out.stage.is_import and not out.stage.is_repo_import:
+        if (
+            out.stage.is_import
+            and not out.stage.is_repo_import
+            and not out.stage.is_db_import
+        ):
             dep = out.stage.deps[0]
             entry.meta = dep.meta
             if out.hash_info:
@@ -183,6 +187,9 @@ def _load_storage_from_out(storage_map, key, out):
     except NoRemoteError:
         pass
 
+    if out.stage.is_db_import:
+        return
+
     if out.stage.is_import:
         dep = out.stage.deps[0]
         if not out.hash_info:

dvc/repo/worktree.py

@@ -101,7 +101,7 @@ def _get_remote(
     return repo.cloud.get_remote(name, command)
 
 
-def _merge_push_meta(
+def _merge_push_meta(  # noqa: C901
     out: "Output",
     index: Union["DataIndex", "DataIndexView"],
     remote: Optional[str] = None,
@@ -116,7 +116,11 @@ def _merge_push_meta(
     from dvc_data.index.save import build_tree
 
     _, key = out.index_key
-    entry = index[key]
+
+    entry = index.get(key)
+    if entry is None:
+        return
+
     repo = out.stage.repo
     if out.isdir():
         old_tree = out.get_obj()

dvc/scm.py

@@ -24,6 +24,8 @@ from dvc.progress import Tqdm
 if TYPE_CHECKING:
     from scmrepo.progress import GitProgressEvent
 
+    from dvc.fs import FileSystem
+
 
 class SCMError(DvcException):
     """Base class for source control management errors."""
@@ -262,3 +264,30 @@ def iter_revs(
 
     rev_resolver = partial(resolve_rev, scm)
     return group_by(rev_resolver, results)
+
+
+def lfs_prefetch(fs: "FileSystem", paths: List[str]):
+    from scmrepo.git.lfs import fetch as _lfs_fetch
+
+    from dvc.fs.dvc import DVCFileSystem
+    from dvc.fs.git import GitFileSystem
+
+    if isinstance(fs, DVCFileSystem) and isinstance(fs.repo.fs, GitFileSystem):
+        git_fs = fs.repo.fs
+        scm = fs.repo.scm
+        assert isinstance(scm, Git)
+    else:
+        return
+
+    try:
+        if "filter=lfs" not in git_fs.open(".gitattributes").read():
+            return
+    except OSError:
+        return
+    with TqdmGit(desc="Checking for Git-LFS objects") as pbar:
+        _lfs_fetch(
+            scm,
+            [git_fs.rev],
+            include=[(path if path.startswith("/") else f"/{path}") for path in paths],
+            progress=pbar.update_git,
+        )

dvc/stage/__init__.py

@@ -105,7 +105,7 @@ def create_stage(cls: Type[_T], repo, path, **kwargs) -> _T:
     fill_stage_outputs(stage, **kwargs)
     check_no_externals(stage)
     fill_stage_dependencies(
-        stage, **project(kwargs, ["deps", "erepo", "params", "fs_config"])
+        stage, **project(kwargs, ["deps", "erepo", "params", "fs_config", "db"])
     )
     check_circular_dependency(stage)
     check_duplicated_arguments(stage)
@@ -281,9 +281,24 @@ class Stage(params.StageParams):
 
         return isinstance(self.deps[0], RepoDependency)
 
+    @property
+    def is_db_import(self) -> bool:
+        if not self.is_import:
+            return False
+
+        from dvc.dependency import DbDependency, DbtDependency
+
+        return isinstance(self.deps[0], (DbDependency, DbtDependency))
+
     @property
     def is_versioned_import(self) -> bool:
-        return self.is_import and self.deps[0].fs.version_aware
+        from dvc.dependency import DbDependency, DbtDependency
+
+        return (
+            self.is_import
+            and not isinstance(self.deps[0], (DbDependency, DbtDependency))
+            and self.deps[0].fs.version_aware
+        )
 
     def short_description(self) -> Optional["str"]:
         desc: Optional["str"] = None
@@ -446,6 +461,9 @@ class Stage(params.StageParams):
     ) -> None:
         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
+        force = self.is_db_import
         update_import(
             self,
             rev=rev,
@@ -453,6 +471,7 @@ class Stage(params.StageParams):
             remote=remote,
             no_download=no_download,
             jobs=jobs,
+            force=force,
         )
 
     def reload(self) -> "Stage":

dvc/stage/decorators.py

@@ -7,6 +7,7 @@ from funcy import decorator
 def rwlocked(call, read=None, write=None):
     import sys
 
+    from dvc.dependency.db import AbstractDependency
     from dvc.dependency.repo import RepoDependency
     from dvc.rwlock import rwlock
 
@@ -27,7 +28,7 @@ def rwlocked(call, read=None, write=None):
             for item in getattr(stage, attr)
             # There is no need to lock RepoDependency deps, as there is no
             # corresponding OutputREPO, so we can't even write it.
-            if not isinstance(item, RepoDependency)
+            if not isinstance(item, (RepoDependency, AbstractDependency))
         ]
 
     cmd = " ".join(sys.argv)

dvc/stage/imports.py

@@ -18,7 +18,13 @@ def _update_import_on_remote(stage, remote, jobs):
 
 
 def update_import(
-    stage, rev=None, to_remote=False, remote=None, no_download=None, jobs=None
+    stage,
+    rev=None,
+    to_remote=False,
+    remote=None,
+    no_download=None,
+    jobs=None,
+    force=False,
 ):
     stage.deps[0].update(rev=rev)
 
@@ -30,7 +36,7 @@ def update_import(
         if to_remote:
             _update_import_on_remote(stage, remote, jobs)
         else:
-            stage.reproduce(no_download=no_download, jobs=jobs)
+            stage.reproduce(no_download=no_download, jobs=jobs, force=force)
     finally:
         if no_download and changed:
             # Avoid retaining stale information
@@ -55,7 +61,7 @@ def sync_import(
     else:
         stage.save_deps()
         if no_download:
-            if stage.is_repo_import:
+            if stage.is_repo_import or stage.is_db_import:
                 stage.deps[0].update()
         else:
             stage.deps[0].download(

dvc/stage/utils.py

@@ -77,12 +77,14 @@ def fill_stage_outputs(stage, **kwargs):
         )
 
 
-def fill_stage_dependencies(stage, deps=None, erepo=None, params=None, fs_config=None):
+def fill_stage_dependencies(
+    stage, deps=None, erepo=None, params=None, fs_config=None, db=None
+):
     from dvc.dependency import loads_from, loads_params
 
     assert not stage.deps
     stage.deps = []
-    stage.deps += loads_from(stage, deps or [], erepo=erepo, fs_config=fs_config)
+    stage.deps += loads_from(stage, deps or [], erepo=erepo, fs_config=fs_config, db=db)
     stage.deps += loads_params(stage, params or [])
 
 

dvc/testing/benchmarks/cli/commands/test_fetch.py

@@ -0,0 +1,3 @@
+def test_fetch(bench_dvc, tmp_dir, dvc, make_dataset, remote):
+    make_dataset(cache=False, dvcfile=True, files=False, remote=True)
+    bench_dvc("fetch")

dvc/testing/benchmarks/cli/commands/test_pull.py

@@ -1,3 +0,0 @@
-def test_pull(bench_dvc, tmp_dir, dvc, make_dataset, remote):
-    make_dataset(cache=False, dvcfile=True, files=False, remote=True)
-    bench_dvc("pull")

dvc/testing/benchmarks/cli/stories/use_cases/test_sharing.py

@@ -11,7 +11,8 @@ def test_sharing(bench_dvc, tmp_dir, dvc, dataset, remote):
     shutil.rmtree(dataset)
     shutil.rmtree(tmp_dir / ".dvc" / "cache")
 
-    bench_dvc("pull")
-    bench_dvc("pull", name="noop")
+    bench_dvc("fetch")
+    bench_dvc("fetch", name="noop")
 
+    bench_dvc("checkout")
     bench_dvc("checkout", name="noop")

dvc/testing/benchmarks/fixtures.py

@@ -2,26 +2,48 @@ import os
 import shutil
 from pathlib import Path
 from subprocess import check_output
+from typing import Dict, Optional
 
 import pytest
 from dulwich.porcelain import clone
 from funcy import first
 from packaging import version
 
+from dvc.types import StrPath
+
 
 @pytest.fixture(scope="session")
 def bench_config(request):
     return request.config.bench_config
 
 
+class VirtualEnv:
+    def __init__(self, path: StrPath) -> None:
+        self.path = Path(path)
+        self.bin = self.path / ("Scripts" if os.name == "nt" else "bin")
+
+    def create(self) -> None:
+        import virtualenv
+
+        virtualenv.cli_run([os.fspath(self.path)])
+
+    def run(self, cmd: str, *args: str, env: Optional[Dict[str, str]] = None) -> None:
+        exe = self.which(cmd)
+        check_output([exe, *args], env=env)  # noqa: S603
+
+    def which(self, cmd: str) -> str:
+        assert self.bin.exists()
+        return shutil.which(cmd, path=self.bin) or cmd
+
+
 @pytest.fixture(scope="session")
 def make_dvc_venv(tmp_path_factory):
     def _make_dvc_venv(name):
-        from pytest_virtualenv import VirtualEnv
-
         name = _sanitize_venv_name(name)
         venv_dir = tmp_path_factory.mktemp(f"dvc-venv-{name}")
-        return VirtualEnv(workspace=venv_dir)
+        venv = VirtualEnv(venv_dir)
+        venv.create()
+        return venv
 
     return _make_dvc_venv
 
@@ -74,16 +96,16 @@ def make_dvc_bin(
         venv = dvc_venvs.get(dvc_rev)
         if not venv:
             venv = make_dvc_venv(dvc_rev)
-            venv.run("pip install -U pip")
+            venv.run("pip", "install", "-U", "pip")
             if bench_config.dvc_install_deps:
                 egg = f"dvc[{bench_config.dvc_install_deps}]"
             else:
                 egg = "dvc"
-            venv.run(f"pip install git+file://{dvc_git_repo}@{dvc_rev}#egg={egg}")
+            venv.run("pip", "install", f"git+file://{dvc_git_repo}@{dvc_rev}#egg={egg}")
             if dvc_rev in ["2.18.1", "2.11.0", "2.6.3"]:
-                venv.run("pip install fsspec==2022.11.0")
+                venv.run("pip", "install", "fsspec==2022.11.0")
             dvc_venvs[dvc_rev] = venv
-        dvc_bin = venv.virtualenv / "bin" / "dvc"
+        dvc_bin = venv.which("dvc")
     else:
         dvc_bin = bench_config.dvc_bin
 
@@ -171,7 +193,7 @@ def make_dataset(request, bench_config, tmp_dir, dvc_bench_git_repo):
 
         path = tmp_dir / "dataset"
         root = dvc_bench_git_repo
-        src = root / "data" / bench_config.size / "dataset"
+        src = root / "data" / bench_config.dataset / "dataset"
         src_dvc = src.with_suffix(".dvc")
 
         dvc = Repo(root)

dvc/testing/benchmarks/plugin.py

@@ -4,6 +4,11 @@ DEFAULT_DVC_BENCH_GIT_REPO = "https://github.com/iterative/dvc-bench"
 DEFAULT_PROJECT_GIT_REPO = "https://github.com/iterative/example-get-started"
 
 
+def pytest_report_header(config):
+    bconf = config.bench_config
+    return "dvc-bench: (" f"dataset: '{bconf.dataset}'," f"revs: '{bconf.dvc_revs}'" ")"
+
+
 def pytest_generate_tests(metafunc):
     str_revs = metafunc.config.getoption("--dvc-revs")
     revs = str_revs.split(",") if str_revs else [None]
@@ -13,7 +18,7 @@ def pytest_generate_tests(metafunc):
 
 class DVCBenchConfig:
     def __init__(self):
-        self.size = "small"
+        self.dataset = "small"
         self.dvc_bin = DEFAULT_DVC_BIN
         self.dvc_revs = None
         self.dvc_git_repo = DEFAULT_DVC_GIT_REPO
@@ -29,7 +34,7 @@ def pytest_configure(config):
     )
 
     config.bench_config = DVCBenchConfig()
-    config.bench_config.size = config.getoption("--size")
+    config.bench_config.dataset = config.getoption("--dataset")
     config.bench_config.dvc_bin = config.getoption("--dvc-bin")
     config.bench_config.dvc_revs = config.getoption("--dvc-revs")
     config.bench_config.dvc_git_repo = config.getoption("--dvc-git-repo")
@@ -41,10 +46,10 @@ def pytest_configure(config):
 
 def pytest_addoption(parser):
     parser.addoption(
-        "--size",
-        choices=["tiny", "small", "large", "mnist"],
+        "--dataset",
+        type=str,
         default="small",
-        help="Size of the dataset/datafile to use in tests",
+        help="Dataset name to use in tests (e.g. tiny/small/large/mnist/etc)",
     )
 
     parser.addoption(

dvc/testing/tmp_dir.py

@@ -70,16 +70,20 @@ class TmpDir(pathlib.Path):
     def __new__(cls, *args, **kwargs):
         if cls is TmpDir:
             cls = WindowsTmpDir if os.name == "nt" else PosixTmpDir
-        # init parameter and `_init` method has been removed in Python 3.10.
-        kw = {"init": False} if sys.version_info < (3, 10) else {}
-        self = cls._from_parts(args, **kw)  # type: ignore[attr-defined]
-        if not self._flavour.is_supported:
-            raise NotImplementedError(
-                f"cannot instantiate {cls.__name__!r} on your system"
-            )
-        if sys.version_info < (3, 10):
-            self._init()
-        return self
+
+        if sys.version_info >= (3, 12):
+            return object.__new__(cls)
+        else:  # noqa: RET505
+            # init parameter and `_init` method has been removed in Python 3.10.
+            kw = {"init": False} if sys.version_info < (3, 10) else {}
+            self = cls._from_parts(args, **kw)  # type: ignore[attr-defined]
+            if not self._flavour.is_supported:
+                raise NotImplementedError(
+                    f"cannot instantiate {cls.__name__!r} on your system"
+                )
+            if sys.version_info < (3, 10):
+                self._init()
+            return self
 
     def init(self, *, scm=False, dvc=False, subdir=False):
         from dvc.repo import Repo

dvc/ui/_rich_progress.py

@@ -4,6 +4,7 @@ from rich.progress import (
     DownloadColumn,
     MofNCompleteColumn,
     Progress,
+    SpinnerColumn,
     TextColumn,
     TimeElapsedColumn,
     TimeRemainingColumn,
@@ -52,3 +53,24 @@ class RichTransferProgress(RichProgress):
         yield self.make_tasks_table(summary_tasks)
         self.columns = self.TRANSFER_COLS
         yield self.make_tasks_table(other_tasks)
+
+
+class DbDownloadProgress(RichProgress):
+    PROGRESS_COLS = (
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TextColumn("[progress.download]{task.completed:,} rows"),
+        TextColumn("to [repr.filename]{task.fields[output]}"),
+    )
+    STATUS_COLS = (
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TextColumn("to [repr.filename]{task.fields[output]}"),
+    )
+
+    def get_renderables(self):
+        if self.tasks:
+            (task, *_) = self.tasks
+            cols = self.PROGRESS_COLS if task.completed else self.STATUS_COLS
+            self.columns = cols[1:] if task.finished else cols
+        yield self.make_tasks_table(self.tasks)

dvc/utils/packaging.py

@@ -0,0 +1,32 @@
+import logging
+
+from funcy import once_per_args
+
+from dvc.log import logger
+
+logger = logger.getChild(__name__)
+
+
+@once_per_args
+def check_required_version(pkg: str, dist: str = "dvc", log_level=logging.WARNING):
+    from importlib import metadata
+
+    from packaging.requirements import InvalidRequirement, Requirement
+
+    try:
+        reqs = {
+            r.name: r.specifier for r in map(Requirement, metadata.requires(dist) or [])
+        }
+        version = metadata.version(pkg)
+    except (metadata.PackageNotFoundError, InvalidRequirement):
+        return
+
+    specifier = reqs.get(pkg)
+    if specifier and version and version not in specifier:
+        logger.log(
+            log_level,
+            "%s%s is required, but you have %r installed which is incompatible.",
+            pkg,
+            specifier,
+            version,
+        )

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=2.21.2,<2.22.0",
+    "dvc-data>=2.22.1,<2.23.0",
     "dvc-http>=2.29.0",
     "dvc-render>=0.3.1,<1",
     "dvc-studio-client>=0.13.0,<1",
@@ -58,7 +58,7 @@ dependencies = [
     "requests>=2.22",
     "rich>=12",
     "ruamel.yaml>=0.17.11",
-    "scmrepo>=1.4.1,<2",
+    "scmrepo>=1.5.0,<2",
     "shortuuid>=0.5",
     "shtab<2,>=1.3.4",
     "tabulate>=0.8.7",
@@ -77,6 +77,7 @@ gs = ["dvc-gs==2.22.1"]
 hdfs = ["dvc-hdfs==2.19"]
 lint = [
     "mypy==1.7.0",
+    "pandas-stubs",
     "types-colorama",
     "types-psutil",
     "types-pyinstaller",
@@ -90,9 +91,9 @@ s3 = ["dvc-s3==2.23.0"]
 ssh = ["dvc-ssh>=2.22.1,<3"]
 ssh_gssapi = ["dvc-ssh[gssapi]>=2.22.1,<3"]
 testing = [
-  "pytest-test-utils",
   "pytest-benchmark[histogram]",
-  "pytest-virtualenv",
+  "pytest-test-utils",
+  "virtualenv",
 ]
 tests = [
     "dvc[testing]",
@@ -110,6 +111,9 @@ 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",
 ]
 webdav = ["dvc-webdav==2.19.1"]
 webhdfs = ["dvc-webhdfs==2.19"]
@@ -126,9 +130,6 @@ dvc = "dvc.cli:main"
 [project.entry-points."fsspec.specs"]
 dvc = "dvc.api:DVCFileSystem"
 
-[project.entry-points."pytest11"]
-dvc-testing = "dvc.testing.plugin"
-
 [project.entry-points."pyinstaller40"]
 hook-dirs = "dvc.__pyinstaller:get_hook_dirs"
 tests = "dvc.__pyinstaller:get_PyInstaller_tests"
@@ -144,7 +145,7 @@ namespaces = false
 write_to = "dvc/_dvc_version.py"
 
 [tool.pytest.ini_options]
-addopts = "-ra --cov-config pyproject.toml --dist worksteal"
+addopts = "-ra --cov-config pyproject.toml --dist worksteal --benchmark-skip"
 filterwarnings = [
     "error::ResourceWarning",
     "error::pytest.PytestUnraisableExceptionWarning",
@@ -155,8 +156,6 @@ filterwarnings = [
     # google.logging: https://github.com/googleapis/python-logging/issues/730
     # Also happens with `zc.lockfile`.
     "ignore:Deprecated call to `pkg_resources.declare_namespace:DeprecationWarning",
-    # tpi imports pkg_resources
-    "ignore:pkg_resources is deprecated as an API:DeprecationWarning",
     # see https://github.com/networkx/networkx/issues/5723.
     "ignore:nx.nx_pydot.* depends on the pydot package, which has.*known issues and is not actively maintained:DeprecationWarning",
 ]
@@ -207,6 +206,7 @@ warn_unused_configs = true
 [[tool.mypy.overrides]]
 ignore_missing_imports = true
 module = [
+    "agate.*",
     "celery.*",
     "configobj.*",
     "dpath.*",
@@ -228,12 +228,11 @@ module = [
     "pyparsing",
     "pytest_benchmark.*",
     "pytest_docker.plugin",
-    "pytest_virtualenv.*",
     "ruamel.*",
     "ruamel.yaml.*",
     "shortuuid",
     "shtab",
-    "tpi.*",
+    "virtualenv",
     "viztracer",
     "voluptuous",
     "yappi",

tests/func/test_cli.py

@@ -96,7 +96,7 @@ def test_config_unset(dvc):
     cmd = args.func(args)
     assert isinstance(cmd, CmdConfig)
     assert args.unset
-    assert args.name == (False, "section", "option")
+    assert args.name == (None, "section", "option")
     assert args.value == value
 
 

tests/func/test_import_db.py

@@ -0,0 +1,128 @@
+"""import_db tests without involving any databases."""
+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)
+
+
+@pytest.fixture
+def adapter(mocker):
+    m = mocker.patch("dvc.database.get_adapter")
+    adapter = mocker.MagicMock()
+    adapter.query.return_value.__enter__.side_effect = serializers
+    m.return_value.__enter__.return_value = adapter
+    return m
+
+
+def test_sql(adapter, tmp_dir, dvc):
+    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_sql_conn_string(adapter, tmp_dir, dvc):
+    with dvc.config.edit(level="local") as conf:
+        conf["db"] = {"conn": {"url": "conn"}}
+
+    stage = dvc.imp_db(
+        sql="select * from model", connection="conn", output_format="json"
+    )
+    db = {"file_format": "json", "connection": "conn", "query": "select * from model"}
+    assert (tmp_dir / stage.relpath).parse() == {
+        "md5": "f0dcca8e8411907adb25aa547d7c432c",
+        "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": "8c0441b206d8ee7c7b9e3d54871c3cbb",
+        "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.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,
+            }
+        ],
+    }

tests/func/test_import_db_e2e.py

@@ -0,0 +1,142 @@
+"""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"))
+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}
+
+    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
+    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)]

tests/func/test_scm.py

@@ -3,7 +3,7 @@ import os
 import pytest
 from git import Repo
 
-from dvc.scm import SCM, Git, NoSCM, SCMError
+from dvc.scm import SCM, Git, NoSCM, SCMError, lfs_prefetch
 
 
 def test_init_none(tmp_dir):
@@ -27,3 +27,20 @@ def test_init_sub_dir(tmp_dir):
 
     scm = SCM(os.fspath(subdir))
     assert scm.root_dir == os.fspath(tmp_dir)
+
+
+def test_lfs_prefetch(tmp_dir, dvc, scm, mocker):
+    mock_fetch = mocker.patch("scmrepo.git.lfs.fetch")
+    rev = scm.get_rev()
+
+    with dvc.switch(rev):
+        lfs_prefetch(dvc.dvcfs, ["foo"])
+        mock_fetch.assert_not_called()
+
+    tmp_dir.scm_gen(
+        ".gitattributes", ".lfs filter=lfs diff=lfs merge=lfs -text", commit="init lfs"
+    )
+    rev = scm.get_rev()
+    with dvc.switch(rev):
+        lfs_prefetch(dvc.dvcfs, ["foo"])
+        mock_fetch.assert_called_once()

tests/unit/stage/test_stage.py

@@ -69,7 +69,7 @@ def test_stage_update(dvc, mocker):
     is_repo_import.return_value = True
     with dvc.lock:
         stage.update()
-    reproduce.assert_called_once_with(no_download=None, jobs=None)
+    reproduce.assert_called_once_with(no_download=None, jobs=None, force=False)
 
     is_repo_import.return_value = False
     with pytest.raises(StageUpdateError):

tests/unit/test_dirs.py

@@ -0,0 +1,8 @@
+from dvc.dirs import global_config_dir
+from dvc.env import DVC_GLOBAL_CONFIG_DIR
+
+
+def test_global_config_dir_respects_env_var(monkeypatch):
+    path = "/some/random/path"
+    monkeypatch.setenv(DVC_GLOBAL_CONFIG_DIR, path)
+    assert global_config_dir() == path