Repository: iterative/dvc
Pull request #10219 (2024-01-04 06:42:30)

import-db: support importing a table
I have also made `import-db` command visible in the help message and hidden the dbt-specific options.

```console
$ dvc import-db --table customer --conn pgsql
```

(Note that I haven't added support for downloading a table with dbt's connection profile).

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) Report
Attention: `19 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`d218b6a`)](https://app.codecov.io/gh/iterative/dvc/commit/d218b6ac2be82fe4c915cf2b2ed726ddb7467880?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) 90.50% compared to head [(`736b66f`)](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) 90.07%.
> Report is 3 commits behind head on main.

| [Files](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) | Patch % | Lines |
|---|---|---|
| [dvc/commands/imp\_db.py](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#diff-ZHZjL2NvbW1hbmRzL2ltcF9kYi5weQ==) | 37.50% | [15 Missing :warning: ](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) |
| [dvc/dependency/db.py](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#diff-ZHZjL2RlcGVuZGVuY3kvZGIucHk=) | 80.00% | [2 Missing and 2 partials :warning: ](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main   #10219      +/-   ##
==========================================
- Coverage   90.50%   90.07%   -0.44%     
==========================================
  Files         499      499              
  Lines       37982    38042      +60     
  Branches     5518     5521       +3     
==========================================
- Hits        34375    34265     -110     
- Misses       2965     3087     +122     
- Partials      642      690      +48     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iterative/dvc/pull/10219?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).

dberenbaum: Could be a separate PR, but should we drop the dbt part at this point?

Patch:
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="<commit>",
     )
     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/dvc/repo/index.py b/dvc/repo/index.py
index ee1d6a798..35eea98f2 100644
--- a/dvc/repo/index.py
+++ b/dvc/repo/index.py
@@ -255,6 +255,33 @@ def _load_storage_from_out(storage_map, key, out):
         _load_storage_from_import(storage_map, key, out)
 
 
+def _build_tree_from_outs(outs):
+    from dvc_data.hashfile.tree import Tree
+
+    tree = Tree()
+    for out in outs:
+        if not out.use_cache:
+            continue
+
+        ws, key = out.index_key
+
+        if not out.stage.is_partial_import:
+            tree.add((ws, *key), out.meta, out.hash_info)
+            continue
+
+        dep = out.stage.deps[0]
+        if not dep.files:
+            tree.add((ws, *key), dep.meta, dep.hash_info)
+            continue
+
+        for okey, ometa, ohi in dep.get_obj():
+            tree.add((ws, *key, *okey), ometa, ohi)
+
+    tree.digest()
+
+    return tree
+
+
 class Index:
     def __init__(
         self,
@@ -504,20 +531,7 @@ class Index:
 
     @cached_property
     def data_tree(self):
-        from dvc_data.hashfile.tree import Tree
-
-        tree = Tree()
-        for out in self.outs:
-            if not out.use_cache:
-                continue
-
-            ws, key = out.index_key
-
-            tree.add((ws, *key), out.meta, out.hash_info)
-
-        tree.digest()
-
-        return tree
+        return _build_tree_from_outs(self.outs)
 
     @cached_property
     def data(self) -> "Dict[str, DataIndex]":
@@ -772,20 +786,7 @@ class IndexView:
 
     @cached_property
     def data_tree(self):
-        from dvc_data.hashfile.tree import Tree
-
-        tree = Tree()
-        for out in self.outs:
-            if not out.use_cache:
-                continue
-
-            ws, key = out.index_key
-
-            tree.add((ws, *key), out.meta, out.hash_info)
-
-        tree.digest()
-
-        return tree
+        return _build_tree_from_outs(self.outs)
 
     @cached_property
     def data(self) -> Dict[str, Union["DataIndex", "DataIndexView"]]:
diff --git a/pyproject.toml b/pyproject.toml
index 5bf04c37b..1eb843141 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,7 +35,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=3.6,<3.7",
+    "dvc-data>=3.7,<3.8",
     "dvc-http>=2.29.0",
     "dvc-render>=1.0.0,<2",
     "dvc-studio-client>=0.17.1,<1",


Test patch:
diff --git a/dvc/testing/workspace_tests.py b/dvc/testing/workspace_tests.py
index c9846873c..13bdaf7d3 100644
--- a/dvc/testing/workspace_tests.py
+++ b/dvc/testing/workspace_tests.py
@@ -149,14 +149,33 @@ class TestImportURLVersionAware:
         assert (tmp_dir / "data_dir" / "subdir" / "file").read_text() == "modified"
         assert (tmp_dir / "data_dir" / "new_file").read_text() == "new"
 
-    def test_import_no_download(self, tmp_dir, dvc, remote_version_aware):
+    def test_import_no_download(self, tmp_dir, dvc, remote_version_aware, scm):
         remote_version_aware.gen({"data_dir": {"subdir": {"file": "file"}}})
         dvc.imp_url("remote://upstream/data_dir", version_aware=True, no_download=True)
+        scm.add(["data_dir.dvc", ".gitignore"])
+        scm.commit("v1")
+        scm.tag("v1")
+
         stage = first(dvc.index.stages)
         assert not stage.outs[0].can_push
 
-        dvc.pull()
-        assert (tmp_dir / "data_dir" / "subdir" / "file").read_text() == "file"
+        (remote_version_aware / "data_dir" / "foo").write_text("foo")
+        dvc.update(no_download=True)
+        assert dvc.pull()["fetched"] == 2
+        assert (tmp_dir / "data_dir").read_text() == {
+            "foo": "foo",
+            "subdir": {"file": "file"},
+        }
+        scm.add(["data_dir.dvc", ".gitignore"])
+        scm.commit("update")
+
+        scm.checkout("v1")
+        dvc.cache.local.clear()
+        remove(tmp_dir / "data_dir")
+        assert dvc.pull()["fetched"] == 1
+        assert (tmp_dir / "data_dir").read_text() == {
+            "subdir": {"file": "file"},
+        }
 
         dvc.commit(force=True)
         assert dvc.status() == {}
diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py
index 17f15f70a..5e58edf53 100644
--- a/tests/func/test_data_cloud.py
+++ b/tests/func/test_data_cloud.py
@@ -271,7 +271,7 @@ def test_pull_partial_import(tmp_dir, dvc, local_workspace):
     stage = dvc.imp_url("remote://workspace/file", os.fspath(dst), no_download=True)
 
     result = dvc.pull("file")
-    assert result["fetched"] == 0
+    assert result["fetched"] == 1
     assert dst.exists()
 
     assert stage.outs[0].get_hash().value == "d10b4c3ff123b26dc068d43a8bef2d23"
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)