config: azure: don't forget Coerce() for timeouts (#10247)

Co-authored-by: Dave Berenbaum <[email protected]>

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
       - id: sort-simple-yaml
       - id: trailing-whitespace
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.1.8'
+    rev: 'v0.1.11'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

dvc/cli/formatter.py

@@ -0,0 +1,14 @@
+import argparse
+
+
+class HelpFormatter(argparse.HelpFormatter):
+    def _get_default_metavar_for_optional(self, action: argparse.Action) -> str:
+        return action.dest
+
+
+class RawTextHelpFormatter(HelpFormatter, argparse.RawTextHelpFormatter):
+    pass
+
+
+class RawDescriptionHelpFormatter(HelpFormatter, argparse.RawDescriptionHelpFormatter):
+    pass

dvc/cli/parser.py

@@ -50,7 +50,7 @@ from dvc.commands import (
 )
 from dvc.log import logger
 
-from . import DvcParserError
+from . import DvcParserError, formatter
 
 logger = logger.getChild(__name__)
 
@@ -165,7 +165,7 @@ def get_main_parser():
         prog="dvc",
         description=desc,
         parents=[parent_parser],
-        formatter_class=argparse.RawTextHelpFormatter,
+        formatter_class=formatter.RawTextHelpFormatter,
         add_help=False,
     )
 
@@ -200,15 +200,12 @@ def get_main_parser():
     # Sub commands
     subparsers = parser.add_subparsers(
         title="Available Commands",
-        metavar="COMMAND",
+        metavar="command",
         dest="cmd",
-        help="Use `dvc COMMAND --help` for command-specific help.",
+        help="Use `dvc command --help` for command-specific help.",
+        required=True,
     )
 
-    from .utils import fix_subparsers
-
-    fix_subparsers(subparsers)
-
     for cmd in COMMANDS:
         cmd.add_parser(subparsers, parent_parser)
 

dvc/cli/utils.py

@@ -26,18 +26,6 @@ class DictAction(argparse.Action):
         setattr(args, self.dest, d)
 
 
-def fix_subparsers(subparsers):
-    """Workaround for bug in Python 3. See more info at:
-    https://bugs.python.org/issue16308
-    https://github.com/iterative/dvc/issues/769
-
-    Args:
-        subparsers: subparsers to fix.
-    """
-    subparsers.required = True
-    subparsers.dest = "cmd"
-
-
 def append_doc_link(help_message, path):
     from dvc.utils import format_link
 

