experiments: auto push experiments (#10323)

* add auto_push for experiment on run and on save.
add config.exp.auto_push and config.exp.git_remote

* add func tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* separate auto_push behaviour out of commit

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Saugat Pachhai (सौगात) <[email protected]>

.github/workflows/build.yaml

@@ -25,7 +25,7 @@ jobs:
 
     - uses: actions/setup-python@v5
       with:
-        python-version: 3.8
+        python-version: 3.12
 
     - run: python -m pip install --upgrade pip setuptools_scm build twine
 

.github/workflows/tests.yaml

@@ -23,7 +23,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-20.04, windows-latest, macos-latest]
-        pyv: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        pyv: ["3.9", "3.10", "3.11", "3.12"]
 
     steps:
     - uses: actions/checkout@v4
@@ -39,7 +39,7 @@ jobs:
       run: |
         pip install --upgrade pip wheel
         pip install -e ".[dev]"
-    - uses: pre-commit/[email protected].0
+    - uses: pre-commit/[email protected].1
     - name: run tests
       timeout-minutes: 40
       env:

.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.2.2'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

dvc/__pyinstaller/hook-celery.py

@@ -5,10 +5,7 @@ from PyInstaller.utils.hooks import collect_submodules, is_module_or_submodule
 # Celery dynamically imports most celery internals at runtime
 # pyinstaller hook must expose all modules loaded by
 # kombu.utils.imports:symbol_by_name()