dvc/commands/add.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -71,7 +69,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(ADD_HELP, "add"),
         help=ADD_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     parser.add_argument(
         "--no-commit",

dvc/commands/artifacts.py

@@ -1,8 +1,6 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBaseNoRepo
-from dvc.cli.utils import DictAction, append_doc_link, fix_subparsers
+from dvc.cli.utils import DictAction, append_doc_link
 from dvc.exceptions import DvcException
 from dvc.log import logger
 
@@ -47,13 +45,13 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(ARTIFACTS_HELP, "artifacts"),
         help=ARTIFACTS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     artifacts_subparsers = artifacts_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc artifacts CMD --help` to display command-specific help.",
+        required=True,
     )
-    fix_subparsers(artifacts_subparsers)
 
     ARTIFACTS_GET_HELP = "Download an artifact from a DVC project."
     get_parser = artifacts_subparsers.add_parser(
@@ -61,7 +59,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(ARTIFACTS_GET_HELP, "artifacts/get"),
         help=ARTIFACTS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     get_parser.add_argument("url", help="Location of DVC repository to download from")
     get_parser.add_argument(

dvc/commands/cache.py

@@ -1,9 +1,9 @@
 import argparse
 import os
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli.utils import append_doc_link
 from dvc.commands.config import CmdConfig
 from dvc.ui import ui
 
@@ -58,12 +58,13 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(CACHE_HELP, "cache"),
         help=CACHE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     cache_subparsers = cache_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc cache CMD --help` for command-specific help.",
+        required=True,
     )
 
     parent_cache_config_parser = argparse.ArgumentParser(
@@ -76,7 +77,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser, parent_cache_config_parser],
         description=append_doc_link(CACHE_HELP, "cache/dir"),
         help=CACHE_DIR_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     cache_dir_parser.add_argument(
         "-u",
@@ -103,7 +104,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(CACHE_HELP, "cache/migrate"),
         help=CACHE_MIGRATE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     cache_migrate_parser.add_argument(
         "--dvc-files",
@@ -122,5 +123,3 @@ def add_parser(subparsers, parent_parser):
         action="store_true",
     )
     cache_migrate_parser.set_defaults(func=CmdCacheMigrate)
-
-    fix_subparsers(cache_subparsers)

dvc/commands/check_ignore.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.ui import ui
@@ -85,7 +83,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(ADD_HELP, "check-ignore"),
         help=ADD_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     parser.add_argument(
         "-d",

dvc/commands/checkout.py

@@ -1,7 +1,6 @@
-import argparse
 import operator
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import CheckoutError
@@ -67,7 +66,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(CHECKOUT_HELP, "checkout"),
         help=CHECKOUT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     checkout_parser.add_argument(
         "--summary",

dvc/commands/commit.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -40,7 +38,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(COMMIT_HELP, "commit"),
         help=COMMIT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     commit_parser.add_argument(
         "-f",

dvc/commands/completion.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.completion import PREAMBLE
 from dvc.cli.utils import append_doc_link
@@ -31,7 +30,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(COMPLETION_DESCRIPTION, "completion"),
         help=COMPLETION_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     completion_parser.add_argument(
         "-s",

dvc/commands/config.py

@@ -3,6 +3,7 @@ import os
 
 from funcy import set_in
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -209,7 +210,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(CONFIG_HELP, "config"),
         help=CONFIG_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     config_parser.add_argument(
         "-u",

dvc/commands/daemon.py

@@ -1,6 +1,5 @@
 from dvc.cli import completion
 from dvc.cli.command import CmdBaseNoRepo
-from dvc.cli.utils import fix_subparsers
 from dvc.log import logger
 
 logger = logger.getChild(__name__)
@@ -53,10 +52,9 @@ def add_parser(subparsers, parent_parser):
     daemon_subparsers = daemon_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc daemon CMD --help` for command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(daemon_subparsers)
-
     DAEMON_UPDATER_HELP = "Fetch latest available version."
     daemon_updater_parser = daemon_subparsers.add_parser(
         "updater",

dvc/commands/dag.py

@@ -1,6 +1,6 @@
-import argparse
 from typing import TYPE_CHECKING
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.ui import ui
@@ -160,7 +160,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(DAG_HELP, "dag"),
         help=DAG_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     dag_parser.add_argument(
         "--dot",

dvc/commands/data.py

@@ -1,10 +1,10 @@
-import argparse
 from typing import TYPE_CHECKING
 
 from funcy import chunks, compact, log_durations
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli.utils import append_doc_link
 from dvc.log import logger
 from dvc.ui import ui
 from dvc.utils import colorize
@@ -129,13 +129,13 @@ def add_parser(subparsers, parent_parser):
     data_parser = subparsers.add_parser(
         "data",
         parents=[parent_parser],
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     data_subparsers = data_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc data CMD --help` to display command-specific help.",
+        required=True,
     )
-    fix_subparsers(data_subparsers)
 
     DATA_STATUS_HELP = (
         "Show changes between the last git commit, the dvcfiles and the workspace."
@@ -144,7 +144,7 @@ def add_parser(subparsers, parent_parser):
         "status",
         parents=[parent_parser],
         description=append_doc_link(DATA_STATUS_HELP, "data/status"),
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
         help=DATA_STATUS_HELP,
     )
     data_status_parser.add_argument(

dvc/commands/data_sync.py

@@ -1,6 +1,6 @@
 import argparse
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -145,7 +145,7 @@ def add_parser(subparsers, _parent_parser):
         parents=[shared_parent_parser()],
         description=append_doc_link(PULL_HELP, "pull"),
         help=PULL_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     pull_parser.add_argument(
         "-r", "--remote", help="Remote storage to pull from", metavar="<name>"
@@ -220,7 +220,7 @@ def add_parser(subparsers, _parent_parser):
         parents=[shared_parent_parser()],
         description=append_doc_link(PUSH_HELP, "push"),
         help=PUSH_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     push_parser.add_argument(
         "-r", "--remote", help="Remote storage to push to", metavar="<name>"
@@ -282,7 +282,7 @@ def add_parser(subparsers, _parent_parser):
         parents=[shared_parent_parser()],
         description=append_doc_link(FETCH_HELP, "fetch"),
         help=FETCH_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     fetch_parser.add_argument(
         "-r", "--remote", help="Remote storage to fetch from", metavar="<name>"
@@ -355,7 +355,7 @@ def add_parser(subparsers, _parent_parser):
         description=append_doc_link(STATUS_HELP, "status"),
         help=STATUS_HELP,
         conflict_handler="resolve",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     status_parser.add_argument(
         "-q",

dvc/commands/destroy.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -41,7 +40,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(DESTROY_HELP, "destroy"),
         help=DESTROY_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     destroy_parser.add_argument(
         "-f",

dvc/commands/diff.py

@@ -1,7 +1,6 @@
-import argparse
 import os
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -163,7 +162,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(DIFF_DESCRIPTION, "diff"),
         help=DIFF_DESCRIPTION,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     diff_parser.add_argument(
         "--targets",

dvc/commands/du.py

@@ -1,7 +1,6 @@
-import argparse
 import logging
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.ui import ui
@@ -34,7 +33,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(DU_HELP, "du"),
         help=DU_HELP,
-        formatter_class=argparse.RawTextHelpFormatter,
+        formatter_class=formatter.RawTextHelpFormatter,
     )
     du_parser.add_argument("url", help="Location of DVC repository")
     du_parser.add_argument(

dvc/commands/experiments/__init__.py

@@ -1,6 +1,5 @@
-import argparse
-
-from dvc.cli.utils import append_doc_link, fix_subparsers, hide_subparsers_from_help
+from dvc.cli import formatter
+from dvc.cli.utils import append_doc_link, hide_subparsers_from_help
 from dvc.commands.experiments import (
     apply,
     branch,
@@ -44,16 +43,16 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         aliases=["exp"],
         description=append_doc_link(EXPERIMENTS_HELP, "exp"),
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
         help=EXPERIMENTS_HELP,
     )
 
     experiments_subparsers = experiments_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc experiments CMD --help` to display command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(experiments_subparsers)
     for cmd in SUB_COMMANDS:
         cmd.add_parser(experiments_subparsers, parent_parser)
     hide_subparsers_from_help(experiments_subparsers)

dvc/commands/experiments/apply.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -30,7 +28,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_APPLY_HELP, "exp/apply"),
         help=EXPERIMENTS_APPLY_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     experiments_apply_parser.add_argument(
         "--no-force",

dvc/commands/experiments/branch.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -21,7 +20,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_BRANCH_HELP, "exp/branch"),
         help=EXPERIMENTS_BRANCH_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     experiments_branch_parser.add_argument(
         "experiment", help="Experiment to be promoted."

dvc/commands/experiments/clean.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -20,6 +19,6 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_CLEAN_HELP, "exp/clean"),
         help=EXPERIMENTS_CLEAN_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     experiments_clean_parser.set_defaults(func=CmdExperimentsClean)

dvc/commands/experiments/diff.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.commands.metrics import DEFAULT_PRECISION
@@ -59,7 +57,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_DIFF_HELP, "exp/diff"),
         help=EXPERIMENTS_DIFF_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     experiments_diff_parser.add_argument(
         "a_rev", nargs="?", help="Old experiment to compare (defaults to HEAD)"

dvc/commands/experiments/ls.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import InvalidArgumentError
@@ -63,7 +62,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_LIST_HELP, "exp/list"),
         help=EXPERIMENTS_LIST_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     add_rev_selection_flags(experiments_list_parser, "List")
     display_group = experiments_list_parser.add_mutually_exclusive_group()

dvc/commands/experiments/pull.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -49,7 +48,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_PULL_HELP, "exp/pull"),
         help=EXPERIMENTS_PULL_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     add_rev_selection_flags(experiments_pull_parser, "Pull", True)
     experiments_pull_parser.add_argument(

dvc/commands/experiments/push.py

@@ -1,7 +1,6 @@
-import argparse
 from typing import Any, Dict
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -90,7 +89,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_PUSH_HELP, "exp/push"),
         help=EXPERIMENTS_PUSH_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     add_rev_selection_flags(experiments_push_parser, "Push", True)
     experiments_push_parser.add_argument(

dvc/commands/experiments/remove.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import InvalidArgumentError
@@ -54,7 +53,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_REMOVE_HELP, "exp/remove"),
         help=EXPERIMENTS_REMOVE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remove_group = experiments_remove_parser.add_mutually_exclusive_group()
     add_rev_selection_flags(experiments_remove_parser, "Remove", False)

dvc/commands/experiments/rename.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import InvalidArgumentError
@@ -38,7 +37,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_RENAME_HELP, "exp/rename"),
         help=EXPERIMENTS_RENAME_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     rename_group = experiments_rename_parser.add_mutually_exclusive_group()
     rename_group.add_argument(

dvc/commands/experiments/run.py

@@ -1,5 +1,6 @@
 import argparse
 
+from dvc.cli import formatter
 from dvc.cli.utils import append_doc_link
 from dvc.commands.repro import CmdRepro
 from dvc.commands.repro import add_arguments as add_repro_arguments
@@ -32,7 +33,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_RUN_HELP, "exp/run"),
         help=EXPERIMENTS_RUN_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     _add_run_common(experiments_run_parser)
     experiments_run_parser.set_defaults(func=CmdExperimentsRun)

dvc/commands/experiments/save.py

@@ -1,5 +1,6 @@
 import argparse
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -38,7 +39,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_SAVE_HELP, "exp/save"),
         help=EXPERIMENTS_SAVE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     save_parser.add_argument(
         "-f",

dvc/commands/experiments/show.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Dict, Iterable
 
 from funcy import lmap
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.commands.metrics import DEFAULT_PRECISION
@@ -200,7 +201,7 @@ def add_parser(experiments_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(EXPERIMENTS_SHOW_HELP, "exp/show"),
         help=EXPERIMENTS_SHOW_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     add_rev_selection_flags(experiments_show_parser, "Show")
     experiments_show_parser.add_argument(

dvc/commands/freeze.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -38,7 +36,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(FREEZE_HELP, "freeze"),
         help=FREEZE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     freeze_parser.add_argument(
         "targets", nargs="+", help="Stages or .dvc files to freeze"
@@ -51,7 +49,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(UNFREEZE_HELP, "unfreeze"),
         help=UNFREEZE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     unfreeze_parser.add_argument(
         "targets", nargs="+", help="Stages or .dvc files to unfreeze"

dvc/commands/gc.py

@@ -1,6 +1,6 @@
-import argparse
 import os
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -102,7 +102,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(GC_DESCRIPTION, "gc"),
         help=GC_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     gc_parser.add_argument(
         "-w",
@@ -156,12 +156,12 @@ def add_parser(subparsers, parent_parser):
         "--date",
         type=str,
         dest="commit_date",
-        metavar="<YYYY-MM-DD>",
+        metavar="<yyyy-mm-dd>",
         default=None,
         help=(
             "Keep cached data referenced in the commits after ( inclusive )"
             " a certain time. Date must match the extended ISO 8601 format "
-            "(YYYY-MM-DD)."
+            "(yyyy-mm-dd)."
         ),
     )
     gc_parser.add_argument(

dvc/commands/get.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.exceptions import DvcException
@@ -59,7 +57,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(GET_HELP, "get"),
         help=GET_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     get_parser.add_argument(
         "url", help="Location of DVC or Git repository to download from"

dvc/commands/get_url.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.exceptions import DvcException
@@ -36,7 +34,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(GET_HELP, "get-url"),
         help=GET_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     get_parser.add_argument(
         "url", help="See `dvc import-url -h` for full list of supported URLs."

dvc/commands/git_hook.py

@@ -1,7 +1,7 @@
 import os
 
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
-from dvc.cli.utils import fix_subparsers
 from dvc.exceptions import NotDvcRepoError
 from dvc.log import logger
 
@@ -94,10 +94,9 @@ def add_parser(subparsers, parent_parser):
     git_hook_subparsers = git_hook_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc daemon CMD --help` for command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(git_hook_subparsers)
-
     PRE_COMMIT_HELP = "Run pre-commit GIT hook."
     pre_commit_parser = git_hook_subparsers.add_parser(
         "pre-commit",
@@ -140,6 +139,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=MERGE_DRIVER_HELP,
         help=MERGE_DRIVER_HELP,
+        formatter_class=formatter.HelpFormatter,
     )
     merge_driver_parser.add_argument(
         "--ancestor",

dvc/commands/imp.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.exceptions import DvcException
@@ -51,7 +49,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(IMPORT_HELP, "import"),
         help=IMPORT_HELP,
-        formatter_class=argparse.RawTextHelpFormatter,
+        formatter_class=formatter.RawTextHelpFormatter,
     )
     import_parser.add_argument(
         "url", help="Location of DVC or Git repository to download from"

dvc/commands/imp_db.py

@@ -1,9 +1,5 @@
-import argparse
-
-from funcy import compact, merge
-
-from dvc.cli import completion
-from dvc.cli.command import CmdBase
+from dvc.cli import completion, formatter
+from dvc.cli.command import CmdBase, CmdBaseNoRepo
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
 from dvc.ui import ui
@@ -11,115 +7,55 @@ from dvc.ui import ui
 logger = logger.getChild(__name__)
 
 
-class CmdTestDb(CmdBase):
+class CmdTestDb(CmdBaseNoRepo):
     def run(self):
-        from dvc.database import get_client
-        from dvc.database.dbt_utils import DBT_PROJECT_FILE, is_dbt_project
-        from dvc.dependency.db import _get_dbt_config
+        from dvc.config import Config
+        from dvc.database import client
         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:
+        connection = self.args.conn
+        db_config = Config.from_cwd().get("db", {})
+        if connection not in db_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_client(conn_config, project_dir=project_dir, **dbt_config)
-        with adapter as db:
+        config = db_config.get(connection, {})
+        if self.args.url:
+            config["url"] = self.args.url
+        if self.args.username:
+            config["username"] = self.args.username
+        if self.args.password:
+            config["password"] = self.args.password
+        with client(config) 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,
+            table=self.args.table,
             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=formatter.RawTextHelpFormatter,
     )
-    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")
+    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")
     import_parser.add_argument(
         "--output-format",
         default="csv",
@@ -144,8 +80,7 @@ def add_parser(subparsers, parent_parser):
     )
     import_parser.add_argument(
         "--conn",
-        dest="connection",
-        nargs="?",
+        required=True,
         help="Database connection to use, needs to be set in config",
     )
 
@@ -158,10 +93,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", required=True)
     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/commands/imp_url.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.exceptions import DvcException
@@ -44,7 +42,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(IMPORT_HELP, "import-url"),
         help=IMPORT_HELP,
-        formatter_class=argparse.RawTextHelpFormatter,
+        formatter_class=formatter.RawTextHelpFormatter,
     )
     import_parser.add_argument(
         "url",

dvc/commands/init.py

@@ -1,8 +1,7 @@
-import argparse
-
 import colorama
 
 from dvc import analytics
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -69,7 +68,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(INIT_DESCRIPTION, "init"),
         help=INIT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     init_parser.add_argument(
         "--no-scm",

dvc/commands/install.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -25,7 +24,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(INSTALL_HELP, "install"),
         help=INSTALL_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     install_parser.add_argument(
         "--use-pre-commit-tool",

dvc/commands/ls/__init__.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.commands.ls.ls_colors import LsColors
@@ -75,7 +73,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(LIST_HELP, "list"),
         help=LIST_HELP,
-        formatter_class=argparse.RawTextHelpFormatter,
+        formatter_class=formatter.RawTextHelpFormatter,
     )
     list_parser.add_argument("url", help="Location of DVC repository to list")
     list_parser.add_argument(

dvc/commands/ls_url.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import DictAction, append_doc_link
 from dvc.log import logger
@@ -33,7 +32,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(LS_HELP, "list-url"),
         help=LS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     lsurl_parser.add_argument(
         "url", help="See `dvc import-url -h` for full list of supported URLs."

dvc/commands/metrics.py

@@ -1,8 +1,6 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli.utils import append_doc_link
 from dvc.log import logger
 from dvc.ui import ui
 from dvc.utils.serialize import encode_exception
@@ -107,23 +105,22 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(METRICS_HELP, "metrics"),
         help=METRICS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     metrics_subparsers = metrics_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc metrics CMD --help` to display command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(metrics_subparsers)
-
     METRICS_SHOW_HELP = "Print metrics, with optional formatting."
     metrics_show_parser = metrics_subparsers.add_parser(
         "show",
         parents=[parent_parser],
         description=append_doc_link(METRICS_SHOW_HELP, "metrics/show"),
         help=METRICS_SHOW_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     metrics_show_parser.add_argument(
         "targets",
@@ -197,7 +194,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(METRICS_DIFF_HELP, "metrics/diff"),
         help=METRICS_DIFF_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     metrics_diff_parser.add_argument(
         "a_rev",

dvc/commands/move.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -34,7 +32,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(MOVE_DESCRIPTION, "move"),
         help=MOVE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     move_parser.add_argument(
         "src", help="Source path to a data file or directory."

dvc/commands/params.py

@@ -1,8 +1,6 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli.utils import append_doc_link
 from dvc.log import logger
 from dvc.ui import ui
 
@@ -61,16 +59,15 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(PARAMS_HELP, "params"),
         help=PARAMS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     params_subparsers = params_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc params CMD --help` to display command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(params_subparsers)
-
     PARAMS_DIFF_HELP = (
         "Show changes in params between commits in the DVC repository, or "
         "between a commit and the workspace."
@@ -80,7 +77,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(PARAMS_DIFF_HELP, "params/diff"),
         help=PARAMS_DIFF_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     params_diff_parser.add_argument(
         "a_rev",

dvc/commands/plots.py

@@ -4,9 +4,9 @@ from typing import TYPE_CHECKING, Dict, List, Optional
 
 from funcy import compact, first, get_in
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
 from dvc.log import logger
 from dvc.ui import ui
@@ -223,15 +223,14 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(PLOTS_HELP, "plots"),
         help=PLOTS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     plots_subparsers = plots_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc plots CMD --help` to display command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(plots_subparsers)
-
     SHOW_HELP = (
         "Generate plots from target files or from `plots` definitions in `dvc.yaml`."
     )
@@ -240,7 +239,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(SHOW_HELP, "plots/show"),
         help=SHOW_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     plots_show_parser.add_argument(
         "targets",
@@ -263,7 +262,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(PLOTS_DIFF_HELP, "plots/diff"),
         help=PLOTS_DIFF_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     plots_diff_parser.add_argument(
         "--targets",
@@ -299,7 +298,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(PLOTS_MODIFY_HELP, "plots/modify"),
         help=PLOTS_MODIFY_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     plots_modify_parser.add_argument(
         "target",
@@ -320,7 +319,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(TEMPLATES_HELP, "plots/templates"),
         help=TEMPLATES_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     plots_templates_parser.add_argument(
         "template",

dvc/commands/queue/__init__.py

@@ -1,6 +1,5 @@
-import argparse
-
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli import formatter
+from dvc.cli.utils import append_doc_link
 from dvc.commands.queue import kill, logs, remove, start, status, stop
 
 SUB_COMMANDS = [
@@ -20,15 +19,15 @@ def add_parser(subparsers, parent_parser):
         "queue",
         parents=[parent_parser],
         description=append_doc_link(QUEUE_HELP, "queue"),
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
         help=QUEUE_HELP,
     )
 
     queue_subparsers = queue_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc queue CMD --help` to display command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(queue_subparsers)
     for cmd in SUB_COMMANDS:
         cmd.add_parser(queue_subparsers, parent_parser)

dvc/commands/queue/kill.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -27,7 +26,7 @@ def add_parser(queue_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(QUEUE_KILL_HELP, "queue/kill"),
         help=QUEUE_KILL_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     queue_kill_parser.add_argument(
         "-f",

dvc/commands/queue/logs.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -29,7 +28,7 @@ def add_parser(queue_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(QUEUE_LOGS_HELP, "queue/logs"),
         help=QUEUE_LOGS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     queue_logs_parser.add_argument(
         "-e",

dvc/commands/queue/remove.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import InvalidArgumentError
@@ -62,7 +61,7 @@ def add_parser(queue_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(QUEUE_REMOVE_HELP, "queue/remove"),
         help=QUEUE_REMOVE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     queue_remove_parser.add_argument(
         "--all",

dvc/commands/queue/start.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -27,7 +26,7 @@ def add_parser(queue_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(QUEUE_START_HELP, "queue/start"),
         help=QUEUE_START_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     queue_start_parser.add_argument(
         "-j",

dvc/commands/queue/status.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.compare import TabularData
@@ -52,6 +51,6 @@ def add_parser(queue_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(QUEUE_STATUS_HELP, "queue/status"),
         help=QUEUE_STATUS_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     queue_status_parser.set_defaults(func=CmdQueueStatus)

dvc/commands/queue/stop.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -32,7 +31,7 @@ def add_parser(queue_subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(QUEUE_STOP_HELP, "queue/stop"),
         help=QUEUE_STOP_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     queue_stop_parser.add_argument(
         "--kill",

dvc/commands/remote.py

@@ -1,6 +1,5 @@
-import argparse
-
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli import formatter
+from dvc.cli.utils import append_doc_link
 from dvc.commands.config import CmdConfig
 from dvc.ui import ui
 from dvc.utils import format_link
@@ -155,23 +154,22 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(REMOTE_HELP, "remote"),
         help=REMOTE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     remote_subparsers = remote_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc remote CMD --help` for command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(remote_subparsers)
-
     REMOTE_ADD_HELP = "Add a new data remote."
     remote_add_parser = remote_subparsers.add_parser(
         "add",
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(REMOTE_ADD_HELP, "remote/add"),
         help=REMOTE_ADD_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remote_add_parser.add_argument("name", help="Name of the remote")
     remote_add_parser.add_argument(
@@ -202,7 +200,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(REMOTE_DEFAULT_HELP, "remote/default"),
         help=REMOTE_DEFAULT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remote_default_parser.add_argument("name", nargs="?", help="Name of the remote")
     remote_default_parser.add_argument(
@@ -220,7 +218,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(REMOTE_MODIFY_HELP, "remote/modify"),
         help=REMOTE_MODIFY_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remote_modify_parser.add_argument("name", help="Name of the remote")
     remote_modify_parser.add_argument("option", help="Name of the option to modify.")
@@ -242,7 +240,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(REMOTE_LIST_HELP, "remote/list"),
         help=REMOTE_LIST_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remote_list_parser.set_defaults(func=CmdRemoteList)
 
@@ -252,7 +250,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(REMOTE_REMOVE_HELP, "remote/remove"),
         help=REMOTE_REMOVE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remote_remove_parser.add_argument("name", help="Name of the remote to remove.")
     remote_remove_parser.set_defaults(func=CmdRemoteRemove)
@@ -262,7 +260,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_config_parser, parent_parser],
         description=append_doc_link(REMOTE_RENAME_HELP, "remote/rename"),
         help=REMOTE_RENAME_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remote_rename_parser.add_argument("name", help="Remote to be renamed")
     remote_rename_parser.add_argument("new", help="New name of the remote")

dvc/commands/remove.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -30,7 +28,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(REMOVE_HELP, "remove"),
         help=REMOVE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     remove_parser.add_argument(
         "--outs",

dvc/commands/repro.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.commands.status import CmdDataStatus
@@ -167,7 +165,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(REPRO_HELP, "repro"),
         help=REPRO_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     # repro/exp run shared args
     add_arguments(repro_parser)

dvc/commands/root.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -24,6 +23,6 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(ROOT_HELP, "root"),
         help=ROOT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     root_parser.set_defaults(func=CmdRoot)

dvc/commands/stage.py

@@ -4,9 +4,9 @@ from contextlib import contextmanager
 from itertools import chain, filterfalse
 from typing import TYPE_CHECKING, Dict, Iterable, List
 
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli.utils import append_doc_link
 from dvc.log import logger
 from dvc.utils.cli_parse import parse_params
 from dvc.utils.humanize import truncate_text
@@ -284,23 +284,22 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(STAGES_HELP, "stage"),
         help=STAGES_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     stage_subparsers = stage_parser.add_subparsers(
         dest="cmd",
         help="Use `dvc stage CMD --help` to display command-specific help.",
+        required=True,
     )
 
-    fix_subparsers(stage_subparsers)
-
     STAGE_ADD_HELP = "Create stage"
     stage_add_parser = stage_subparsers.add_parser(
         "add",
         parents=[parent_parser],
         description=append_doc_link(STAGE_ADD_HELP, "stage/add"),
         help=STAGE_ADD_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     stage_add_parser.add_argument(
         "-n", "--name", help="Name of the stage to add", required=True
@@ -314,7 +313,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(STAGE_LIST_HELP, "stage/list"),
         help=STAGE_LIST_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     stage_list_parser.add_argument(
         "targets",

dvc/commands/studio.py

@@ -1,9 +1,9 @@
-import argparse
 import os
 
 from funcy import get_in
 
-from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.cli import formatter
+from dvc.cli.utils import append_doc_link
 from dvc.commands.config import CmdConfig
 from dvc.log import logger
 
@@ -92,13 +92,13 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(STUDIO_DESCRIPTION, "studio"),
         help=STUDIO_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     studio_subparser = studio_parser.add_subparsers(
         dest="cmd",
         help="Use `DVC studio CMD --help` to display command-specific help.",
+        required=True,
     )
-    fix_subparsers(studio_subparser)
 
     STUDIO_LOGIN_HELP = "Authenticate DVC with Studio host"
     STUDIO_LOGIN_DESCRIPTION = (
@@ -110,7 +110,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(STUDIO_LOGIN_DESCRIPTION, "studio/login"),
         help=STUDIO_LOGIN_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     login_parser.add_argument(
@@ -157,7 +157,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(STUDIO_LOGOUT_DESCRIPTION, "studio/logout"),
         help=STUDIO_LOGOUT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     logout_parser.set_defaults(func=CmdStudioLogout)
@@ -169,7 +169,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(STUDIO_TOKEN_HELP, "studio/token"),
         help=STUDIO_TOKEN_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
 
     logout_parser.set_defaults(func=CmdStudioToken)

dvc/commands/unprotect.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -31,7 +29,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(UNPROTECT_HELP, "unprotect"),
         help=UNPROTECT_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     unprotect_parser.add_argument(
         "targets", nargs="+", help="Data files/directories to unprotect."

dvc/commands/update.py

@@ -1,6 +1,4 @@
-import argparse
-
-from dvc.cli import completion
+from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
@@ -38,7 +36,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(UPDATE_HELP, "update"),
         help=UPDATE_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
     )
     update_parser.add_argument(
         "targets", nargs="+", help=".dvc files to update."

dvc/commands/version.py

@@ -1,5 +1,4 @@
-import argparse
-
+from dvc.cli import formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
@@ -27,7 +26,7 @@ def add_parser(subparsers, parent_parser):
         parents=[parent_parser],
         description=append_doc_link(VERSION_HELP, "version"),
         help=VERSION_HELP,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=formatter.RawDescriptionHelpFormatter,
         aliases=["doctor"],
     )
     version_parser.set_defaults(func=CmdVersion)

dvc/config_schema.py

@@ -240,6 +240,9 @@ SCHEMA = {
                     "ssl_verify": Any(Bool, str),
                     "token": str,
                     "use_https": Bool,
+                    "user": str,
+                    "password": str,
+                    "data_proxy_target": str,
                     Optional("verify", default=False): Bool,
                     **REMOTE_COMMON,
                 },
@@ -257,9 +260,9 @@ SCHEMA = {
                     "exclude_shared_token_cache_credential": Bool,
                     "exclude_managed_identity_credential": Bool,
                     Optional("verify", default=False): Bool,
-                    "timeout": int,
-                    "read_timeout": int,
-                    "connection_timeout": int,
+                    "timeout": Coerce(int),
+                    "read_timeout": Coerce(int),
+                    "connection_timeout": Coerce(int),
                     **REMOTE_COMMON,
                 },
                 "oss": {
@@ -318,8 +321,6 @@ SCHEMA = {
     "feature": FeatureSchema(
         {
             Optional("machine", default=False): Bool,
-            "dbt_profile": str,
-            "dbt_target": str,
         },
     ),
     "plots": {

dvc/database.py

@@ -0,0 +1,133 @@
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from tempfile import NamedTemporaryFile
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import make_url as _make_url
+from sqlalchemy.exc import NoSuchModuleError
+
+from dvc import env
+from dvc.exceptions import DvcException
+from dvc.log import logger
+from dvc.types import StrOrBytesPath
+from dvc.utils import env2bool
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import URL, Connectable, Engine
+    from sqlalchemy.sql.expression import Selectable
+
+
+logger = logger.getChild(__name__)
+
+
+def noop(_):
+    pass
+
+
+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, dict):
+        return make_url(**config)
+    return make_url(config)
+
+
+@contextmanager
+def atomic_file(file: StrOrBytesPath, mode: str = "w+b"):
+    head, tail = os.path.split(os.fsdecode(file))
+    with NamedTemporaryFile(mode, prefix=tail + "-", dir=head, delete=False) as f:
+        yield f
+    os.replace(f.name, file)
+
+
+@dataclass
+class Serializer:
+    sql: "Union[str, Selectable]"
+    con: "Union[str, Connectable]"
+    chunksize: int = 10_000
+
+    def to_csv(self, file: StrOrBytesPath, progress=noop):
+        import pandas as pd
+
+        idfs = pd.read_sql(self.sql, self.con, chunksize=self.chunksize)
+        with atomic_file(file) as f:
+            for i, df in enumerate(idfs):
+                df.to_csv(f, header=i == 0, index=False)
+                progress(len(df))
+
+    def to_json(self, file: StrOrBytesPath, progress=noop):  # noqa: ARG002
+        import pandas as pd
+
+        df = pd.read_sql(self.sql, self.con)
+        with atomic_file(file) as f:
+            df.to_json(f, orient="records")
+
+    def export(self, file: StrOrBytesPath, format: str = "csv", progress=noop):  # noqa: A002
+        if format == "json":
+            return self.to_json(file, progress=progress)
+        return self.to_csv(file, progress=progress)
+
+
+@dataclass
+class Client:
+    engine: "Engine"
+
+    def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None:
+        try:
+            with self.engine.connect() as conn:
+                conn.exec_driver_sql("select 1")
+        except Exception as exc:
+            if callable(onerror):
+                onerror()
+            logger.exception(
+                "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
+
+    def export(
+        self,
+        sql: "Union[str, Selectable]",
+        file: StrOrBytesPath,
+        format: str = "csv",  # noqa: A002
+        progress=noop,
+    ) -> None:
+        con = self.engine.connect()
+        if format == "csv":
+            con = con.execution_options(stream_results=True)  # use server-side cursors
+
+        with con:
+            serializer = Serializer(sql, con)
+            return serializer.export(file, format=format, progress=progress)
+
+
+@contextmanager
+def handle_error(url: "URL"):
+    try:
+        yield
+    except (ModuleNotFoundError, NoSuchModuleError) as e:
+        # TODO: write installation instructions
+        driver = url.drivername
+        raise DvcException(f"Could not load database driver for {driver!r}") from e
+
+
+@contextmanager
+def client(
+    url_or_config: Union[str, "URL", Dict[str, str]], **engine_kwargs: Any
+) -> Iterator[Client]:
+    url = url_from_config(url_or_config)
+    echo = env2bool(env.DVC_SQLALCHEMY_ECHO, False)
+    engine_kwargs.setdefault("echo", echo)
+
+    with handle_error(url):
+        engine = create_engine(url, **engine_kwargs)
+
+    try:
+        yield Client(engine)
+    finally:
+        engine.dispose()

dvc/database/__init__.py

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

dvc/database/dbt_models.py

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

@@ -1,88 +0,0 @@
-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 DbtClient:
-    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 client(
-    project_dir: Optional[str] = None,
-    profiles_dir: Optional[str] = None,
-    profile: Optional[str] = None,
-    target: Optional[str] = None,
-) -> Iterator["DbtClient"]:
-    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 DbtClient(adapter, creds)

dvc/database/dbt_runner.py

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

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

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

@@ -1,73 +0,0 @@
-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 SQLAlchemyClient:
-    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 client(
-    url_or_config: Union[str, URL, Dict[str, str]], **engine_kwargs: Any
-) -> Iterator[SQLAlchemyClient]:
-    url = url_from_config(url_or_config)
-    with handle_error(url):
-        engine = create_engine(url, **engine_kwargs)
-    try:
-        yield SQLAlchemyClient(engine)
-    finally:
-        engine.dispose()

dvc/dependency/__init__.py

@@ -4,7 +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 .db import DbDependency
 from .param import ParamsDependency
 from .repo import RepoDependency
 
@@ -15,7 +15,7 @@ from .repo import RepoDependency
 SCHEMA: Mapping[str, Any] = {
     **ARTIFACT_SCHEMA,
     **RepoDependency.REPO_SCHEMA,
-    **DB_SCHEMA,
+    **DbDependency.DB_SCHEMA,
     Output.PARAM_FILES: [DIR_FILES_SCHEMA],
     Output.PARAM_FS_CONFIG: dict,
 }
@@ -23,16 +23,13 @@ SCHEMA: Mapping[str, Any] = {
 
 def _get(stage, p, info, **kwargs):
     d = info or {}
-    db = d.get(PARAM_DB, {})
     params = d.pop(ParamsDependency.PARAM_PARAMS, None)
     repo = d.pop(RepoDependency.PARAM_REPO, None)
 
     if params:
         return ParamsDependency(stage, p, params)
-    if DbDependency.PARAM_QUERY in db:
-        return DbDependency(stage, info)
-    if db:
-        return DbtDependency(repo, stage, info)
+    if DbDependency.PARAM_DB in d:
+        return DbDependency(stage, d)
 
     assert p
     if repo:

dvc/dependency/db.py

@@ -1,58 +1,20 @@
-import os
-from contextlib import contextmanager, nullcontext
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional
 
 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
@@ -79,13 +41,28 @@ class AbstractDependency(Dependency):
 
 
 class DbDependency(AbstractDependency):
-    PARAM_QUERY = "query"
     PARAM_CONNECTION = "connection"
-    QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str}
+    PARAM_DB = "db"
+    PARAM_QUERY = "query"
+    PARAM_TABLE = "table"
+    PARAM_FILE_FORMAT = "file_format"
+    DB_SCHEMA = {
+        PARAM_DB: {
+            PARAM_QUERY: str,
+            PARAM_CONNECTION: str,
+            PARAM_FILE_FORMAT: str,
+            PARAM_TABLE: str,
+        }
+    }
 
     def __init__(self, stage: "Stage", info, *args, **kwargs):
         super().__init__(stage, info, *args, **kwargs)
-        self.target = None
+        self.db_info: Dict[str, str] = self.info.get(self.PARAM_DB, {})
+        self.connection = self.db_info.get(self.PARAM_CONNECTION)
+
+    @property
+    def sql(self) -> Optional[str]:
+        return self.db_info.get(self.PARAM_QUERY) or self.db_info.get(self.PARAM_TABLE)
 
     def __repr__(self):
         return "{}:{}".format(
@@ -95,9 +72,7 @@ class DbDependency(AbstractDependency):
     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)
+        return truncate_text(self.sql or "", 50)
 
     def workspace_status(self):
         return False  # no workspace to check
@@ -109,8 +84,8 @@ class DbDependency(AbstractDependency):
         """nothing to save."""
 
     def dumpd(self, **kwargs):
-        db_info = compact(self.info.get(PARAM_DB, {}))
-        return {PARAM_DB: db_info} if db_info else {}
+        db_info = compact(self.db_info)
+        return {self.PARAM_DB: db_info} if db_info else {}
 
     def update(self, rev=None):
         """nothing to update."""
@@ -122,203 +97,22 @@ class DbDependency(AbstractDependency):
         file_format: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
-        from dvc.database import export, get_client
+        from dvc.database import client
+        from dvc.ui import ui
 
-        db_info = self.info.get(PARAM_DB, {})
-        query = db_info.get(self.PARAM_QUERY)
-        if not query:
-            raise DvcException("Cannot download: no query specified")
+        sql = self.sql
+        if not sql:
+            raise DvcException("Cannot download: no query or table 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")
+        db_config = self.repo.config.get(self.PARAM_DB, {})
+        config = db_config.get(self.connection)
+        if not config:
+            raise DvcException(f"connection {self.connection} not found in config")
 
-        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_client(
-            config, project_dir=project_dir, profile=profile, target=target
-        ) as db:
-            logger.debug("using adapter: %s", db)
-            with log_status("Testing connection") as status:
+        file_format = file_format or self.db_info.get(self.PARAM_FILE_FORMAT, "csv")
+        with client(config) as db:
+            msg = "Testing connection"
+            with log_durations(logger.debug, msg), ui.status(msg) 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,
-    },
-}
+            with download_progress(to) as progress:
+                db.export(sql, to.fs_path, format=file_format, progress=progress)

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"

dvc/output.py

@@ -407,7 +407,6 @@ class Output:
         self.fs_path = self._parse_path(self.fs, fs_path)
         self.obj: Optional["HashFile"] = None
 
-        self.odb: Optional["HashFileDB"] = None
         self.remote = remote
 
         if self.fs.version_aware:
@@ -444,8 +443,6 @@ class Output:
         if self.files:
             tree = Tree.from_list(self.files, hash_name=self.hash_name)
             tree.digest(with_meta=True)
-            self.odb = HashFileDB(tree.fs, tree.path + ".odb")
-            self.odb.add(tree.path, tree.fs, tree.hash_info.value)
 
             self.hash_info = tree.hash_info
             self.meta.isdir = True

dvc/repo/imp_db.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Dict, Optional
+
+from funcy import compact
 
 from dvc.exceptions import OutputDuplicationError
 from dvc.repo.scm_context import scm_context
@@ -13,17 +15,11 @@ from . import locked
 
 @locked
 @scm_context
-def imp_db(  # noqa: PLR0913
+def imp_db(
     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,
+    table: Optional[str] = None,
     frozen: bool = True,
-    profile: Optional[str] = None,
-    target: Optional[str] = None,
     output_format: str = "csv",
     out: Optional[str] = None,
     force: bool = False,
@@ -34,40 +30,33 @@ def imp_db(  # noqa: PLR0913
         "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 sql or table
     assert output_format in ("csv", "json")
 
-    db: Dict[str, Any] = {"file_format": output_format}
-    if profile:
-        db["profile"] = profile
+    db: Dict[str, str] = compact(
+        {
+            "connection": connection,
+            "file_format": output_format,
+            "query": sql,
+            "table": table,
+        }
+    )
 
-    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})
+    file_name = table or "results"
+    out = out or f"{file_name}.{output_format}"
+    out = resolve_output(".", out, force=force)
 
-    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,
+        deps=[None],
         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:

dvc/repo/imp_url.py

@@ -77,6 +77,7 @@ def imp_url(  # noqa: PLR0913
     elif to_remote:
         remote_odb = self.cloud.get_remote_odb(remote, "import-url")
         stage.outs[0].transfer(url, odb=remote_odb, jobs=jobs)
+        stage.outs[0].ignore()
         stage.save_deps()
         stage.md5 = stage.compute_md5()
     else:

dvc/repo/index.py

@@ -116,6 +116,28 @@ def collect_files(
         dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
 
 
+def _load_data_from_tree(index, prefix, ws, key, tree):
+    from dvc_data.index import DataIndexEntry, Meta
+
+    parents = set()
+
+    for okey, ometa, ohi in tree:
+        for key_len in range(1, len(okey)):
+            parents.add((*key, *okey[:key_len]))
+
+        fkey = (*key, *okey)
+        index[(*prefix, ws, *fkey)] = DataIndexEntry(
+            key=fkey,
+            meta=ometa,
+            hash_info=ohi,
+        )
+
+    for parent in parents:
+        index[(*prefix, ws, *parent)] = DataIndexEntry(
+            key=parent, meta=Meta(isdir=True), loaded=True
+        )
+
+
 def _load_data_from_outs(index, prefix, outs):
     from dvc_data.index import DataIndexEntry, Meta
 
@@ -129,25 +151,25 @@ def _load_data_from_outs(index, prefix, outs):
         for key_len in range(1, len(key)):
             parents.add((ws, key[:key_len]))
 
-        loaded = None
-        if out.files:
-            loaded = True
-            for okey, ometa, ohi in out.get_obj():
-                for key_len in range(1, len(okey)):
-                    parents.add((ws, (*key, *okey[:key_len])))
+        tree = None
+        if (
+            out.stage.is_import
+            and not out.stage.is_repo_import
+            and not out.stage.is_db_import
+            and out.stage.deps[0].files
+        ):
+            tree = out.stage.deps[0].get_obj()
+        elif out.files:
+            tree = out.get_obj()
 
-                fkey = (*key, *okey)
-                index[(*prefix, ws, *fkey)] = DataIndexEntry(
-                    key=fkey,
-                    meta=ometa,
-                    hash_info=ohi,
-                )
+        if tree is not None:
+            _load_data_from_tree(index, prefix, ws, key, tree)
 
         entry = DataIndexEntry(
             key=key,
             meta=out.meta,
             hash_info=out.hash_info,
-            loaded=loaded,
+            loaded=None if tree is None else True,
         )
 
         if (
@@ -174,13 +196,35 @@ def _load_data_from_outs(index, prefix, outs):
         )
 
 
+def _load_storage_from_import(storage_map, key, out):
+    from fsspec.utils import tokenize
+
+    from dvc_data.index import FileStorage
+
+    if out.stage.is_db_import:
+        return
+
+    dep = out.stage.deps[0]
+    if not out.hash_info:
+        # partial import
+        fs_cache = out.repo.cache.fs_cache
+        storage_map.add_cache(
+            FileStorage(
+                key,
+                fs_cache.fs,
+                fs_cache.fs.join(fs_cache.path, dep.fs.protocol, tokenize(dep.fs_path)),
+            )
+        )
+
+    if out.stage.is_repo_import or not out.hash_info or dep.fs.version_aware:
+        storage_map.add_remote(FileStorage(key, dep.fs, dep.fs_path, read_only=True))
+
+
 def _load_storage_from_out(storage_map, key, out):
     from dvc.cachemgr import LEGACY_HASH_NAMES
     from dvc.config import NoRemoteError
     from dvc_data.index import FileStorage, ObjectStorage
 
-    if out.odb:
-        storage_map.add_data(ObjectStorage(key, out.odb))
     storage_map.add_cache(ObjectStorage(key, out.cache))
     try:
         remote = out.repo.cloud.get_remote(out.remote)
@@ -192,36 +236,50 @@ def _load_storage_from_out(storage_map, key, out):
                     path=remote.path,
                     index=remote.index,
                     prefix=(),
+                    read_only=(not out.can_push),
                 )
             )
         else:
             odb = (
                 remote.legacy_odb if out.hash_name in LEGACY_HASH_NAMES else remote.odb
             )
-            storage_map.add_remote(ObjectStorage(key, odb, index=remote.index))
+            storage_map.add_remote(
+                ObjectStorage(
+                    key, odb, index=remote.index, read_only=(not out.can_push)
+                )
+            )
     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:
-            from fsspec.utils import tokenize
+        _load_storage_from_import(storage_map, key, out)
 
-            # partial import
-            fs_cache = out.repo.cache.fs_cache
-            storage_map.add_cache(
-                FileStorage(
-                    key,
-                    fs_cache.fs,
-                    fs_cache.fs.join(
-                        fs_cache.path, dep.fs.protocol, tokenize(dep.fs_path)
-                    ),
-                )
-            )
-        storage_map.add_remote(FileStorage(key, dep.fs, dep.fs_path))
+
+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:
@@ -473,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]":
@@ -741,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"]]:

dvc/repo/metrics/show.py

@@ -17,7 +17,6 @@ from funcy import ldistinct
 from scmrepo.exceptions import SCMError
 
 from dvc.log import logger
-from dvc.repo import locked
 from dvc.scm import NoSCMError
 from dvc.utils import as_posix, expand_paths
 from dvc.utils.collections import ensure_list
@@ -185,7 +184,6 @@ def _hide_workspace(
     return res
 
 
-@locked
 def show(
     repo: "Repo",
     targets: Optional[List[str]] = None,

dvc/repo/params/show.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Un
 
 from dvc.dependency.param import ParamsDependency, read_param_file
 from dvc.log import logger
-from dvc.repo import locked
 from dvc.repo.metrics.show import FileResult, Result
 from dvc.stage import PipelineStage
 from dvc.utils import as_posix, expand_paths
@@ -158,7 +157,6 @@ def _gather_params(
     return data
 
 
-@locked
 def show(
     repo: "Repo",
     targets: Optional[List[str]] = None,

dvc/repo/plots/__init__.py

@@ -487,6 +487,7 @@ def _collect_pipeline_files(repo, targets: List[str], props, onerror=None):
             if isinstance(elem, str):
                 dvcfile_defs_dict[elem] = None
             else:
+                assert elem
                 k, v = next(iter(elem.items()))
                 dvcfile_defs_dict[k] = v
 

dvc/repo/push.py

@@ -37,6 +37,9 @@ def _update_meta(index, **kwargs):
     for remote_name, idx in worktree_view_by_remotes(index, push=True, **kwargs):
         remote = index.repo.cloud.get_remote(remote_name)
 
+        if not remote.fs.version_aware:
+            continue
+
         with ui.progress(
             desc=f"Collecting {remote.path} on {remote.fs.protocol}",
             unit="entry",
@@ -46,9 +49,6 @@ def _update_meta(index, **kwargs):
             new = _rebuild(idx.data["repo"], remote.path, remote.fs, cb)
 
         for out in idx.outs:
-            if not remote.fs.version_aware:
-                continue
-
             _merge_push_meta(out, new, remote.name)
             stages.add(out.stage)
 
@@ -138,7 +138,7 @@ def push(  # noqa: PLR0913
     try:
         with ui.progress(
             desc="Pushing",
-            unit="file",
+            bar_format="{desc}",
             leave=True,
         ) as pb:
             push_transferred, push_failed = ipush(

dvc/schema.py

@@ -112,7 +112,7 @@ PLOT_DEFINITION = {
     Output.PARAM_PLOT_TITLE: str,
     Output.PARAM_PLOT_TEMPLATE: str,
 }
-SINGLE_PLOT_SCHEMA = {str: vol.Any(PLOT_DEFINITION, None)}
+SINGLE_PLOT_SCHEMA = {vol.Required(str): vol.Any(PLOT_DEFINITION, None)}
 ARTIFACTS = "artifacts"
 SINGLE_ARTIFACT_SCHEMA = vol.Schema({str: ARTIFACT_SCHEMA})
 FOREACH_IN = {

dvc/stage/__init__.py

@@ -286,17 +286,17 @@ class Stage(params.StageParams):
         if not self.is_import:
             return False
 
-        from dvc.dependency import DbDependency, DbtDependency
+        from dvc.dependency import DbDependency
 
-        return isinstance(self.deps[0], (DbDependency, DbtDependency))
+        return isinstance(self.deps[0], DbDependency)
 
     @property
     def is_versioned_import(self) -> bool:
-        from dvc.dependency import DbDependency, DbtDependency
+        from dvc.dependency import DbDependency
 
         return (
             self.is_import
-            and not isinstance(self.deps[0], (DbDependency, DbtDependency))
+            and not isinstance(self.deps[0], DbDependency)
             and self.deps[0].fs.version_aware
         )
 
@@ -462,7 +462,7 @@ class Stage(params.StageParams):
         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
+        # always force update DbDep since we don't know if it's changed
         force = self.is_db_import
         update_import(
             self,
@@ -613,7 +613,7 @@ class Stage(params.StageParams):
         ) or self.is_partial_import:
             self._sync_import(dry, force, kwargs.get("jobs", None), no_download)
         elif not self.frozen and self.cmd:
-            self._run_stage(dry, force, allow_missing=allow_missing, **kwargs)
+            self._run_stage(dry, force, **kwargs)
         elif kwargs.get("pull"):
             logger.info("Pulling data for %s", self)
             self.repo.pull(self.addressing, jobs=kwargs.get("jobs", None))

dvc/stage/run.py

@@ -124,11 +124,10 @@ def run_stage(
     dry=False,
     force=False,
     run_env=None,
-    allow_missing: bool = False,
     **kwargs,
 ):
     if not force:
-        if allow_missing and kwargs.get("pull") and not dry:
+        if kwargs.get("pull") and not dry:
             _pull_missing_deps(stage)
 
         from .cache import RunCacheNotFoundError

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() == {}

pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=3.3,<3.4",
+    "dvc-data>=3.7,<3.8",
     "dvc-http>=2.29.0",
     "dvc-render>=1.0.0,<2",
     "dvc-studio-client>=0.17.1,<1",
@@ -77,7 +77,7 @@ gdrive = ["dvc-gdrive>=3,<4"]
 gs = ["dvc-gs>=3,<4"]
 hdfs = ["dvc-hdfs>=3,<4"]
 lint = [
-    "mypy==1.7.1",
+    "mypy==1.8.0",
     "pandas-stubs",
     "types-colorama",
     "types-psutil",
@@ -112,13 +112,12 @@ 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",
+    "pandas>=1", # optional dependency for `import-db`
 ]
 webdav = ["dvc-webdav>=3,<4"]
-webhdfs = ["dvc-webhdfs>=3,<4"]
-webhdfs_kerberos = ["dvc-webhdfs[kerberos]>=3,<4"]
+webhdfs = ["dvc-webhdfs>=3.1,<4"]
+webhdfs_kerberos = ["dvc-webhdfs[kerberos]>=3.1,<4"]
 
 [project.urls]
 Documentation = "https://dvc.org/doc"

tests/func/repro/test_repro_pull.py

@@ -56,3 +56,15 @@ def test_repro_skip_pull_if_no_run_cache_is_passed(tmp_dir, dvc, mocker, local_r
 
     assert dvc.reproduce(pull=True, run_cache=False)
     assert not spy_pull.called
+
+
+def test_repro_skip_pull_if_single_item_is_passed(tmp_dir, dvc, mocker, local_remote):
+    (foo,) = tmp_dir.dvc_gen("foo", "foo")
+
+    dvc.push()
+
+    dvc.stage.add(name="copy-foo", cmd="cp foo bar", deps=["foo"], outs=["bar"])
+    remove("foo")
+    remove(foo.outs[0].cache_path)
+
+    assert dvc.reproduce(pull=True, single_item=True)

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"
@@ -481,12 +481,13 @@ def test_push_pull_fetch_pipeline_stages(tmp_dir, dvc, run_copy, local_remote):
 
 
 def test_pull_partial(tmp_dir, dvc, local_remote):
-    tmp_dir.dvc_gen({"foo": {"bar": {"baz": "baz"}, "spam": "spam"}})
+    other_files = {f"spam{i}": f"spam{i}" for i in range(10)}
+    tmp_dir.dvc_gen({"foo": {"bar": {"baz": "baz"}, **other_files}})
     dvc.push()
     clean(["foo"], dvc)
 
     stats = dvc.pull(os.path.join("foo", "bar"))
-    assert stats["fetched"] == 3
+    assert stats["fetched"] == 2
     assert (tmp_dir / "foo").read_text() == {"bar": {"baz": "baz"}}
 
 

tests/func/test_import_db.py

@@ -1,128 +1,116 @@
-"""import_db tests without involving any databases."""
+import os
+import sqlite3
+from contextlib import closing
+
+import pandas as pd
 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)
+from funcy import compact
 
 
 @pytest.fixture
-def client(mocker):
-    m = mocker.patch("dvc.database.get_client")
-    client = mocker.MagicMock()
-    client.query.return_value.__enter__.side_effect = serializers
-    m.return_value.__enter__.return_value = client
-    return m
+def db_path(tmp_dir):
+    return tmp_dir / "main.db"
 
 
-def test_sql(client, tmp_dir, dvc):
-    stage = dvc.imp_db(
-        sql="select * from model", profile="profile", output_format="json"
-    )
+@pytest.fixture
+def seed_db(db_path):
+    conn = sqlite3.connect(db_path)
+    conn.execute("CREATE TABLE model (id INTEGER PRIMARY KEY, value INTEGER)")
 
-    db = {"file_format": "json", "profile": "profile", "query": "select * from model"}
+    def inner(values):
+        conn.executemany("INSERT INTO model(value) VALUES(?)", [(i,) for i in values])
+        conn.commit()
 
-    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,
-            }
-        ],
-    }
+    with closing(conn):
+        yield inner
 
 
-def test_sql_conn_string(client, tmp_dir, dvc):
+@pytest.fixture
+def db_connection(dvc, db_path):
     with dvc.config.edit(level="local") as conf:
-        conf["db"] = {"conn": {"url": "conn"}}
+        conf["db"] = {"conn": {"url": f"sqlite:///{db_path.fs_path}"}}
+    return "conn"
 
-    stage = dvc.imp_db(
-        sql="select * from model", connection="conn", output_format="json"
+
+def load_data(file, output_format):
+    if output_format == "json":
+        return pd.read_json(file, orient="records")
+    return pd.read_csv(file)
+
+
+@pytest.mark.parametrize("output_format", ("csv", "json"))
+@pytest.mark.parametrize(
+    "args,file_name",
+    [
+        ({"sql": "select * from model"}, "results"),
+        ({"table": "model"}, "model"),
+    ],
+)
+def test(
+    M,
+    tmp_dir,
+    scm,
+    dvc,
+    db_connection,
+    seed_db,
+    output_format,
+    args,
+    file_name,
+):
+    seed_db(values=range(5))
+    if output_format == "json":
+        file_size = 96, 192
+        md5 = "6039fe7565d212b339aaa446ca234e5d", "e1b8adf4d9eb9ab2b64d3ab0bb5f65ac"
+    elif os.name == "nt":
+        file_size = 35, 61
+        md5 = "14c34db5ddd184345c06f74718539f04", "3bb836e6d43c9afa43a9d73b36bbbab4"
+    else:
+        file_size = 29, 50
+        md5 = "6f7fc0d701d1ac13eec83d79fffaf427", "c04f712f8167496a2fb43f289f2b7e28"
+
+    db = compact(
+        {
+            "file_format": output_format,
+            "connection": db_connection,
+            "table": args.get("table"),
+            "query": args.get("sql"),
+        }
     )
-    db = {"file_format": "json", "connection": "conn", "query": "select * from model"}
+    stage = dvc.imp_db(**args, connection=db_connection, output_format=output_format)
+
+    output_file = f"{file_name}.{output_format}"
+    df = load_data(output_file, output_format)
+    assert df.values.tolist() == [[i + 1, i] for i in range(5)]
     assert (tmp_dir / stage.relpath).parse() == {
-        "md5": "f0dcca8e8411907adb25aa547d7c432c",
+        "md5": M.instance_of(str),
         "frozen": True,
         "deps": [{"db": db}],
         "outs": [
             {
+                "md5": md5[0],
+                "size": file_size[0],
                 "hash": "md5",
-                "md5": "56226d443067b195c90b427db557e1f2",
-                "path": "results.json",
-                "size": 243,
+                "path": output_file,
             }
         ],
     }
 
+    seed_db(values=range(5, 10))
+
     dvc.update(stage.addressing)
+
+    df = load_data(output_file, output_format)
+    assert df.values.tolist() == [[i + 1, i] for i in range(10)]
     assert (tmp_dir / stage.relpath).parse() == {
-        "md5": "8c0441b206d8ee7c7b9e3d54871c3cbb",
+        "md5": M.instance_of(str),
         "frozen": True,
         "deps": [{"db": db}],
         "outs": [
             {
+                "md5": md5[1],
+                "size": file_size[1],
                 "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,
+                "path": output_file,
             }
         ],
     }

tests/func/test_import_db_e2e.py

@@ -1,142 +0,0 @@
-"""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_import_url.py

@@ -10,6 +10,7 @@ from dvc.dvcfile import load_file
 from dvc.exceptions import InvalidArgumentError
 from dvc.stage import Stage
 from dvc.testing.workspace_tests import TestImport as _TestImport
+from tests.utils import get_gitignore_content
 
 
 def test_cmd_import(tmp_dir, dvc):
@@ -170,7 +171,7 @@ def test_import_url_preserve_fields(tmp_dir, dvc):
     )
 
 
-def test_import_url_to_remote_absolute(tmp_dir, make_tmp_dir, dvc, local_remote):
+def test_import_url_to_remote_absolute(tmp_dir, make_tmp_dir, dvc, scm, local_remote):
     tmp_abs_dir = make_tmp_dir("abs")
     tmp_foo = tmp_abs_dir / "foo"
     tmp_foo.write_text("foo")
@@ -181,6 +182,7 @@ def test_import_url_to_remote_absolute(tmp_dir, make_tmp_dir, dvc, local_remote)
     assert stage.deps[0].fspath == str(tmp_foo)
     assert stage.outs[0].fspath == os.fspath(foo)
     assert foo.with_suffix(".dvc").exists()
+    assert get_gitignore_content() == ["/foo"]
 
 
 def test_import_url_to_remote_invalid_combinations(dvc):

tests/func/test_repo_index.py

@@ -1,3 +1,4 @@
+import os
 from itertools import chain
 
 import pytest
@@ -338,3 +339,81 @@ params:
     tmp_dir.gen("dvc.yaml", top_level_params)
     index = Index.from_repo(dvc)
     assert index.param_keys == {"repo": {("classifier", "custom_params_file.yaml")}}
+
+
+def test_data_index(tmp_dir, dvc, local_cloud, erepo_dir):
+    tmp_dir.dvc_gen(
+        {
+            "foo": b"foo",
+            "dir": {"bar": b"bar", "subdir": {"baz": b"baz"}},
+        }
+    )
+
+    with erepo_dir.chdir():
+        erepo_dir.dvc_gen("efoo", b"efoo", commit="create efoo")
+        erepo_dir.dvc_gen(
+            "edir",
+            {"ebar": b"ebar", "esubdir": {"ebaz": b"ebaz"}},
+            commit="create edir",
+        )
+
+    dvc.imp(os.fspath(erepo_dir), "efoo")
+    dvc.imp(os.fspath(erepo_dir), "edir")
+
+    dvc.imp(os.fspath(erepo_dir), "efoo", "efoo_partial", no_download=True)
+    dvc.imp(os.fspath(erepo_dir), "edir", "edir_partial", no_download=True)
+
+    local_cloud.gen("ifoo", b"ifoo")
+    local_cloud.gen("idir", {"ibar": b"ibar", "isubdir": {"ibaz": b"ibaz"}})
+
+    dvc.imp_url(str(local_cloud / "ifoo"))
+    dvc.imp_url(str(local_cloud / "idir"))
+
+    dvc.imp_url(str(local_cloud / "ifoo"), "ifoo_partial", no_download=True)
+    dvc.imp_url(str(local_cloud / "idir"), "idir_partial", no_download=True)
+
+    index = Index.from_repo(dvc)
+    assert index.data_keys == {
+        "local": set(),
+        "repo": {
+            ("foo",),
+            ("dir",),
+            ("efoo",),
+            ("edir",),
+            ("efoo_partial",),
+            ("edir_partial",),
+            ("ifoo",),
+            ("idir",),
+            ("ifoo_partial",),
+            ("idir_partial",),
+        },
+    }
+
+    data = index.data["repo"]
+    assert set(data.keys()) == {
+        ("foo",),
+        ("dir",),
+        ("efoo",),
+        ("edir",),
+        ("efoo_partial",),
+        ("edir_partial",),
+        ("ifoo",),
+        ("idir",),
+        ("ifoo_partial",),
+        ("idir_partial",),
+    }
+
+    assert not data.storage_map[("foo",)].remote
+    assert not data.storage_map[("dir",)].remote
+
+    assert data.storage_map[("efoo",)].remote.read_only
+    assert data.storage_map[("edir",)].remote.read_only
+
+    assert data.storage_map[("efoo_partial",)].remote.read_only
+    assert data.storage_map[("edir_partial",)].remote.read_only
+
+    assert not data.storage_map[("ifoo",)].remote
+    assert not data.storage_map[("idir",)].remote
+
+    assert data.storage_map[("ifoo_partial",)].remote.read_only
+    assert data.storage_map[("idir_partial",)].remote.read_only

tests/func/test_virtual_directory.py

@@ -166,32 +166,27 @@ def test_virtual_update_noop(tmp_dir, dvc, remote):
 
 
 def test_partial_checkout_and_update(M, tmp_dir, dvc, remote):
-    tmp_dir.gen({"dir": {"foo": "foo", "subdir": {"lorem": "lorem"}}})
+    dir1 = {f"{i}.txt": f"dir1 {i}" for i in range(10)}
+    dir2 = {f"{i}.txt": f"dir2 {i}" for i in range(10)}
+    tmp_dir.gen({"dir": {"foo": "foo", "subdir": dir1, "subdir2": dir2}})
 
     (stage,) = dvc.add("dir")
     out = stage.outs[0]
 
-    assert out.hash_info == HashInfo(
-        name="md5", value="22a16c9bf84b3068bc2206d88a6b5776.dir"
-    )
-    assert out.meta == Meta(isdir=True, size=8, nfiles=2)
+    assert out.hash_info == HashInfo("md5", "9899dc38082ee0ec33de077df62a4a12.dir")
+    assert out.meta == Meta(isdir=True, size=123, nfiles=21)
 
-    assert dvc.push() == 3
+    assert dvc.push() == 22
     dvc.cache.local.clear()
     shutil.rmtree("dir")
 
-    assert dvc.pull("dir/subdir") == M.dict(
-        added=[join("dir", "")],
-        fetched=3,
-    )
-    assert (tmp_dir / "dir").read_text() == {"subdir": {"lorem": "lorem"}}
+    assert dvc.pull("dir/subdir") == M.dict(added=[join("dir", "")], fetched=11)
+    assert (tmp_dir / "dir").read_text() == {"subdir": dir1}
 
-    tmp_dir.gen({"dir": {"subdir": {"ipsum": "ipsum"}}})
-    (stage,) = dvc.add("dir/subdir/ipsum")
+    tmp_dir.gen({"dir": {"subdir": {"file": "file"}}})
+    (stage,) = dvc.add(join("dir", "subdir", "file"))
 
     out = stage.outs[0]
-    assert out.hash_info == HashInfo(
-        name="md5", value="06d953a10e0b0ffacba04876a9351e39.dir"
-    )
-    assert out.meta == Meta(isdir=True, size=13, nfiles=3)
+    assert out.hash_info == HashInfo("md5", "e7531c3930f28a00edc25e6cef91db03.dir")
+    assert out.meta == Meta(isdir=True, size=127, nfiles=22)
     assert dvc.push() == 2

tests/unit/command/test_help.py

@@ -1,5 +1,7 @@
 import logging
+import re
 from argparse import SUPPRESS, ArgumentParser
+from itertools import takewhile
 from typing import Tuple
 
 import pytest
@@ -28,7 +30,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)
 
 
@@ -44,5 +46,19 @@ def test_help(caplog, capsys, command_tuples):
     assert not caplog.text
 
     out, err = capsys.readouterr()
+
+    # validate metavars are all in lowercase
+    usage = "\n".join(takewhile(lambda o: bool(o), out.splitlines()))
+
+    message = (
+        "metavars not lowercased, you are likely missing formatter_class=XXX "
+        "in the command, where XXX should be any of the classes from "
+        "`dvc.cli.formatter`, which automatically lowercases metavars.\n"
+        "\nExample:\n"
+        "\nfrom dvc.cli import formatter\n"
+        "\nparser.add_parser(..., formatter_class=formatter.TextHelpFormatter)\n"
+    )
+    assert not re.findall(r"\b[A-Z]{2,}\b", usage, re.MULTILINE), message
+
     assert not err
     assert out