-_EXCLUDES = (
-    "celery.bin",
-    "celery.contrib",
-)
+_EXCLUDES = ("celery.bin", "celery.contrib")
 hiddenimports = collect_submodules(
     "celery",
     filter=lambda name: not any(

dvc/annotations.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict, dataclass, field, fields
-from typing import Any, ClassVar, Dict, List, Optional
+from typing import Any, ClassVar, Optional
 
 from funcy import compact
 from voluptuous import Required
@@ -13,11 +13,11 @@ class Annotation:
     PARAM_META: ClassVar[str] = "meta"
 
     desc: Optional[str] = None
-    type: Optional[str] = None  # noqa: A003
-    labels: List[str] = field(default_factory=list)
-    meta: Dict[str, Any] = field(default_factory=dict)
+    type: Optional[str] = None
+    labels: list[str] = field(default_factory=list)
+    meta: dict[str, Any] = field(default_factory=dict)
 
-    def to_dict(self) -> Dict[str, str]:
+    def to_dict(self) -> dict[str, str]:
         return compact(asdict(self))
 
 
@@ -31,11 +31,11 @@ class Artifact:
 
     path: str
     desc: Optional[str] = None
-    type: Optional[str] = None  # noqa: A003
-    labels: List[str] = field(default_factory=list)
-    meta: Dict[str, Any] = field(default_factory=dict)
+    type: Optional[str] = None
+    labels: list[str] = field(default_factory=list)
+    meta: dict[str, Any] = field(default_factory=dict)
 
-    def to_dict(self) -> Dict[str, str]:
+    def to_dict(self) -> dict[str, str]:
         return compact(asdict(self))
 
 
@@ -46,7 +46,7 @@ ANNOTATION_SCHEMA = {
     Annotation.PARAM_LABELS: [str],
     Annotation.PARAM_META: object,
 }
-ARTIFACT_SCHEMA: Dict[Any, Any] = {
+ARTIFACT_SCHEMA: dict[Any, Any] = {
     Required(Artifact.PARAM_PATH): str,
     **ANNOTATION_SCHEMA,  # type: ignore[arg-type]
 }

dvc/api/__init__.py

@@ -2,21 +2,23 @@ from dvc.fs.dvc import _DVCFileSystem as DVCFileSystem
 
 from .artifacts import artifacts_show
 from .data import get_url, open, read
+from .dataset import get as get_dataset
 from .experiments import exp_save, exp_show
 from .scm import all_branches, all_commits, all_tags
 from .show import metrics_show, params_show
 
 __all__ = [
+    "DVCFileSystem",
     "all_branches",
     "all_commits",
     "all_tags",
     "artifacts_show",
     "exp_save",
     "exp_show",
+    "get_dataset",
     "get_url",
+    "metrics_show",
     "open",
     "params_show",
-    "metrics_show",
     "read",
-    "DVCFileSystem",
 ]

dvc/api/artifacts.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from dvc.repo import Repo
 
@@ -8,7 +8,7 @@ def artifacts_show(
     version: Optional[str] = None,
     stage: Optional[str] = None,
     repo: Optional[str] = None,
-) -> Dict[str, str]:
+) -> dict[str, str]:
     """
     Return path and Git revision for an artifact in a DVC project.
 
@@ -36,7 +36,7 @@ def artifacts_show(
     if version and stage:
         raise ValueError("Artifact version and stage are mutually exclusive.")
 
-    repo_kwargs: Dict[str, Any] = {
+    repo_kwargs: dict[str, Any] = {
         "subrepos": True,
         "uninitialized": True,
     }

dvc/api/data.py

@@ -1,6 +1,6 @@
 from contextlib import _GeneratorContextManager as GCM
 from contextlib import contextmanager
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from funcy import reraise
 
@@ -39,7 +39,7 @@ def get_url(path, repo=None, rev=None, remote=None):
     from dvc.config import NoRemoteError
     from dvc_data.index import StorageKeyError
 
-    repo_kwargs: Dict[str, Any] = {}
+    repo_kwargs: dict[str, Any] = {}
     if remote:
         repo_kwargs["config"] = {"core": {"remote": remote}}
     with Repo.open(
@@ -57,11 +57,7 @@ def get_url(path, repo=None, rev=None, remote=None):
 class _OpenContextManager(GCM):
     def __init__(self, func, args, kwds):
         self.gen = func(*args, **kwds)
-        self.func, self.args, self.kwds = (  # type: ignore[assignment]
-            func,
-            args,
-            kwds,
-        )
+        self.func, self.args, self.kwds = (func, args, kwds)  # type: ignore[assignment]
 
     def __getattr__(self, name):
         raise AttributeError("dvc.api.open() should be used in a with statement.")
@@ -74,8 +70,8 @@ def open(  # noqa: A001
     remote: Optional[str] = None,
     mode: str = "r",
     encoding: Optional[str] = None,
-    config: Optional[Dict[str, Any]] = None,
-    remote_config: Optional[Dict[str, Any]] = None,
+    config: Optional[dict[str, Any]] = None,
+    remote_config: Optional[dict[str, Any]] = None,
 ):
     """
     Opens a file tracked in a DVC project.
@@ -233,7 +229,7 @@ def _open(
     config=None,
     remote_config=None,
 ):
-    repo_kwargs: Dict[str, Any] = {
+    repo_kwargs: dict[str, Any] = {
         "subrepos": True,
         "uninitialized": True,
         "remote": remote,
@@ -262,11 +258,7 @@ def _open(
                 fs_path = fs.from_os_path(path)
 
             try:
-                with fs.open(
-                    fs_path,
-                    mode=mode,
-                    encoding=encoding,
-                ) as fobj:
+                with fs.open(fs_path, mode=mode, encoding=encoding) as fobj:
                     yield fobj
             except FileNotFoundError as exc:
                 raise FileMissingError(path) from exc

dvc/api/dataset.py

@@ -0,0 +1,66 @@
+from typing import Literal, TypedDict, Union
+
+
+class DVCXDataset(TypedDict):
+    type: Literal["dvcx"]
+    name: str
+    version: int
+
+
+class DVCDataset(TypedDict):
+    type: Literal["dvc"]
+    url: str
+    path: str
+    sha: str
+
+
+class URLDataset(TypedDict):
+    type: Literal["url"]
+    files: list[str]
+    path: str
+
+
+def get(name: str) -> Union[DVCXDataset, DVCDataset, URLDataset]:
+    from difflib import get_close_matches
+
+    from dvc.fs import get_cloud_fs
+    from dvc.repo import Repo, datasets
+
+    repo = Repo()
+    try:
+        dataset = repo.datasets[name]
+    except datasets.DatasetNotFoundError as e:
+        add_note = getattr(e, "add_note", lambda _: None)
+        if matches := get_close_matches(name, repo.datasets):
+            add_note(f"Did you mean: {matches[0]!r}?")
+        raise
+
+    if dataset._invalidated:
+        raise ValueError(f"dataset not in sync. Sync with 'dvc ds update {name}'.")
+    if not dataset.lock:
+        raise ValueError("missing lock information")
+    if dataset.type == "dvc":
+        return DVCDataset(
+            type="dvc",
+            url=dataset.lock.url,
+            path=dataset.lock.path,
+            sha=dataset.lock.rev_lock,
+        )
+    if dataset.type == "dvcx":
+        return DVCXDataset(
+            type="dvcx", name=dataset.name_version[0], version=dataset.lock.version
+        )
+    if dataset.type == "url":
+        fs_cls, _, path = get_cloud_fs(repo.config, url=dataset.lock.url)
+        assert fs_cls
+        join_version = getattr(fs_cls, "join_version", lambda path, _: path)
+        protocol = fs_cls.protocol
+        versioned_path = join_version(path, dataset.lock.meta.version_id)
+        versioned_path = f"{protocol}://{versioned_path}"
+        files = [
+            join_version(
+                fs_cls.join(versioned_path, file.relpath), file.meta.version_id
+            )
+            for file in dataset.lock.files
+        ]
+        return URLDataset(type="url", files=files, path=versioned_path)

dvc/api/experiments.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union
+from typing import Optional, Union
 
 from rich.text import Text
 
@@ -9,7 +9,7 @@ from dvc.repo.experiments.show import tabulate
 def exp_save(
     name: Optional[str] = None,
     force: bool = False,
-    include_untracked: Optional[List[str]] = None,
+    include_untracked: Optional[list[str]] = None,
 ):
     """
     Create a new DVC experiment using `exp save`.
@@ -61,12 +61,12 @@ def _postprocess(exp_rows):
 
 def exp_show(
     repo: Optional[str] = None,
-    revs: Optional[Union[str, List[str]]] = None,
+    revs: Optional[Union[str, list[str]]] = None,
     num: int = 1,
     param_deps: bool = False,
     force: bool = False,
-    config: Optional[Dict] = None,
-) -> List[Dict]:
+    config: Optional[dict] = None,
+) -> list[dict]:
     """Get DVC experiments tracked in `repo`.
 
     Without arguments, this function will retrieve all experiments derived from

dvc/api/scm.py

@@ -1,9 +1,9 @@
-from typing import List, Optional
+from typing import Optional
 
 from dvc.repo import Repo
 
 
-def all_branches(repo: Optional[str] = None) -> List[str]:
+def all_branches(repo: Optional[str] = None) -> list[str]:
     """Get all Git branches in a DVC repository.
 
     Args:
@@ -20,7 +20,7 @@ def all_branches(repo: Optional[str] = None) -> List[str]:
         return _repo.scm.list_branches()
 
 
-def all_commits(repo: Optional[str] = None) -> List[str]:
+def all_commits(repo: Optional[str] = None) -> list[str]:
     """Get all Git commits in a DVC repository.
 
     Args:
@@ -37,7 +37,7 @@ def all_commits(repo: Optional[str] = None) -> List[str]:
         return _repo.scm.list_all_commits()
 
 
-def all_tags(repo: Optional[str] = None) -> List[str]:
+def all_tags(repo: Optional[str] = None) -> list[str]:
     """Get all Git tags in a DVC repository.
 
     Args:

dvc/api/show.py

@@ -1,6 +1,7 @@
 import typing
 from collections import Counter
-from typing import Dict, Iterable, Optional, Union
+from collections.abc import Iterable
+from typing import Optional, Union
 
 from funcy import first
 
@@ -8,7 +9,7 @@ from dvc.repo import Repo
 
 
 def _postprocess(results):
-    processed: Dict[str, Dict] = {}
+    processed: dict[str, dict] = {}
     for rev, rev_data in results.items():
         if not rev_data:
             continue
@@ -24,7 +25,7 @@ def _postprocess(results):
                 (k if counts[k] == 1 else f"{file_name}:{k}"): v
                 for k, v in file_data["data"].items()
             }
-            processed[rev] = {**processed[rev], **to_merge}
+            processed[rev] = processed[rev] | to_merge
 
     if "workspace" in processed:
         del processed["workspace"]
@@ -36,8 +37,8 @@ def metrics_show(
     *targets: str,
     repo: Optional[str] = None,
     rev: Optional[str] = None,
-    config: Optional[Dict] = None,
-) -> Dict:
+    config: Optional[dict] = None,
+) -> dict:
     """Get metrics tracked in `repo`.
 
     Without arguments, this function will retrieve all metrics from all tracked
@@ -159,8 +160,8 @@ def params_show(
     stages: Optional[Union[str, Iterable[str]]] = None,
     rev: Optional[str] = None,
     deps: bool = False,
-    config: Optional[Dict] = None,
-) -> Dict:
+    config: Optional[dict] = None,
+) -> dict:
     """Get parameters tracked in `repo`.
 
     Without arguments, this function will retrieve all params from all tracked

dvc/cachemgr.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Optional
 
 from dvc.fs import GitFileSystem, Schemes
 from dvc_data.hashfile.db import get_odb
@@ -15,7 +15,7 @@ def _get_odb(
     repo,
     settings,
     fs=None,
-    prefix: Optional[Tuple[str, ...]] = None,
+    prefix: Optional[tuple[str, ...]] = None,
     hash_name: Optional[str] = None,
     **kwargs,
 ):
@@ -144,15 +144,9 @@ def migrate_2_to_3(repo: "Repo", dry: bool = False):
         )
         return
 
-    with TqdmCallback(
-        desc="Computing DVC 3.0 hashes",
-        unit="files",
-    ) as cb:
+    with TqdmCallback(desc="Computing DVC 3.0 hashes", unit="files") as cb:
         migration = prepare(src, dest, callback=cb)
 
-    with TqdmCallback(
-        desc="Migrating to DVC 3.0 cache",
-        unit="files",
-    ) as cb:
+    with TqdmCallback(desc="Migrating to DVC 3.0 cache", unit="files") as cb:
         count = migrate(migration, callback=cb)
     ui.write(f"Migrated {count} files to DVC 3.0 cache location.")

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

@@ -17,6 +17,7 @@ from dvc.commands import (
     dag,
     data,
     data_sync,
+    dataset,
     destroy,
     diff,
     du,
@@ -50,53 +51,54 @@ from dvc.commands import (
 )
 from dvc.log import logger
 
-from . import DvcParserError
+from . import DvcParserError, formatter
 
 logger = logger.getChild(__name__)
 
 COMMANDS = [
-    init,
-    queue,
+    add,
+    artifacts,
+    cache,
+    check_ignore,
+    checkout,
+    commit,
+    completion,
+    config,
+    daemon,
+    dag,
+    data,
+    data_sync,
+    dataset,
+    destroy,
+    diff,
+    du,
+    experiments,
+    freeze,
+    gc,
     get,
     get_url,
-    destroy,
-    add,
-    remove,
-    move,
-    unprotect,
-    repro,
-    data_sync,
-    gc,
+    git_hook,
     imp,
     imp_db,
     imp_url,
-    config,
-    checkout,
-    remote,
-    cache,
-    metrics,
-    params,
+    init,
     install,
-    root,
     ls,
     ls_url,
-    freeze,
-    dag,
-    daemon,
-    commit,
-    completion,
-    diff,
-    version,
-    update,
-    git_hook,
+    metrics,
+    move,
+    params,
     plots,
+    queue,
+    remote,
+    remove,
+    repro,
+    root,
     stage,
-    experiments,
-    check_ignore,
-    data,
-    artifacts,
     studio,
-    du,
+    unprotect,
+    update,
+    version,
 ]
 
 
@@ -165,7 +167,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 +202,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 typing import TYPE_CHECKING, ClassVar
 
 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
@@ -17,14 +17,14 @@ logger = logger.getChild(__name__)
 
 
 class CmdDataStatus(CmdBase):
-    COLORS = {
+    COLORS: ClassVar[dict[str, str]] = {
         "not_in_remote": "red",
         "not_in_cache": "red",
         "committed": "green",
         "uncommitted": "yellow",
         "untracked": "cyan",
     }
-    LABELS = {
+    LABELS: ClassVar[dict[str, str]] = {
         "not_in_remote": "Not in remote",
         "not_in_cache": "Not in cache",
         "committed": "DVC committed changes",
@@ -32,7 +32,7 @@ class CmdDataStatus(CmdBase):
         "untracked": "Untracked files",
         "unchanged": "DVC unchanged files",
     }
-    HINTS = {
+    HINTS: ClassVar[dict[str, tuple[str, ...]]] = {
         "not_in_remote": ('use "dvc push <file>..." to upload files',),
         "not_in_cache": ('use "dvc fetch <file>..." to download files',),
         "committed": ("git commit the corresponding dvc files to update the repo",),
@@ -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",
@@ -429,5 +429,11 @@ def add_parser(subparsers, _parent_parser):
         default=False,
         help="Show status in JSON format.",
     )
+    status_parser.add_argument(
+        "--no-updates",
+        dest="check_updates",
+        action="store_false",
+        help="Ignore updates to imported data.",
+    )
 
     status_parser.set_defaults(func=CmdDataStatus)

dvc/commands/dataset.py

@@ -0,0 +1,225 @@
+from typing import TYPE_CHECKING, Optional
+
+from dvc.cli import formatter
+from dvc.cli.command import CmdBase
+from dvc.cli.utils import append_doc_link
+from dvc.exceptions import DvcException
+from dvc.log import logger
+
+if TYPE_CHECKING:
+    from rich.text import Text
+
+    from dvc.repo.datasets import Dataset, FileInfo
+
+logger = logger.getChild(__name__)
+
+
+def diff_files(old: list["FileInfo"], new: list["FileInfo"]) -> dict[str, list[str]]:
+    old_files = {d.relpath: d for d in old}
+    new_files = {d.relpath: d for d in new}
+    rest = old_files.keys() & new_files.keys()
+    return {
+        "added": list(new_files.keys() - old_files.keys()),
+        "deleted": list(old_files.keys() - new_files.keys()),
+        "modified": [p for p in rest if new_files[p] != old_files[p]],
+    }
+
+
+class CmdDatasetAdd(CmdBase):
+    @classmethod
+    def display(cls, name: str, dataset: "Dataset", action: str = "Adding"):
+        from dvc.ui import ui
+
+        assert dataset.lock
+
+        url = dataset.spec.url
+        ver: str = ""
+        if dataset.type == "dvcx":
+            ver = f"v{dataset.lock.version}"
+        if dataset.type == "dvc":
+            if dataset.lock.path:
+                url = f"{url}:/{dataset.lock.path.lstrip('/')}"
+            if rev := dataset.lock.rev:
+                ver = rev
+
+        ver_part: Optional["Text"] = None
+        if ver:
+            ver_part = ui.rich_text.assemble(" @ ", (ver, "repr.number"))
+        text = ui.rich_text.assemble("(", (url, "repr.url"), ver_part or "", ")")
+        ui.write(action, ui.rich_text(name, "cyan"), text, styled=True)
+
+    def run(self):
+        if not self.args.dvc and self.args.rev:
+            raise DvcException("--rev can't be used without --dvc")
+        if not self.args.dvc and self.args.path:
+            raise DvcException("--path can't be used without --dvc")
+
+        d = vars(self.args)
+        for key in ["dvc", "dvcx", "url"]:
+            if url := d.pop(key, None):
+                d.update({"type": key, "url": url})
+                break
+
+        existing = self.repo.datasets.get(self.args.name)
+        with self.repo.scm_context:
+            if not self.args.force and existing:
+                path = self.repo.fs.relpath(existing.manifest_path)
+                raise DvcException(
+                    f"{self.args.name} already exists in {path}, "
+                    "use the --force to overwrite"
+                )
+            dataset = self.repo.datasets.add(**d)
+            self.display(self.args.name, dataset)
+            return 0
+
+
+class CmdDatasetUpdate(CmdBase):
+    def display(self, name: str, dataset: "Dataset", new: "Dataset"):
+        from dvc.commands.checkout import log_changes
+        from dvc.ui import ui
+
+        action = "Updating"
+        if not dataset.lock:
+            return CmdDatasetAdd.display(name, new, action)
+        if dataset == new:
+            ui.write("[yellow]Nothing to update[/]", styled=True)
+            return
+
+        assert new.lock
+
+        v: Optional[tuple[str, str]] = None
+        if dataset.type == "dvcx":
+            assert new.type == "dvcx"
+            if new.lock.version < dataset.lock.version:
+                action = "Downgrading"
+
+            v = (f"v{dataset.lock.version}", f"v{new.lock.version}")
+        if dataset.type == "dvc":
+            assert new.type == "dvc"
+            v = (f"{dataset.lock.rev_lock[:9]}", f"{new.lock.rev_lock[:9]}")
+
+        if v:
+            part = ui.rich_text.assemble(
+                (v[0], "repr.number"),
+                " -> ",
+                (v[1], "repr.number"),
+            )
+        else:
+            part = ui.rich_text(dataset.spec.url, "repr.url")
+        changes = ui.rich_text.assemble("(", part, ")")
+        ui.write(action, ui.rich_text(name, "cyan"), changes, styled=True)
+        if dataset.type == "url":
+            assert new.type == "url"
+            stats = diff_files(dataset.lock.files, new.lock.files)
+            log_changes(stats)
+
+    def run(self):
+        from difflib import get_close_matches
+
+        from dvc.repo.datasets import DatasetNotFoundError
+        from dvc.ui import ui
+
+        version = None
+        if self.args.rev:
+            try:
+                version = int(self.args.rev.lstrip("v"))
+            except ValueError:
+                version = self.args.rev
+
+        d = vars(self.args) | {"version": version}
+        with self.repo.scm_context:
+            try:
+                dataset, new = self.repo.datasets.update(**d)
+            except DatasetNotFoundError:
+                logger.exception("")
+                if matches := get_close_matches(self.args.name, self.repo.datasets):
+                    ui.write(
+                        "did you mean?",
+                        ui.rich_text(matches[0], "cyan"),
+                        stderr=True,
+                        styled=True,
+                    )
+                return 1
+            self.display(self.args.name, dataset, new)
+            return 0
+
+
+def add_parser(subparsers, parent_parser):
+    ds_parser = subparsers.add_parser(
+        "dataset",
+        aliases=["ds"],
+        parents=[parent_parser],
+        formatter_class=formatter.RawDescriptionHelpFormatter,
+    )
+    ds_subparsers = ds_parser.add_subparsers(
+        dest="cmd",
+        help="Use `dvc dataset CMD --help` to display command-specific help.",
+        required=True,
+    )
+
+    dataset_add_help = "Add a dataset."
+    ds_add_parser = ds_subparsers.add_parser(
+        "add",
+        parents=[parent_parser],
+        description=append_doc_link(dataset_add_help, "dataset/add"),
+        formatter_class=formatter.RawTextHelpFormatter,
+        help=dataset_add_help,
+    )
+
+    url_exclusive_group = ds_add_parser.add_mutually_exclusive_group(required=True)
+    url_exclusive_group.add_argument(
+        "--dvcx", metavar="name", help="Name of the dvcx dataset to track"
+    )
+    url_exclusive_group.add_argument(
+        "--dvc",
+        help="Path or URL to a Git/DVC repository to track",
+        metavar="url",
+    )
+    url_exclusive_group.add_argument(
+        "--url",
+        help="""\
+URL of a cloud-versioned remote to track. Supported URLs:
+
+s3://bucket/key/path
+gs://bucket/path/to/file/or/dir
+azure://mycontainer/path
+remote://remote_name/path/to/file/or/dir (see `dvc remote`)
+""",
+    )
+    ds_add_parser.add_argument("name", help="Name of the dataset to add")
+    ds_add_parser.add_argument(
+        "--rev",
+        help="Git revision, e.g. SHA, branch, tag (only applicable with --dvc)",
+        metavar="<commit>",
+    )
+    ds_add_parser.add_argument(
+        "--path",
+        help="Path to a file or a directory within a git repository "
+        "(only applicable with --dvc)",
+    )
+    ds_add_parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        default=False,
+        help="Overwrite existing dataset",
+    )
+    ds_add_parser.set_defaults(func=CmdDatasetAdd)
+
+    dataset_update_help = "Update a dataset."
+    ds_update_parser = ds_subparsers.add_parser(
+        "update",
+        parents=[parent_parser],
+        description=append_doc_link(dataset_update_help, "dataset/update"),
+        formatter_class=formatter.RawDescriptionHelpFormatter,
+        help=dataset_update_help,
+    )
+    ds_update_parser.add_argument("name", help="Name of the dataset to update")
+    ds_update_parser.add_argument(
+        "--rev",
+        "--version",
+        nargs="?",
+        help="DVCX dataset version or Git revision (e.g. SHA, branch, tag)",
+        metavar="<version>",
+    )
+    ds_update_parser.set_defaults(func=CmdDatasetUpdate)

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 typing import Any
 
-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
@@ -12,7 +11,7 @@ logger = logger.getChild(__name__)
 
 class CmdExperimentsPush(CmdBase):
     @staticmethod
-    def log_result(result: Dict[str, Any], remote: str):
+    def log_result(result: dict[str, Any], remote: str):
         from dvc.utils import humanize
 
         def join_exps(exps):
@@ -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)
@@ -106,8 +107,4 @@ def _add_run_common(parser):
         default=None,
         help="Custom commit message to use when committing the experiment.",
     )
-    parser.add_argument(
-        "-M",  # obsolete
-        dest="message",
-        help=argparse.SUPPRESS,
-    )
+    parser.add_argument("-M", dest="message", help=argparse.SUPPRESS)  # obsolete

dvc/commands/experiments/save.py

@@ -1,5 +1,6 @@
 import argparse
 
+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
@@ -13,7 +14,9 @@ class CmdExperimentsSave(CmdBase):
     def run(self):
         try:
             ref = self.repo.experiments.save(
+                targets=self.args.targets,
                 name=self.args.name,
+                recursive=self.args.recursive,
                 force=self.args.force,
                 include_untracked=self.args.include_untracked,
                 message=self.args.message,
@@ -38,7 +41,19 @@ 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(
+        "targets",
+        nargs="*",
+        help=("Limit DVC caching to these .dvc files and stage names."),
+    ).complete = completion.DVCFILES_AND_STAGE
+    save_parser.add_argument(
+        "-R",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="Cache subdirectories of the specified directory.",
     )
     save_parser.add_argument(
         "-f",
@@ -78,9 +93,5 @@ def add_parser(experiments_subparsers, parent_parser):
         default=None,
         help="Custom commit message to use when committing the experiment.",
     )
-    save_parser.add_argument(
-        "-M",  # obsolete
-        dest="message",
-        help=argparse.SUPPRESS,
-    )
+    save_parser.add_argument("-M", dest="message", help=argparse.SUPPRESS)  # obsolete
     save_parser.set_defaults(func=CmdExperimentsSave)

dvc/commands/experiments/show.py

@@ -1,10 +1,12 @@
 import argparse
 import re
+from collections.abc import Iterable
 from datetime import date, datetime
-from typing import TYPE_CHECKING, Dict, Iterable
+from typing import TYPE_CHECKING
 
 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
@@ -56,7 +58,7 @@ def baseline_styler(typ):
 
 def show_experiments(
     td: "TabularData",
-    headers: Dict[str, Iterable[str]],
+    headers: dict[str, Iterable[str]],
     keep=None,
     drop=None,
     pager=True,
@@ -200,7 +202,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
@@ -86,6 +86,7 @@ class CmdGC(CmdBase):
             num=self.args.num,
             not_in_remote=self.args.not_in_remote,
             dry=self.args.dry,
+            skip_failed=self.args.skip_failed,
         )
         return 0
 
@@ -101,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",
@@ -155,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(
@@ -188,6 +189,12 @@ def add_parser(subparsers, parent_parser):
         help="Remote storage to collect garbage in",
         metavar="<name>",
     )
+    gc_parser.add_argument(
+        "--skip-failed",
+        action="store_true",
+        default=False,
+        help="Skip revisions that fail when collected.",
+    )
     gc_parser.add_argument(
         "-f",
         "--force",

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
@@ -59,6 +57,9 @@ class CmdList(CmdBaseNoRepo):
             elif entries:
                 show_entries(entries, with_color=True, with_size=self.args.size)
             return 0
+        except FileNotFoundError:
+            logger.exception("")
+            return 1
         except DvcException:
             logger.exception("failed to list '%s'", self.args.url)
             return 1
@@ -75,7 +76,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(
@@ -121,11 +122,7 @@ def add_parser(subparsers, parent_parser):
             "specified by '--remote') in the target repository."
         ),
     )
-    list_parser.add_argument(
-        "--size",
-        action="store_true",
-        help="Show sizes.",
-    )
+    list_parser.add_argument("--size", action="store_true", help="Show sizes.")
     list_parser.add_argument(
         "path",
         nargs="?",

dvc/commands/ls/ls_colors.py

@@ -20,7 +20,7 @@ class LsColors:
             else:
                 self._codes[code] = color
 
-    def format(self, entry):  # noqa: A003
+    def format(self, entry):
         text = entry["path"]
 
         if entry.get("isout", False) and "out" in self._codes:

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,22 +32,15 @@ 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."
     )
     lsurl_parser.add_argument(
-        "-R",
-        "--recursive",
-        action="store_true",
-        help="Recursively list files.",
-    )
-    lsurl_parser.add_argument(
-        "--size",
-        action="store_true",
-        help="Show sizes.",
+        "-R", "--recursive", action="store_true", help="Recursively list files."
     )
+    lsurl_parser.add_argument("--size", action="store_true", help="Show sizes.")
     lsurl_parser.add_argument(
         "--fs-config",
         type=str,

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

@@ -1,12 +1,12 @@
 import argparse
 import os
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, 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
@@ -20,35 +20,26 @@ logger = logger.getChild(__name__)
 
 
 def _show_json(
-    renderers_with_errors: List["RendererWithErrors"],
+    renderers_with_errors: list["RendererWithErrors"],
     split=False,
-    errors: Optional[Dict[str, Exception]] = None,
+    errors: Optional[dict[str, Exception]] = None,
 ):
     from dvc.render.convert import to_json
     from dvc.utils.serialize import encode_exception
 
-    all_errors: List[Dict] = []
+    all_errors: list[dict] = []
     data = {}
 
     for renderer, src_errors, def_errors in renderers_with_errors:
         name = renderer.name
         data[name] = to_json(renderer, split)
         all_errors.extend(
-            {
-                "name": name,
-                "rev": rev,
-                "source": source,
-                **encode_exception(e),
-            }
+            {"name": name, "rev": rev, "source": source, **encode_exception(e)}
             for rev, per_rev_src_errors in src_errors.items()
             for source, e in per_rev_src_errors.items()
         )
         all_errors.extend(
-            {
-                "name": name,
-                "rev": rev,
-                **encode_exception(e),
-            }
+            {"name": name, "rev": rev, **encode_exception(e)}
             for rev, e in def_errors.items()
         )
 
@@ -101,10 +92,7 @@ class CmdPlots(CmdBase):
                 return 1
 
         try:
-            plots_data = self._func(
-                targets=self.args.targets,
-                props=self._props(),
-            )
+            plots_data = self._func(targets=self.args.targets, props=self._props())
 
             if not plots_data and not self.args.json:
                 ui.error_write(
@@ -223,15 +211,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 +227,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 +250,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 +286,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 +307,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,16 +1,8 @@
-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 = [
-    start,
-    stop,
-    status,
-    logs,
-    remove,
-    kill,
-]
+SUB_COMMANDS = [start, stop, status, logs, remove, kill]
 
 
 def add_parser(subparsers, parent_parser):
@@ -20,15 +12,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",
@@ -46,9 +45,5 @@ def add_parser(queue_subparsers, parent_parser):
         ),
         action="store_true",
     )
-    queue_logs_parser.add_argument(
-        "task",
-        help="Task to show.",
-        metavar="<task>",
-    )
+    queue_logs_parser.add_argument("task", help="Task to show.", metavar="<task>")
     queue_logs_parser.set_defaults(func=CmdQueueLogs)

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

@@ -1,12 +1,13 @@
 import argparse
 import logging
+from collections.abc import Iterable
 from contextlib import contextmanager
 from itertools import chain, filterfalse
-from typing import TYPE_CHECKING, Dict, Iterable, List
+from typing import TYPE_CHECKING
 
-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
@@ -33,7 +34,7 @@ def generate_description(stage: "Stage") -> str:
     def is_plot_or_metric(out: "Output"):
         return bool(out.plot) or bool(out.metric)
 
-    desc: List[str] = []
+    desc: list[str] = []
 
     outs = list(filterfalse(is_plot_or_metric, stage.outs))
     if outs:
@@ -55,7 +56,7 @@ def prepare_stages_data(
     stages: Iterable["Stage"],
     description: bool = True,
     max_length: int = MAX_TEXT_LENGTH,
-) -> Dict[str, str]:
+) -> dict[str, str]:
     return {
         stage.addressing: (
             prepare_description(stage, max_length=max_length) if description else ""
@@ -67,7 +68,7 @@ def prepare_stages_data(
 class CmdStageList(CmdBase):
     def _get_stages(self) -> Iterable["Stage"]:
         if self.args.all:
-            stages: List["Stage"] = self.repo.index.stages
+            stages: list["Stage"] = self.repo.index.stages
             logger.trace("%d no. of stages found", len(stages))
             return stages
 
@@ -96,7 +97,7 @@ class CmdStageList(CmdBase):
         return 0
 
 
-def parse_cmd(commands: List[str]) -> str:
+def parse_cmd(commands: list[str]) -> str:
     """
     We need to take into account two cases:
 
@@ -284,23 +285,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 +314,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/status.py

@@ -60,18 +60,19 @@ class CmdDataStatus(CmdDataBase):
                     all_commits=self.args.all_commits,
                     with_deps=self.args.with_deps,
                     recursive=self.args.recursive,
+                    check_updates=self.args.check_updates,
                 )
             except DvcException:
                 logger.exception("")
                 return 1
 
-            if self.args.quiet:
-                return bool(st)
-
             if self.args.json:
                 ui.write_json(st)
                 return 0
 
+            if self.args.quiet:
+                return bool(st)
+
             if st:
                 self._show(st, indent)
                 return 0

dvc/commands/studio.py

@@ -1,10 +1,11 @@
-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.exceptions import DvcException
 from dvc.log import logger
 
 logger = logger.getChild(__name__)
@@ -22,12 +23,21 @@ class CmdStudioLogin(CmdConfig):
         hostname = self.args.hostname or os.environ.get(DVC_STUDIO_URL) or STUDIO_URL
         scopes = self.args.scopes
 
+        studio = self.config.get("studio", {})
+        if studio.get("url", hostname) == hostname and "token" in studio:
+            raise DvcException(
+                "Token already exists. "
+                "To login with a different token, "
+                "logout using 'dvc studio logout'."
+            )
+
+        open_browser = not self.args.no_open
         try:
-            token_name, access_token = get_access_token(
+            _, access_token = get_access_token(
                 token_name=name,
                 hostname=hostname,
                 scopes=scopes,
-                use_device_code=self.args.no_open,
+                open_browser=open_browser,
                 client_name="DVC",
             )
         except StudioAuthError as e:
@@ -35,11 +45,9 @@ class CmdStudioLogin(CmdConfig):
             return 1
 
         self.save_config(hostname, access_token)
-        ui.write(
-            "Authentication has been successfully completed."
-            "The generated token will now be accessible as"
-            f" {token_name} in the user's Studio profile."
-        )
+
+        config_path = self.config.files["global"]
+        ui.write(f"Authentication complete. Saved token to {config_path}.")
         return 0
 
     def save_config(self, hostname, token):
@@ -54,12 +62,14 @@ class CmdStudioLogout(CmdConfig):
 
         with self.config.edit("global") as conf:
             if not get_in(conf, ["studio", "token"]):
-                ui.error_write("Not logged in to Studio.")
+                ui.error_write(
+                    "Not logged in to Studio. Log in with 'dvc studio login'."
+                )
                 return 1
 
             del conf["studio"]["token"]
 
-        ui.write("Logged out from Studio")
+        ui.write("Logged out from Studio (you can log back in with 'dvc studio login')")
         return 0
 
 
@@ -92,13 +102,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 +120,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 +167,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 +179,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/compare.py

@@ -1,23 +1,15 @@
 from collections import abc
-from itertools import chain, repeat, zip_longest
-from operator import itemgetter
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
+from collections.abc import (
     ItemsView,
     Iterable,
     Iterator,
-    List,
     Mapping,
     MutableSequence,
-    Optional,
     Sequence,
-    Set,
-    Tuple,
-    Union,
-    overload,
 )
+from itertools import chain, repeat, zip_longest
+from operator import itemgetter
+from typing import TYPE_CHECKING, Any, Optional, Union, overload
 
 from funcy import reraise
 
@@ -25,7 +17,7 @@ if TYPE_CHECKING:
     from dvc.ui.table import CellT
 
 
-class Column(List["CellT"]):
+class Column(list["CellT"]):
     pass
 
 
@@ -35,13 +27,13 @@ def with_value(value, default):
 
 class TabularData(MutableSequence[Sequence["CellT"]]):
     def __init__(self, columns: Sequence[str], fill_value: Optional[str] = ""):
-        self._columns: Dict[str, Column] = {name: Column() for name in columns}
-        self._keys: List[str] = list(columns)
+        self._columns: dict[str, Column] = {name: Column() for name in columns}
+        self._keys: list[str] = list(columns)
         self._fill_value = fill_value
-        self._protected: Set[str] = set()
+        self._protected: set[str] = set()
 
     @property
-    def columns(self) -> List[Column]:
+    def columns(self) -> list[Column]:
         return list(map(self.column, self.keys()))
 
     def is_protected(self, col_name) -> bool:
@@ -60,10 +52,10 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         projection = {k: self.column(k) for k in self.keys()}
         return projection.items()
 
-    def keys(self) -> List[str]:
+    def keys(self) -> list[str]:
         return self._keys
 
-    def _iter_col_row(self, row: Sequence["CellT"]) -> Iterator[Tuple["CellT", Column]]:
+    def _iter_col_row(self, row: Sequence["CellT"]) -> Iterator[tuple["CellT", Column]]:
         for val, col in zip_longest(row, self.columns):
             if col is None:
                 break
@@ -81,7 +73,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         for val, col in self._iter_col_row(value):
             col.insert(index, val)
 
-    def __iter__(self) -> Iterator[List["CellT"]]:
+    def __iter__(self) -> Iterator[list["CellT"]]:
         return map(list, zip(*self.columns))
 
     def __getattr__(self, item: str) -> Column:
@@ -127,7 +119,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         return len(self.columns[0])
 
     @property
-    def shape(self) -> Tuple[int, int]:
+    def shape(self) -> tuple[int, int]:
         return len(self.columns), len(self)
 
     def drop(self, *col_names: str) -> None:
@@ -170,7 +162,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
             if key not in keys:
                 self.add_column(key)
 
-        row: List["CellT"] = [
+        row: list["CellT"] = [
             with_value(d.get(key), self._fill_value) for key in self.keys()
         ]
         self.append(row)
@@ -185,7 +177,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
 
     def as_dict(
         self, cols: Optional[Iterable[str]] = None
-    ) -> Iterable[Dict[str, "CellT"]]:
+    ) -> Iterable[dict[str, "CellT"]]:
         keys = self.keys() if cols is None else set(cols)
         return [{k: self._columns[k][i] for k in keys} for i in range(len(self))]
 
@@ -202,7 +194,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         if how not in ["any", "all"]:
             raise ValueError(f"Invalid 'how' value {how}. Choose one of ['any', 'all']")
 
-        match_line: Set = set()
+        match_line: set = set()
         match_any = True
         if how == "all":
             match_any = False
@@ -245,7 +237,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
             )
 
         if axis == "cols":
-            cols_to_drop: List[str] = []
+            cols_to_drop: list[str] = []
             for n_col, col in enumerate(self.columns):
                 if subset and self.keys()[n_col] not in subset:
                     continue
@@ -259,7 +251,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
 
         elif axis == "rows":
             unique_rows = []
-            rows_to_drop: List[int] = []
+            rows_to_drop: list[int] = []
             for n_row, row in enumerate(self):
                 if subset:
                     row = [
@@ -319,7 +311,7 @@ def diff_table(
 ) -> TabularData:
     a_rev = a_rev or "HEAD"
     b_rev = b_rev or "workspace"
-    headers: List[str] = ["Path", title, a_rev, b_rev, "Change"]
+    headers: list[str] = ["Path", title, a_rev, b_rev, "Change"]
     fill_value = "-"
     td = TabularData(headers, fill_value=fill_value)
 
@@ -394,7 +386,7 @@ def metrics_table(
 
     for branch, val in metrics.items():
         for fname, metric in val.get("data", {}).items():
-            row_data: Dict[str, str] = {"Revision": branch, "Path": fname}
+            row_data: dict[str, str] = {"Revision": branch, "Path": fname}
             metric = metric.get("data", {})
             flattened = (
                 flatten(format_dict(metric))

dvc/config.py

@@ -3,7 +3,7 @@ import os
 import re
 from contextlib import contextmanager
 from functools import partial
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING, Optional
 
 from funcy import compact, memoize, re_find
 
@@ -138,7 +138,7 @@ class Config(dict):
             return system_config_dir()
 
     @cached_property
-    def files(self) -> Dict[str, str]:
+    def files(self) -> dict[str, str]:
         files = {
             level: os.path.join(self.get_dir(level), self.CONFIG)
             for level in ("system", "global")
@@ -344,7 +344,7 @@ class Config(dict):
         return Schema(dirs_schema, extra=ALLOW_EXTRA)(conf)
 
     def load_config_to_level(self, level=None):
-        merged_conf: Dict = {}
+        merged_conf: dict = {}
         for merge_level in self.LEVELS:
             if merge_level == level:
                 break
@@ -390,7 +390,7 @@ class Config(dict):
 
 
 def _parse_named(conf):
-    result: Dict[str, Dict] = {"remote": {}, "machine": {}, "db": {}}
+    result: dict[str, dict] = {"remote": {}, "machine": {}, "db": {}}
 
     for section, val in conf.items():
         match = re_find(r'^\s*(remote|machine|db)\s*"(.*)"\s*$', section)

dvc/config_schema.py

@@ -181,7 +181,7 @@ SCHEMA = {
     "remote": {
         str: ByUrl(
             {
-                "": {**LOCAL_COMMON, **REMOTE_COMMON},
+                "": LOCAL_COMMON | REMOTE_COMMON,
                 "s3": {
                     "region": str,
                     "profile": str,
@@ -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": {
@@ -282,10 +285,10 @@ SCHEMA = {
                     Optional("verify", default=True): Bool,
                     **REMOTE_COMMON,
                 },
-                "http": {**HTTP_COMMON, **REMOTE_COMMON},
-                "https": {**HTTP_COMMON, **REMOTE_COMMON},
-                "webdav": {**WEBDAV_COMMON, **REMOTE_COMMON},
-                "webdavs": {**WEBDAV_COMMON, **REMOTE_COMMON},
+                "http": HTTP_COMMON | REMOTE_COMMON,
+                "https": HTTP_COMMON | REMOTE_COMMON,
+                "webdav": WEBDAV_COMMON | REMOTE_COMMON,
+                "webdavs": WEBDAV_COMMON | REMOTE_COMMON,
                 "remote": {str: object},  # Any of the above options are valid
             }
         )
@@ -318,8 +321,6 @@ SCHEMA = {
     "feature": FeatureSchema(
         {
             Optional("machine", default=False): Bool,
-            "dbt_profile": str,
-            "dbt_target": str,
         },
     ),
     "plots": {
@@ -335,6 +336,8 @@ SCHEMA = {
         "params": str,
         "plots": str,
         "live": str,
+        "auto_push": Bool,
+        "git_remote": str,
     },
     "parsing": {
         "bool": All(Lower, Choices("store_true", "boolean_optional")),
@@ -345,6 +348,7 @@ SCHEMA = {
         Exclusive("config_dir", "config_source"): str,
         Exclusive("config_module", "config_source"): str,
         "config_name": str,
+        "plugins_path": str,
     },
     "studio": {
         "token": str,

dvc/daemon.py

@@ -5,13 +5,14 @@ import logging
 import os
 import subprocess
 import sys
+from collections.abc import Mapping, Sequence
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from dvc.log import logger
 
 if TYPE_CHECKING:
-    from typing import ContextManager
+    from contextlib import AbstractContextManager
 
 from dvc.env import DVC_DAEMON, DVC_DAEMON_LOGFILE
 from dvc.utils import fix_env, is_binary
@@ -55,7 +56,7 @@ def _win_detached_subprocess(args: Sequence[str], **kwargs) -> int:
     return popen.pid
 
 
-def _get_dvc_args() -> List[str]:
+def _get_dvc_args() -> list[str]:
     args = [sys.executable]
     if not is_binary():
         root_dir = os.path.abspath(os.path.dirname(__file__))
@@ -106,12 +107,7 @@ def _posix_detached_subprocess(args: Sequence[str], **kwargs) -> int:
         os.close(read_end)
         return int(pid_str)
 
-    proc = subprocess.Popen(
-        args,
-        shell=False,  # noqa: S603
-        close_fds=True,
-        **kwargs,
-    )
+    proc = subprocess.Popen(args, shell=False, close_fds=True, **kwargs)  # noqa: S603
     os.close(read_end)
     os.write(write_end, str(proc.pid).encode("utf8"))
     os.close(write_end)
@@ -136,7 +132,7 @@ def _map_log_level_to_flag() -> Optional[str]:
     return flags.get(logger.getEffectiveLevel())
 
 
-def daemon(args: List[str]) -> None:
+def daemon(args: list[str]) -> None:
     """Launch a `dvc daemon` command in a detached process.
 
     Args:
@@ -148,12 +144,12 @@ def daemon(args: List[str]) -> None:
 
 
 def _spawn(
-    args: List[str],
-    executable: Optional[Union[str, List[str]]] = None,
+    args: list[str],
+    executable: Optional[Union[str, list[str]]] = None,
     env: Optional[Mapping[str, str]] = None,
     output_file: Optional[str] = None,
 ) -> int:
-    file: "ContextManager[Any]" = nullcontext()
+    file: "AbstractContextManager[Any]" = nullcontext()
     kwargs = {}
     if output_file:
         file = open(output_file, "ab")  # noqa: SIM115
@@ -168,9 +164,9 @@ def _spawn(
         return _detached_subprocess(executable + args, env=env, **kwargs)
 
 
-def daemonize(args: List[str], executable: Union[None, str, List[str]] = None) -> None:
+def daemonize(args: list[str], executable: Union[None, str, list[str]] = None) -> None:
     if os.name not in ("posix", "nt"):
-        return None
+        return
 
     if os.environ.get(DVC_DAEMON):
         logger.debug("skipping launching a new daemon.")

dvc/data_cloud.py

@@ -1,6 +1,7 @@
 """Manages dvc remotes that user can use with push/pull/status commands."""
 
-from typing import TYPE_CHECKING, Iterable, Optional, Set, Tuple
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Optional
 
 from dvc.config import NoRemoteError, RemoteConfigError
 from dvc.log import logger
@@ -50,7 +51,7 @@ class Remote:
 
 def _split_legacy_hash_infos(
     hash_infos: Iterable["HashInfo"],
-) -> Tuple[Set["HashInfo"], Set["HashInfo"]]:
+) -> tuple[set["HashInfo"], set["HashInfo"]]:
     from dvc.cachemgr import LEGACY_HASH_NAMES
 
     legacy = set()

dvc/database.py

@@ -0,0 +1,134 @@
+import os
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from tempfile import NamedTemporaryFile
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+
+from sqlalchemy import create_engine  # type: ignore[import]
+from sqlalchemy.engine import make_url as _make_url  # type: ignore[import]
+from sqlalchemy.exc import NoSuchModuleError  # type: ignore[import]
+
+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  # type: ignore[import]
+
+
+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

@@ -1,10 +1,12 @@
 from collections import defaultdict
-from typing import Any, Dict, List, Mapping, Set
+from collections.abc import Mapping
+from typing import Any
 
 from dvc.output import ARTIFACT_SCHEMA, DIR_FILES_SCHEMA, Output
 
 from .base import Dependency
-from .db import DB_SCHEMA, PARAM_DB, DbDependency, DbtDependency
+from .dataset import DatasetDependency
+from .db import DbDependency
 from .param import ParamsDependency
 from .repo import RepoDependency
 
@@ -15,7 +17,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,18 +25,17 @@ 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 DatasetDependency.is_dataset(p):
+        return DatasetDependency(stage, p, info)
     if repo:
         return RepoDependency(repo, stage, p, info)
     return Dependency(stage, p, info, **kwargs)
@@ -58,23 +59,15 @@ def loads_from(stage, s_list, erepo=None, fs_config=None, db=None):
     info = {RepoDependency.PARAM_REPO: erepo} if erepo else {}
     if db:
         info.update({"db": db})
-    return [
-        _get(
-            stage,
-            s,
-            info.copy(),
-            fs_config=fs_config,
-        )
-        for s in s_list
-    ]
+    return [_get(stage, s, info.copy(), fs_config=fs_config) for s in s_list]
 
 
-def _merge_params(s_list) -> Dict[str, List[str]]:
+def _merge_params(s_list) -> dict[str, list[str]]:
     d = defaultdict(list)
     default_file = ParamsDependency.DEFAULT_PARAMS_FILE
 
     # figure out completely tracked params file, and ignore specific keys
-    wholly_tracked: Set[str] = set()
+    wholly_tracked: set[str] = set()
     for key in s_list:
         if not isinstance(key, dict):
             continue

dvc/dependency/base.py

@@ -1,5 +1,3 @@
-from typing import Dict, Type
-
 from dvc.exceptions import DvcException
 from dvc.fs import download as fs_download
 from dvc.output import Output
@@ -25,11 +23,11 @@ class DependencyIsStageFileError(DvcException):
 class Dependency(Output):
     IS_DEPENDENCY = True
 
-    DoesNotExistError: Type[DvcException] = DependencyDoesNotExistError
-    IsNotFileOrDirError: Type[DvcException] = DependencyIsNotFileOrDirError
-    IsStageFileError: Type[DvcException] = DependencyIsStageFileError
+    DoesNotExistError: type[DvcException] = DependencyDoesNotExistError
+    IsNotFileOrDirError: type[DvcException] = DependencyIsNotFileOrDirError
+    IsStageFileError: type[DvcException] = DependencyIsStageFileError
 
-    def workspace_status(self) -> Dict[str, str]:
+    def workspace_status(self) -> dict[str, str]:
         if self.fs.version_aware:
             old_fs_path = self.fs_path
             try:

dvc/dependency/dataset.py

@@ -0,0 +1,84 @@
+from typing import TYPE_CHECKING, Any, ClassVar
+from urllib.parse import urlparse
+
+from funcy import compact, merge
+
+from dvc.exceptions import DvcException
+from dvc_data.hashfile.hash_info import HashInfo
+
+from .db import AbstractDependency
+
+if TYPE_CHECKING:
+    from dvc.stage import Stage
+
+
+class DatasetDependency(AbstractDependency):
+    PARAM_DATASET = "dataset"
+    DATASET_SCHEMA: ClassVar[dict] = {PARAM_DATASET: dict}
+
+    def __init__(self, stage: "Stage", p, info, *args, **kwargs):
+        super().__init__(stage, info, *args, **kwargs)
+        self.def_path = p
+        self.name = urlparse(p).netloc
+        dataset_info = info.get(self.PARAM_DATASET) or {}
+        self.hash_info = HashInfo(self.PARAM_DATASET, dataset_info)  # type: ignore[arg-type]
+        self.hash_name = self.PARAM_DATASET
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.def_path!r})"
+
+    def __str__(self):
+        return self.def_path
+
+    @classmethod
+    def is_dataset(cls, p: str):
+        return urlparse(p).scheme == "ds"
+
+    @property
+    def protocol(self):
+        return None
+
+    def dumpd(self, **kwargs):
+        return compact({self.PARAM_PATH: self.def_path, **self.hash_info.to_dict()})
+
+    def fill_values(self, values=None):
+        """Load params values dynamically."""
+        self.hash_info = HashInfo(
+            self.PARAM_DATASET, merge(self.hash_info.value, values or {})
+        )
+
+    def workspace_status(self):
+        ds = self.repo.datasets[self.name]
+        if not ds.lock:
+            return {str(self): "not in sync"}
+
+        info: dict[str, Any] = self.hash_info.value if self.hash_info else {}  # type: ignore[assignment]
+        lock = self.repo.datasets._lock_from_info(info)
+        if not lock:
+            return {str(self): "new"}
+        if lock != ds.lock:
+            return {str(self): "modified"}
+        return {}
+
+    def status(self):
+        return self.workspace_status()
+
+    def get_hash(self):
+        ds = self.repo.datasets[self.name]
+        if not ds.lock:
+            if ds._invalidated:
+                raise DvcException(
+                    "Dataset information is not in sync. "
+                    f"Run 'dvc ds update {self.name}' to sync."
+                )
+            raise DvcException("Dataset information missing from dvc.lock file")
+        return HashInfo(self.PARAM_DATASET, ds.lock.to_dict())  # type: ignore[arg-type]
+
+    def save(self):
+        self.hash_info = self.get_hash()
+
+    def download(self, to, jobs=None):
+        raise NotImplementedError
+
+    def update(self, rev=None):
+        raise NotImplementedError

dvc/dependency/db.py

@@ -1,58 +1,22 @@
-import os
-from contextlib import contextmanager, nullcontext
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, 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.repo import Repo
     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
@@ -69,8 +33,8 @@ def download_progress(to: "Output") -> Iterator[Callable[[int], Any]]:
 class AbstractDependency(Dependency):
     """Dependency without workspace/fs/fs_path"""
 
-    def __init__(self, stage: "Stage", info: Dict[str, Any], *args, **kwargs):
-        self.repo = stage.repo
+    def __init__(self, stage: "Stage", info: dict[str, Any], *args, **kwargs):
+        self.repo: "Repo" = stage.repo
         self.stage = stage
         self.fs = None
         self.fs_path = None
@@ -79,13 +43,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: ClassVar[dict] = {
+        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 +74,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 +86,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 +99,23 @@ 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:
-                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"
+        file_format = file_format or self.db_info.get(self.PARAM_FILE_FORMAT, "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 client(config) as db:
+            msg = "Testing connection"
+            with log_durations(logger.debug, msg), ui.status(msg) as status:
+                db.test_connection(onerror=status.stop)
+            with download_progress(to) as progress:
+                db.export(sql, to.fs_path, format=file_format, progress=progress)

dvc/dependency/param.py

@@ -1,7 +1,7 @@
 import os
 import typing
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import dpath
 
@@ -37,7 +37,7 @@ class BadParamFileError(DvcException):
 def read_param_file(
     fs: "FileSystem",
     path: str,
-    key_paths: Optional[List[str]] = None,
+    key_paths: Optional[list[str]] = None,
     flatten: bool = False,
     **load_kwargs,
 ) -> Any:
@@ -76,10 +76,7 @@ class ParamsDependency(Dependency):
         self.params = list(params) if params else []
         hash_info = HashInfo()
         if isinstance(params, dict):
-            hash_info = HashInfo(
-                self.PARAM_PARAMS,
-                params,  # type: ignore[arg-type]
-            )
+            hash_info = HashInfo(self.PARAM_PARAMS, params)  # type: ignore[arg-type]
         repo = repo or stage.repo
         path = path or os.path.join(repo.root_dir, self.DEFAULT_PARAMS_FILE)
         super().__init__(stage, path, repo=repo)
@@ -103,14 +100,11 @@ class ParamsDependency(Dependency):
         for param in self.params:
             if param in values:
                 info[param] = values[param]
-        self.hash_info = HashInfo(
-            self.PARAM_PARAMS,
-            info,  # type: ignore[arg-type]
-        )
+        self.hash_info = HashInfo(self.PARAM_PARAMS, info)  # type: ignore[arg-type]
 
     def read_params(
         self, flatten: bool = True, **kwargs: typing.Any
-    ) -> Dict[str, typing.Any]:
+    ) -> dict[str, typing.Any]:
         try:
             self.validate_filepath()
         except MissingParamsFile:
@@ -134,7 +128,7 @@ class ParamsDependency(Dependency):
 
         from funcy import ldistinct
 
-        status: Dict[str, Any] = defaultdict(dict)
+        status: dict[str, Any] = defaultdict(dict)
         info = self.hash_info.value if self.hash_info else {}
         assert isinstance(info, dict)
         actual = self.read_params()

dvc/dependency/repo.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
 import voluptuous as vol
 
@@ -20,7 +20,7 @@ class RepoDependency(Dependency):
     PARAM_CONFIG = "config"
     PARAM_REMOTE = "remote"
 
-    REPO_SCHEMA = {
+    REPO_SCHEMA: ClassVar[dict] = {
         PARAM_REPO: {
             vol.Required(PARAM_URL): str,
             PARAM_REV: str,
@@ -30,7 +30,7 @@ class RepoDependency(Dependency):
         }
     }
 
-    def __init__(self, def_repo: Dict[str, Any], stage: "Stage", *args, **kwargs):
+    def __init__(self, def_repo: dict[str, Any], stage: "Stage", *args, **kwargs):
         self.def_repo = def_repo
         super().__init__(stage, *args, **kwargs)
 
@@ -65,7 +65,7 @@ class RepoDependency(Dependency):
             self.def_repo[self.PARAM_REV_LOCK] = rev
 
     @classmethod
-    def _dump_def_repo(cls, def_repo) -> Dict[str, str]:
+    def _dump_def_repo(cls, def_repo) -> dict[str, str]:
         repo = {cls.PARAM_URL: def_repo[cls.PARAM_URL]}
 
         rev = def_repo.get(cls.PARAM_REV)
@@ -85,7 +85,7 @@ class RepoDependency(Dependency):
             repo[cls.PARAM_REMOTE] = remote
         return repo
 
-    def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]:
+    def dumpd(self, **kwargs) -> dict[str, Union[str, dict[str, str]]]:
         return {
             self.PARAM_PATH: self.def_path,
             self.PARAM_REPO: self._dump_def_repo(self.def_repo),

dvc/dvcfile.py

@@ -1,16 +1,6 @@
 import contextlib
 import os
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, TypeVar, Union
 
 from dvc.exceptions import DvcException
 from dvc.log import logger
@@ -133,7 +123,7 @@ class FileMixin:
         d, _ = self._load(**kwargs)
         return d
 
-    def _load(self, **kwargs: Any) -> Tuple[Any, str]:
+    def _load(self, **kwargs: Any) -> tuple[Any, str]:
         # it raises the proper exceptions by priority:
         # 1. when the file doesn't exists
         # 2. filename is not a DVC file
@@ -156,7 +146,7 @@ class FileMixin:
 
         return validate(d, cls.SCHEMA, path=fname)  # type: ignore[arg-type]
 
-    def _load_yaml(self, **kwargs: Any) -> Tuple[Any, str]:
+    def _load_yaml(self, **kwargs: Any) -> tuple[Any, str]:
         from dvc.utils import strictyaml
 
         return strictyaml.load(
@@ -181,10 +171,12 @@ class SingleStageFile(FileMixin):
     from dvc.schema import COMPILED_SINGLE_STAGE_SCHEMA as SCHEMA
     from dvc.stage.loader import SingleStageLoader as LOADER  # noqa: N814
 
-    metrics: List[str] = []
-    plots: Any = {}
-    params: List[str] = []
-    artifacts: Dict[str, Optional[Dict[str, Any]]] = {}
+    datasets: ClassVar[list[dict[str, Any]]] = []
+    datasets_lock: ClassVar[list[dict[str, Any]]] = []
+    metrics: ClassVar[list[str]] = []
+    plots: ClassVar[Any] = {}
+    params: ClassVar[list[str]] = []
+    artifacts: ClassVar[dict[str, Optional[dict[str, Any]]]] = {}
 
     @property
     def stage(self) -> "Stage":
@@ -249,6 +241,20 @@ class ProjectFile(FileMixin):
         if update_lock:
             self._dump_lockfile(stage, **kwargs)
 
+    def dump_dataset(self, dataset):
+        with modify_yaml(self.path, fs=self.repo.fs) as data:
+            datasets: list[dict] = data.setdefault("datasets", [])
+            loc = next(
+                (i for i, ds in enumerate(datasets) if ds["name"] == dataset["name"]),
+                None,
+            )
+            if loc is not None:
+                apply_diff(dataset, datasets[loc])
+                datasets[loc] = dataset
+            else:
+                datasets.append(dataset)
+        self.repo.scm_context.track_file(self.relpath)
+
     def _dump_lockfile(self, stage, **kwargs):
         self._lockfile.dump(stage, **kwargs)
 
@@ -283,11 +289,11 @@ class ProjectFile(FileMixin):
         raise DvcException("ProjectFile has multiple stages. Please specify it's name.")
 
     @cached_property
-    def contents(self) -> Dict[str, Any]:
+    def contents(self) -> dict[str, Any]:
         return self._load()[0]
 
     @cached_property
-    def lockfile_contents(self) -> Dict[str, Any]:
+    def lockfile_contents(self) -> dict[str, Any]:
         return self._lockfile.load()
 
     @cached_property
@@ -302,7 +308,7 @@ class ProjectFile(FileMixin):
         return self.LOADER(self, self.contents, self.lockfile_contents)
 
     @property
-    def metrics(self) -> List[str]:
+    def metrics(self) -> list[str]:
         return self.contents.get("metrics", [])
 
     @property
@@ -310,11 +316,19 @@ class ProjectFile(FileMixin):
         return self.contents.get("plots", {})
 
     @property
-    def params(self) -> List[str]:
+    def params(self) -> list[str]:
         return self.contents.get("params", [])
 
     @property
-    def artifacts(self) -> Dict[str, Optional[Dict[str, Any]]]:
+    def datasets(self) -> list[dict[str, Any]]:
+        return self.contents.get("datasets", [])
+
+    @property
+    def datasets_lock(self) -> list[dict[str, Any]]:
+        return self.lockfile_contents.get("datasets", [])
+
+    @property
+    def artifacts(self) -> dict[str, Optional[dict[str, Any]]]:
         return self.contents.get("artifacts", {})
 
     def remove(self, force=False):
@@ -362,6 +376,24 @@ class Lockfile(FileMixin):
             self._check_gitignored()
             return {}, ""
 
+    def dump_dataset(self, dataset: dict):
+        with modify_yaml(self.path, fs=self.repo.fs) as data:
+            data.update({"schema": "2.0"})
+            if not data:
+                logger.info("Generating lock file '%s'", self.relpath)
+
+            datasets: list[dict] = data.setdefault("datasets", [])
+            loc = next(
+                (i for i, ds in enumerate(datasets) if ds["name"] == dataset["name"]),
+                None,
+            )
+            if loc is not None:
+                datasets[loc] = dataset
+            else:
+                datasets.append(dataset)
+            data.setdefault("stages", {})
+        self.repo.scm_context.track_file(self.relpath)
+
     def dump(self, stage, **kwargs):
         stage_data = serialize.to_lockfile(stage, **kwargs)
 

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"