deps: bump dvc-data to >=2.22.5 (#10141)

* deps: bump dvc-data to >=2.22.4

* Update test_odb.py

.github/workflows/codeql.yml

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

.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.5'
+    rev: 'v0.1.6'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

dvc/cli/parser.py

@@ -42,6 +42,7 @@ from dvc.commands import (
     repro,
     root,
     stage,
+    studio,
     unprotect,
     update,
     version,
@@ -93,6 +94,7 @@ COMMANDS = [
     check_ignore,
     data,
     artifacts,
+    studio,
 ]
 
 

dvc/commands/config.py

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

dvc/commands/imp_db.py

@@ -1,13 +1,68 @@
 import argparse
 
+from funcy import compact, merge
+
 from dvc.cli import completion
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.log import logger
+from dvc.ui import ui
 
 logger = logger.getChild(__name__)
 
 
+class CmdTestDb(CmdBase):
+    def run(self):
+        from dvc.database import get_adapter
+        from dvc.database.dbt_utils import DBT_PROJECT_FILE, is_dbt_project
+        from dvc.dependency.db import _get_dbt_config
+        from dvc.exceptions import DvcException
+
+        connection = self.args.connection
+
+        db_config = self.repo.config.get("db", {})
+        config = db_config.get(connection, {})
+        if connection and not config:
+            raise DvcException(f"connection {connection} not found in config")
+
+        cli_config = compact(
+            {
+                "url": self.args.url,
+                "username": self.args.username,
+                "password": self.args.password,
+            }
+        )
+        conn_config = merge(config, cli_config)
+
+        cli_dbt_config = {"profile": self.args.profile, "target": self.args.target}
+        dbt_config = merge(_get_dbt_config(self.repo.config), cli_dbt_config)
+
+        project_dir = self.repo.root_dir
+        if not (conn_config or dbt_config):
+            if is_dbt_project(project_dir):
+                ui.write("Using", DBT_PROJECT_FILE, "for testing", styled=True)
+            else:
+                raise DvcException(
+                    f"no config set and {DBT_PROJECT_FILE} is missing; "
+                    "provide arguments or set a configuration"
+                )
+
+        adapter = get_adapter(conn_config, project_dir=project_dir, **dbt_config)
+        with adapter as db:
+            ui.write(f"Testing with {db}", styled=True)
+
+            creds = getattr(db, "creds", {})
+            for k, v in creds.items():
+                ui.write("\t", f"{k}:", v, styled=True)
+
+            if creds:
+                ui.write()
+
+            db.test_connection()
+
+        ui.write("Connection successful", styled=True)
+
+
 class CmdImportDb(CmdBase):
     def run(self):
         if not (self.args.sql or self.args.model):
@@ -24,6 +79,7 @@ class CmdImportDb(CmdBase):
             output_format=self.args.output_format,
             out=self.args.out,
             force=self.args.force,
+            connection=self.args.connection,
         )
         return 0
 
@@ -86,5 +142,26 @@ def add_parser(subparsers, parent_parser):
         default=False,
         help="Override destination file or folder if exists.",
     )
+    import_parser.add_argument(
+        "--conn",
+        dest="connection",
+        nargs="?",
+        help="Database connection to use, needs to be set in config",
+    )
 
     import_parser.set_defaults(func=CmdImportDb)
+
+    TEST_DB_HELP = "Test the database connection"
+    test_db_parser = subparsers.add_parser(
+        "test-db",
+        parents=[parent_parser],
+        description=append_doc_link(TEST_DB_HELP, "test-db"),
+        add_help=False,
+    )
+    test_db_parser.add_argument("--conn", dest="connection")
+    test_db_parser.add_argument("--url")
+    test_db_parser.add_argument("--password")
+    test_db_parser.add_argument("--username")
+    test_db_parser.add_argument("--profile")
+    test_db_parser.add_argument("--target")
+    test_db_parser.set_defaults(func=CmdTestDb)

dvc/commands/plots.py

@@ -59,49 +59,6 @@ def _show_json(
     ui.write_json(compact({"errors": all_errors, "data": data}), highlight=False)
 
 
-def _adjust_vega_renderers(renderers):
-    from dvc.render import REVISION_FIELD, VERSION_FIELD
-    from dvc_render import VegaRenderer
-
-    for r in renderers:
-        if isinstance(r, VegaRenderer):
-            if _data_versions_count(r) > 1:
-                summary = _summarize_version_infos(r)
-                for dp in r.datapoints:
-                    vi = dp.pop(VERSION_FIELD, {})
-                    keys = list(vi.keys())
-                    for key in keys:
-                        if not (len(summary.get(key, set())) > 1):
-                            vi.pop(key)
-                    if vi:
-                        dp["rev"] = "::".join(vi.values())
-            else:
-                for dp in r.datapoints:
-                    dp[REVISION_FIELD] = dp[VERSION_FIELD]["revision"]
-                    dp.pop(VERSION_FIELD, {})
-
-
-def _summarize_version_infos(renderer):
-    from collections import defaultdict
-
-    from dvc.render import VERSION_FIELD
-
-    result = defaultdict(set)
-
-    for dp in renderer.datapoints:
-        for key, value in dp.get(VERSION_FIELD, {}).items():
-            result[key].add(value)
-    return dict(result)
-
-
-def _data_versions_count(renderer):
-    from itertools import product
-
-    summary = _summarize_version_infos(renderer)
-    x = product(summary.get("filename", {None}), summary.get("field", {None}))
-    return len(set(x))
-
-
 class CmdPlots(CmdBase):
     def _func(self, *args, **kwargs):
         raise NotImplementedError
@@ -175,11 +132,10 @@ class CmdPlots(CmdBase):
                 return 0
 
             renderers = [r.renderer for r in renderers_with_errors]
-            _adjust_vega_renderers(renderers)
             if self.args.show_vega:
                 renderer = first(filter(lambda r: r.TYPE == "vega", renderers))
                 if renderer:
-                    ui.write_json(renderer.get_filled_template(as_string=False))
+                    ui.write_json(renderer.get_filled_template())
                 return 0
 
             output_file: Path = (Path.cwd() / out).resolve() / "index.html"

dvc/commands/studio.py

@@ -0,0 +1,175 @@
+import argparse
+import os
+
+from funcy import get_in
+
+from dvc.cli.utils import append_doc_link, fix_subparsers
+from dvc.commands.config import CmdConfig
+from dvc.log import logger
+
+logger = logger.getChild(__name__)
+
+
+class CmdStudioLogin(CmdConfig):
+    def run(self):
+        from dvc_studio_client.auth import StudioAuthError, get_access_token
+
+        from dvc.env import DVC_STUDIO_URL
+        from dvc.ui import ui
+        from dvc.utils.studio import STUDIO_URL
+
+        name = self.args.name
+        hostname = self.args.hostname or os.environ.get(DVC_STUDIO_URL) or STUDIO_URL
+        scopes = self.args.scopes
+
+        try:
+            token_name, access_token = get_access_token(
+                token_name=name,
+                hostname=hostname,
+                scopes=scopes,
+                use_device_code=self.args.no_open,
+                client_name="DVC",
+            )
+        except StudioAuthError as e:
+            ui.error_write(str(e))
+            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."
+        )
+        return 0
+
+    def save_config(self, hostname, token):
+        with self.config.edit("global") as conf:
+            conf["studio"]["token"] = token
+            conf["studio"]["url"] = hostname
+
+
+class CmdStudioLogout(CmdConfig):
+    def run(self):
+        from dvc.ui import ui
+
+        with self.config.edit("global") as conf:
+            if not get_in(conf, ["studio", "token"]):
+                ui.error_write("Not logged in to Studio.")
+                return 1
+
+            del conf["studio"]["token"]
+
+        ui.write("Logged out from Studio")
+        return 0
+
+
+class CmdStudioToken(CmdConfig):
+    def run(self):
+        from dvc.ui import ui
+
+        conf = self.config.read("global")
+        token = get_in(conf, ["studio", "token"])
+        if not token:
+            ui.error_write("Not logged in to Studio.")
+            return 1
+
+        ui.write(token)
+        return 0
+
+
+def add_parser(subparsers, parent_parser):
+    STUDIO_HELP = "Commands to authenticate DVC with Iterative Studio"
+    STUDIO_DESCRIPTION = (
+        "Authenticate DVC with Studio and set the token."
+        " Once this token has been properly configured,\n"
+        " DVC will utilize it for seamlessly sharing live experiments\n"
+        " and sending notifications to Studio regarding any experiments"
+        " that have been pushed."
+    )
+
+    studio_parser = subparsers.add_parser(
+        "studio",
+        parents=[parent_parser],
+        description=append_doc_link(STUDIO_DESCRIPTION, "studio"),
+        help=STUDIO_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    studio_subparser = studio_parser.add_subparsers(
+        dest="cmd",
+        help="Use `DVC studio CMD --help` to display command-specific help.",
+    )
+    fix_subparsers(studio_subparser)
+
+    STUDIO_LOGIN_HELP = "Authenticate DVC with Studio host"
+    STUDIO_LOGIN_DESCRIPTION = (
+        "By default, this command authenticates the DVC with Studio\n"
+        " using default scopes and assigns a random name as the token name."
+    )
+    login_parser = studio_subparser.add_parser(
+        "login",
+        parents=[parent_parser],
+        description=append_doc_link(STUDIO_LOGIN_DESCRIPTION, "studio/login"),
+        help=STUDIO_LOGIN_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    login_parser.add_argument(
+        "-H",
+        "--hostname",
+        action="store",
+        default=None,
+        help="The hostname of the Studio instance to authenticate with.",
+    )
+    login_parser.add_argument(
+        "-s",
+        "--scopes",
+        action="store",
+        default=None,
+        help="The scopes for the authentication token. ",
+    )
+
+    login_parser.add_argument(
+        "-n",
+        "--name",
+        action="store",
+        default=None,
+        help="The name of the authentication token. It will be used to\n"
+        "identify token shown in Studio profile.",
+    )
+
+    login_parser.add_argument(
+        "--no-open",
+        action="store_true",
+        default=False,
+        help="Use authentication flow based on user code.\n"
+        "You will be presented with user code to enter in browser.\n"
+        "DVC will also use this if it cannot launch browser on your behalf.",
+    )
+    login_parser.set_defaults(func=CmdStudioLogin)
+
+    STUDIO_LOGOUT_HELP = "Logout user from Studio"
+    STUDIO_LOGOUT_DESCRIPTION = (
+        "This removes the studio token from your global config.\n"
+    )
+
+    logout_parser = studio_subparser.add_parser(
+        "logout",
+        parents=[parent_parser],
+        description=append_doc_link(STUDIO_LOGOUT_DESCRIPTION, "studio/logout"),
+        help=STUDIO_LOGOUT_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    logout_parser.set_defaults(func=CmdStudioLogout)
+
+    STUDIO_TOKEN_HELP = "View the token dvc uses to contact Studio"  # noqa: S105 # nosec B105
+
+    logout_parser = studio_subparser.add_parser(
+        "token",
+        parents=[parent_parser],
+        description=append_doc_link(STUDIO_TOKEN_HELP, "studio/token"),
+        help=STUDIO_TOKEN_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    logout_parser.set_defaults(func=CmdStudioToken)

dvc/config.py

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

dvc/config_schema.py

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

dvc/database/__init__.py

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

dvc/database/dbt_models.py

@@ -0,0 +1,40 @@
+from typing import List, Optional
+
+from dvc.log import logger
+
+from .dbt_runner import dbt_show
+from .dbt_utils import check_dbt
+from .serializer import AgateSerializer
+
+logger = logger.getChild(__name__)
+
+
+def ref(
+    name: str,
+    package: Optional[str] = None,
+    version: Optional[int] = None,
+) -> str:
+    parts: List[str] = []
+    if package:
+        parts.append(repr(package))
+
+    parts.append(repr(name))
+    if version:
+        parts.append(f"{version=}")
+
+    inner = ",".join(parts)
+    return "{{ ref(" + inner + ") }}"
+
+
+@check_dbt("model")
+def get_model(
+    name: str,
+    package: Optional[str] = None,
+    version: Optional[int] = None,
+    profile: Optional[str] = None,
+    target: Optional[str] = None,
+) -> AgateSerializer:
+    model = ref(name, package, version=version)
+    q = f"select * from {model}"  # noqa: S608
+    table = dbt_show(inline=q, profile=profile, target=target)
+    return AgateSerializer(table)

dvc/database/dbt_query.py

@@ -0,0 +1,88 @@
+import os
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
+
+from attrs import define, field
+
+from dvc.exceptions import DvcException
+from dvc.log import logger
+
+from .dbt_utils import check_dbt
+from .serializer import AgateSerializer, PandasSQLSerializer
+
+if TYPE_CHECKING:
+    from dbt.adapters.base.impl import BaseAdapter
+
+logger = logger.getChild(__name__)
+
+
+@define
+class DbtAdapter:
+    adapter: "BaseAdapter" = field(repr=lambda o: type(o).__qualname__)
+    creds: Dict[str, Any] = field(repr=False)
+
+    @contextmanager
+    def query(self, sql: str) -> Iterator[Union[AgateSerializer, PandasSQLSerializer]]:
+        from dbt.adapters.sql import SQLAdapter
+
+        with self.adapter.connection_named("execute"):
+            if isinstance(self.adapter, SQLAdapter):
+                # try to use connection to stream results
+                # otherwise fallback to fetch all at once
+                try:
+                    connection = self.adapter.connections.get_thread_connection().handle
+                except Exception:  # noqa: BLE001
+                    logger.debug("", exc_info=True)
+                else:
+                    # pandas raises warning that it does not support
+                    # dbapi2-based raw connection, but it does seem to work.
+                    yield PandasSQLSerializer(sql, connection)
+                    return
+            _, table = self.adapter.execute(sql, fetch=True)
+            yield AgateSerializer(table)
+
+    def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None:
+        with self.adapter.connection_named("debug"):
+            try:
+                self.adapter.debug_query()
+            except Exception as exc:
+                if callable(onerror):
+                    onerror()
+
+                logger.exception(
+                    "dvc was unable to connect to the specified database. "
+                    "Please check your database credentials and try again.",
+                    exc_info=False,
+                )
+                raise DvcException("The database returned the following error") from exc
+
+
+@contextmanager
+@check_dbt("query")
+def adapter(
+    project_dir: Optional[str] = None,
+    profiles_dir: Optional[str] = None,
+    profile: Optional[str] = None,
+    target: Optional[str] = None,
+) -> Iterator["DbtAdapter"]:
+    from dbt.adapters import factory as adapters_factory
+    from dbt.adapters.base.impl import BaseAdapter
+
+    from .dbt_utils import get_or_build_profile, get_profiles_dir, init_dbt
+
+    profiles_dir = profiles_dir or get_profiles_dir(project_dir)
+    flags = init_dbt(profiles_dir, os.getcwd(), target=target)
+
+    with flags, adapters_factory.adapter_management():
+        # likely invalid connection profile or no adapter
+        profile_obj = get_or_build_profile(project_dir, profile, target)
+
+        adapters_factory.register_adapter(profile_obj)  # type: ignore[arg-type]
+        adapter = adapters_factory.get_adapter(profile_obj)  # type: ignore[arg-type]
+        assert isinstance(adapter, BaseAdapter)
+        try:
+            creds = dict(profile_obj.credentials.connection_info())
+        except:  # noqa: E722
+            creds = {}
+
+        yield DbtAdapter(adapter, creds)

dvc/database/dbt_runner.py

@@ -0,0 +1,61 @@
+from typing import TYPE_CHECKING, List, Optional, Sequence, Union
+
+from dvc.exceptions import DvcException
+from dvc.log import logger
+
+if TYPE_CHECKING:
+    from agate import Table
+    from dbt.contracts.results import RunResult
+
+logger = logger.getChild(__name__)
+
+
+class DbtInternalError(DvcException):
+    pass
+
+
+def _kw_to_cmd_args(**kwargs: Union[None, bool, str, int]) -> List[str]:
+    args: List[str] = []
+    for key, value in kwargs.items():
+        key = key.replace("_", "-")
+        if value is None:
+            continue  # skip creating a flag in this case
+        if value is True:
+            args.append(f"--{key}")
+        elif value is False:
+            args.append(f"--no-{key}")
+        else:
+            args.extend([f"--{key}", str(value)])
+    return args
+
+
+def _dbt_invoke(*posargs: str, quiet: bool = True, **kw: Union[None, bool, str, int]):
+    from dbt.cli.main import dbtRunner
+
+    args = _kw_to_cmd_args(quiet=quiet or None)  # global options
+    args.extend([*posargs, *_kw_to_cmd_args(**kw)])
+
+    runner = dbtRunner()
+    result = runner.invoke(args)
+    if result.success:
+        return result.result
+    raise DbtInternalError(f"failed to run dbt {posargs[0]}") from result.exception
+
+
+def dbt_show(
+    inline: Optional[str] = None,
+    limit: int = -1,
+    profile: Optional[str] = None,
+    target: Optional[str] = None,
+) -> "Table":
+    from dbt.contracts.results import RunExecutionResult
+
+    result = _dbt_invoke(
+        "show", inline=inline, limit=limit, profile=profile, target=target
+    )
+    assert isinstance(result, RunExecutionResult)
+
+    run_results: Sequence["RunResult"] = result.results
+    run_result, *_ = run_results
+    assert run_result.agate_table is not None
+    return run_result.agate_table

dvc/database/dbt_utils.py

@@ -0,0 +1,100 @@
+import os
+from contextlib import contextmanager
+from importlib.util import find_spec
+from typing import TYPE_CHECKING, Optional
+
+from funcy import cut_prefix
+
+from dvc.exceptions import DvcException
+from dvc.log import logger
+from dvc.utils import packaging
+
+if TYPE_CHECKING:
+    from dbt.config.profile import Profile
+
+logger = logger.getChild(__name__)
+
+DBT_PROJECT_FILE = "dbt_project.yml"
+
+
+@contextmanager
+def check_dbt(action: Optional[str] = None):
+    if not (find_spec("dbt") and find_spec("dbt.cli")):
+        action = f" {action}" if action else ""
+        raise DvcException(f"Could not run{action}. dbt-core is not installed")
+    yield packaging.check_required_version(pkg="dbt-core")
+
+
+@contextmanager
+def init_dbt(
+    profiles_dir: str,
+    project_dir: str,
+    target: Optional[str] = None,
+):
+    from argparse import Namespace
+
+    from dbt import flags
+
+    prev = flags.get_flags()
+    try:
+        args = Namespace(
+            use_colors=True,
+            project_dir=project_dir,
+            profiles_dir=profiles_dir,
+            target=target,
+        )
+        flags.set_from_args(args, None)
+        yield flags.get_flags()
+    finally:
+        flags.set_flags(prev)
+
+
+def get_profiles_dir(project_dir: Optional[str] = None) -> str:
+    from dbt.cli.resolvers import default_profiles_dir
+
+    if profiles_dir := os.getenv("DBT_PROFILES_DIR"):
+        return profiles_dir
+    if project_dir and os.path.isfile(os.path.join(project_dir, "profiles.yml")):
+        return project_dir
+    return os.fspath(default_profiles_dir())
+
+
+@contextmanager
+def _handle_profile_parsing_error():
+    from dbt.exceptions import DbtRuntimeError
+
+    try:
+        yield
+    except DbtRuntimeError as e:
+        cause = e.__cause__ is not None and e.__cause__.__context__
+        if isinstance(cause, ModuleNotFoundError) and (
+            adapter := cut_prefix(cause.name, "dbt.adapters.")
+        ):
+            # DbtRuntimeError is very noisy, so send it to debug
+            logger.debug("", exc_info=True)
+            raise DvcException(f"dbt-{adapter} dependency is missing") from cause
+        raise DvcException("failed to read connection profiles") from e
+
+
+def is_dbt_project(project_dir: str):
+    return os.path.isfile(os.path.join(project_dir, DBT_PROJECT_FILE))
+
+
+@_handle_profile_parsing_error()
+def get_or_build_profile(
+    project_dir: Optional[str], profile: Optional[str], target: Optional[str]
+) -> "Profile":
+    from dbt.config.profile import Profile
+    from dbt.config.renderer import ProfileRenderer
+    from dbt.config.runtime import load_profile
+
+    if project_dir and is_dbt_project(project_dir):
+        return load_profile(
+            project_dir, {}, profile_name_override=profile, target_override=target
+        )
+
+    if not profile:
+        raise DvcException("No profile specified to query from.")
+
+    renderer = ProfileRenderer({})
+    return Profile.render(renderer, profile, target_override=target)

dvc/database/serializer.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import sqlalchemy as sa
+    from agate import Table
+
+
+def noop(_):
+    pass
+
+
+@dataclass
+class PandasSQLSerializer:
+    sql: "Union[sa.TextClause, str]"
+    con: "sa.Connection"
+    chunksize: int = 10_000
+
+    def to_csv(self, file: str, progress=noop) -> None:
+        import pandas as pd
+
+        with open(file, mode="wb") as f:
+            idfs = pd.read_sql_query(self.sql, self.con, chunksize=self.chunksize)
+            for i, df in enumerate(idfs):
+                df.to_csv(f, header=i == 0, index=False)
+                progress(len(df))
+
+    def to_json(self, file: str, progress=noop) -> None:  # noqa: ARG002
+        import pandas as pd
+
+        df = pd.read_sql_query(self.sql, self.con)
+        df.to_json(file, orient="records")
+
+
+@dataclass
+class AgateSerializer:
+    table: "Table"
+
+    def to_csv(self, file: str, progress=noop) -> None:  # noqa: ARG002
+        return self.table.to_csv(file)
+
+    def to_json(self, file: str, progress=noop) -> None:  # noqa: ARG002
+        return self.table.to_json(file)
+
+
+def export(
+    serializer: Union[PandasSQLSerializer, AgateSerializer],
+    file: str,
+    format: str = "csv",  # noqa: A002
+    progress=noop,
+) -> None:
+    if format == "csv":
+        return serializer.to_csv(file, progress=progress)
+    return serializer.to_json(file, progress=progress)

dvc/database/sqla.py

@@ -0,0 +1,73 @@
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterator, Optional, Union
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import URL, Engine
+from sqlalchemy.engine import make_url as _make_url
+from sqlalchemy.exc import NoSuchModuleError
+
+from dvc.exceptions import DvcException
+from dvc.log import logger
+
+from .serializer import PandasSQLSerializer
+
+logger = logger.getChild(__name__)
+
+
+def make_url(url: Union[URL, str], **kwargs: Any) -> URL:
+    return _make_url(url).set(**kwargs)
+
+
+def url_from_config(config: Union[str, URL, Dict[str, str]]) -> URL:
+    if isinstance(config, (str, URL)):
+        return make_url(config)
+    return make_url(**config)
+
+
+@dataclass
+class SQLAlchemyAdapter:
+    engine: Engine
+
+    @contextmanager
+    def query(self, sql: str) -> Iterator[PandasSQLSerializer]:
+        with self.engine.connect().execution_options(stream_results=True) as conn:
+            yield PandasSQLSerializer(text(sql), conn)
+
+    def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None:
+        try:
+            with self.engine.connect() as conn:
+                conn.execute(text("select 1"))
+        except Exception as exc:
+            if callable(onerror):
+                onerror()
+            logger.exception(
+                "dvc was unable to connect to the specified database. "
+                "Please check your database credentials and try again.",
+                exc_info=False,
+            )
+            raise DvcException("The database returned the following error") from exc
+
+
+@contextmanager
+def handle_error(url: URL):
+    try:
+        yield
+    except (ModuleNotFoundError, NoSuchModuleError) as e:
+        # TODO: write installation instructions
+        raise DvcException(
+            f"Could not load database driver for {url.drivername!r}"
+        ) from e
+
+
+@contextmanager
+def adapter(
+    url_or_config: Union[str, URL, Dict[str, str]], **engine_kwargs: Any
+) -> Iterator[SQLAlchemyAdapter]:
+    url = url_from_config(url_or_config)
+    with handle_error(url):
+        engine = create_engine(url, **engine_kwargs)
+    try:
+        yield SQLAlchemyAdapter(engine)
+    finally:
+        engine.dispose()

dvc/dependency/db.py

@@ -1,8 +1,8 @@
 import os
 from contextlib import contextmanager, nullcontext
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
 
-from funcy import compact
+from funcy import compact, log_durations
 
 from dvc.exceptions import DvcException
 from dvc.log import logger
@@ -11,9 +11,9 @@ from dvc.scm import SCM
 from .base import Dependency
 
 if TYPE_CHECKING:
-    from agate import Table
     from rich.status import Status
 
+    from dvc.output import Output
     from dvc.stage import Stage
 
 logger = logger.getChild(__name__)
@@ -24,20 +24,23 @@ PARAM_PROFILE = "profile"
 PARAM_FILE_FORMAT = "file_format"
 
 
-def _get_db_config(config: Dict) -> Dict:
+def _get_dbt_config(config: Dict) -> Dict:
     conf = config.get("feature", {})
-    pref = "db_"
+    pref = "dbt_"
     return {k.lstrip(pref): v for k, v in conf.items() if k.startswith(pref)}
 
 
 @contextmanager
-def log_status(msg, log=logger.debug) -> Iterator["Status"]:
+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), ui.status(msg) as status:
-        yield status
+    with log_durations(log, msg), status or ui.status(msg) as st:
+        st.update(msg)
+        yield st
 
 
 @contextmanager
@@ -50,15 +53,23 @@ def chdir(path):
         os.chdir(wdir)
 
 
-def export_to(table: "Table", to: str, file_format: str = "csv") -> None:
-    exporter = {"csv": table.to_csv, "json": table.to_json}
-    return exporter[file_format](to)
+@contextmanager
+def download_progress(to: "Output") -> Iterator[Callable[[int], Any]]:
+    from dvc.ui import ui
+    from dvc.ui._rich_progress import DbDownloadProgress
+
+    with log_durations(logger.debug, f"Saving to {to}"), DbDownloadProgress(
+        console=ui.error_console,
+    ) as progress:
+        task = progress.add_task("Saving", total=None, output=to)
+        yield lambda n: progress.advance(task, advance=n)
+        progress.update(task, description="Saved", total=0)
 
 
 class AbstractDependency(Dependency):
     """Dependency without workspace/fs/fs_path"""
 
-    def __init__(self, stage: "Stage", info, *args, **kwargs):
+    def __init__(self, stage: "Stage", info: Dict[str, Any], *args, **kwargs):
         self.repo = stage.repo
         self.stage = stage
         self.fs = None
@@ -69,7 +80,8 @@ class AbstractDependency(Dependency):
 
 class DbDependency(AbstractDependency):
     PARAM_QUERY = "query"
-    QUERY_SCHEMA = {PARAM_QUERY: str}
+    PARAM_CONNECTION = "connection"
+    QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str}
 
     def __init__(self, stage: "Stage", info, *args, **kwargs):
         super().__init__(stage, info, *args, **kwargs)
@@ -103,33 +115,47 @@ class DbDependency(AbstractDependency):
     def update(self, rev=None):
         """nothing to update."""
 
-    def download(self, to, jobs=None, file_format=None):  # noqa: ARG002
+    def download(
+        self,
+        to: "Output",
+        jobs: Optional[int] = None,  # noqa: ARG002
+        file_format: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        from dvc.database import export, get_adapter
+
         db_info = self.info.get(PARAM_DB, {})
         query = db_info.get(self.PARAM_QUERY)
         if not query:
             raise DvcException("Cannot download: no query specified")
 
-        from dvc.utils.db import _check_dbt, _profiles_dir, execute_sql
+        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 = _get_db_config(self.repo.config)
-        profile = db_info.get(PARAM_PROFILE) or db_config.get(PARAM_PROFILE)
-        target = self.target or db_config.get("target")
-        file_format = file_format or db_info.get(PARAM_FILE_FORMAT, "csv")
+        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")
 
-        _check_dbt(self.PARAM_QUERY)
-        profiles_dir = _profiles_dir(self.repo.root_dir)
-        with log_status("Executing query") as status:
-            table = execute_sql(
-                query,
-                profiles_dir,
-                self.repo.root_dir,
-                profile,
-                target=target,
-                status=status,
-            )
-        # NOTE: we keep everything in memory, and then export it out later.
-        with log_status(f"Saving to {to}"):
-            return export_to(table, to.fs_path, file_format)
+        project_dir = self.repo.root_dir
+        with get_adapter(
+            config, project_dir=project_dir, profile=profile, target=target
+        ) as db:
+            logger.debug("using adapter: %s", db)
+            with log_status("Testing connection") as status:
+                db.test_connection(onerror=status.stop)
+
+            file_format = file_format or db_info.get(PARAM_FILE_FORMAT, "csv")
+            assert file_format
+            with log_status("Executing query") as status, db.query(query) as serializer:
+                status.stop()
+                logger.debug("using serializer: %s", serializer)
+                with download_progress(to) as progress:
+                    return export(
+                        serializer, to.fs_path, format=file_format, progress=progress
+                    )
 
 
 class DbtDependency(AbstractDependency):
@@ -237,9 +263,12 @@ class DbtDependency(AbstractDependency):
             rev = self.rev
         self.def_repo[RepoDependency.PARAM_REV_LOCK] = self._get_clone(rev).get_rev()
 
-    def download(self, to, jobs=None, file_format=None):  # noqa: ARG002
-        from dvc.ui import ui
-
+    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, "")
@@ -255,27 +284,34 @@ class DbtDependency(AbstractDependency):
         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)
-        ui.write(f"Saved file to {to}", styled=True)
 
-    def _download_db(self, to, version=None, file_format=None):
-        from dvc.utils.db import get_model
+    def _download_db(
+        self,
+        to: "Output",
+        version: Optional[int] = None,
+        file_format: Optional[str] = None,
+    ) -> None:
+        from dvc.database import export, get_model
 
         db_info = self.info.get(PARAM_DB, {})
         model = db_info.get(self.PARAM_MODEL)
+        version = version or db_info.get(self.PARAM_VERSION)
+        file_format = file_format or db_info.get(PARAM_FILE_FORMAT) or "csv"
+        assert file_format
         if not model:
             raise DvcException("Cannot download, no model specified")
 
-        db_config = _get_db_config(self.repo.config)
-        version = version or db_info.get(self.PARAM_VERSION)
-        profile = db_info.get(PARAM_PROFILE) or db_config.get(PARAM_PROFILE)
-        target = self.target or db_info.get("target")
-        file_format = file_format or db_info.get(PARAM_FILE_FORMAT, "csv")
+        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"):
-            table = get_model(model, version=version, profile=profile, target=target)
+            serializer = get_model(
+                model, version=version, profile=profile, target=target
+            )
         # NOTE: we keep everything in memory, and then export it out later.
-        with log_status(f"Saving to {to}"):
-            export_to(table, to.fs_path, file_format=file_format)
+        with download_progress(to) as progress:
+            export(serializer, to.fs_path, format=file_format, progress=progress)
 
 
 DB_SCHEMA = {

dvc/fs/__init__.py

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

dvc/render/__init__.py

@@ -1,7 +1,8 @@
-INDEX_FIELD = "step"
-REVISION_FIELD = "rev"
-FILENAME_FIELD = "filename"
-VERSION_FIELD = "dvc_data_version_info"
-REVISIONS_KEY = "revisions"
+INDEX = "step"
+REVISION = "rev"
+FILENAME = "filename"
+FIELD = "field"
+REVISIONS = "revisions"
+ANCHOR_DEFINITIONS = "anchor_definitions"
 TYPE_KEY = "type"
-SRC_FIELD = "src"
+SRC = "src"

dvc/render/convert.py

@@ -1,7 +1,6 @@
-from collections import defaultdict
 from typing import Dict, List, Union
 
-from dvc.render import REVISION_FIELD, REVISIONS_KEY, SRC_FIELD, TYPE_KEY, VERSION_FIELD
+from dvc.render import REVISION, REVISIONS, SRC, TYPE_KEY
 from dvc.render.converter.image import ImageConverter
 from dvc.render.converter.vega import VegaConverter
 
@@ -19,39 +18,31 @@ def _get_converter(
     raise ValueError(f"Invalid renderer class {renderer_class}")
 
 
-def _group_by_rev(datapoints):
-    grouped = defaultdict(list)
-    for datapoint in datapoints:
-        rev = datapoint.get(VERSION_FIELD, {}).get("revision")
-        grouped[rev].append(datapoint)
-    return dict(grouped)
-
-
 def to_json(renderer, split: bool = False) -> List[Dict]:
     if renderer.TYPE == "vega":
-        grouped = _group_by_rev(renderer.datapoints)
+        if not renderer.datapoints:
+            return []
+        revs = renderer.get_revs()
         if split:
-            content = renderer.get_filled_template(
-                skip_anchors=["data"], as_string=False
-            )
+            content, split_content = renderer.get_partial_filled_template()
         else:
-            content = renderer.get_filled_template(as_string=False)
-        if grouped:
-            return [
-                {
-                    TYPE_KEY: renderer.TYPE,
-                    REVISIONS_KEY: sorted(grouped.keys()),
-                    "content": content,
-                    "datapoints": grouped,
-                }
-            ]
-        return []
+            content = renderer.get_filled_template()
+            split_content = {}
+
+        return [
+            {
+                TYPE_KEY: renderer.TYPE,
+                REVISIONS: revs,
+                "content": content,
+                **split_content,
+            }
+        ]
     if renderer.TYPE == "image":
         return [
             {
                 TYPE_KEY: renderer.TYPE,
-                REVISIONS_KEY: [datapoint.get(REVISION_FIELD)],
-                "url": datapoint.get(SRC_FIELD),
+                REVISIONS: [datapoint.get(REVISION)],
+                "url": datapoint.get(SRC),
             }
             for datapoint in renderer.datapoints
         ]

dvc/render/converter/image.py

@@ -2,7 +2,7 @@ import base64
 import os
 from typing import TYPE_CHECKING, Any, Dict, List, Tuple
 
-from dvc.render import FILENAME_FIELD, REVISION_FIELD, SRC_FIELD
+from dvc.render import FILENAME, REVISION, SRC
 
 from . import Converter
 
@@ -58,9 +58,9 @@ class ImageConverter(Converter):
             else:
                 src = self._encode_image(image_data)
             datapoint = {
-                REVISION_FIELD: revision,
-                FILENAME_FIELD: filename,
-                SRC_FIELD: src,
+                REVISION: revision,
+                FILENAME: filename,
+                SRC: src,
             }
             datapoints.append(datapoint)
         return datapoints, properties

dvc/render/converter/vega.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 from funcy import first, last
 
 from dvc.exceptions import DvcException
-from dvc.render import FILENAME_FIELD, INDEX_FIELD, VERSION_FIELD
+from dvc.render import FIELD, FILENAME, INDEX, REVISION
 
 from . import Converter
 
@@ -112,7 +112,6 @@ class VegaConverter(Converter):
     ):
         super().__init__(plot_id, data, properties)
         self.plot_id = plot_id
-        self.inferred_properties: Dict = {}
 
     def _infer_y_from_data(self):
         if self.plot_id in self.data:
@@ -120,34 +119,38 @@ class VegaConverter(Converter):
                 if all(isinstance(item, dict) for item in lst):
                     datapoint = first(lst)
                     field = last(datapoint.keys())
-                    self.inferred_properties["y"] = {self.plot_id: field}
-                    break
+                    return {self.plot_id: field}
+        return None
 
     def _infer_x_y(self):
         x = self.properties.get("x", None)
         y = self.properties.get("y", None)
 
+        inferred_properties: Dict = {}
+
         # Infer x.
         if isinstance(x, str):
-            self.inferred_properties["x"] = {}
+            inferred_properties["x"] = {}
             # If multiple y files, duplicate x for each file.
             if isinstance(y, dict):
                 for file, fields in y.items():
                     # Duplicate x for each y.
                     if isinstance(fields, list):
-                        self.inferred_properties["x"][file] = [x] * len(fields)
+                        inferred_properties["x"][file] = [x] * len(fields)
                     else:
-                        self.inferred_properties["x"][file] = x
+                        inferred_properties["x"][file] = x
             # Otherwise use plot ID as file.
             else:
-                self.inferred_properties["x"][self.plot_id] = x
+                inferred_properties["x"][self.plot_id] = x
 
         # Infer y.
         if y is None:
-            self._infer_y_from_data()
+            inferred_properties["y"] = self._infer_y_from_data()
         # If y files not provided, use plot ID as file.
         elif not isinstance(y, dict):
-            self.inferred_properties["y"] = {self.plot_id: y}
+            inferred_properties["y"] = {self.plot_id: y}
+
+        return inferred_properties
 
     def _find_datapoints(self):
         result = {}
@@ -182,7 +185,7 @@ class VegaConverter(Converter):
 
         x = properties.get("x", None)
         if not isinstance(x, dict):
-            return INDEX_FIELD
+            return INDEX
 
         fields = {field for _, field in _file_field(x)}
         if len(fields) == 1:
@@ -192,7 +195,7 @@ class VegaConverter(Converter):
     def flat_datapoints(self, revision):  # noqa: C901, PLR0912
         file2datapoints, properties = self.convert()
 
-        props_update = {}
+        props_update: Dict[str, Union[str, List[Dict[str, str]]]] = {}
 
         xs = list(_get_xs(properties, file2datapoints))
 
@@ -200,15 +203,17 @@ class VegaConverter(Converter):
         if not xs:
             x_file, x_field = (
                 None,
-                INDEX_FIELD,
+                INDEX,
             )
         else:
             x_file, x_field = xs[0]
-        props_update["x"] = x_field
+
+        num_xs = len(xs)
+        multiple_x_fields = num_xs > 1 and len({x[1] for x in xs}) > 1
+        props_update["x"] = "dvc_inferred_x_value" if multiple_x_fields else x_field
 
         ys = list(_get_ys(properties, file2datapoints))
 
-        num_xs = len(xs)
         num_ys = len(ys)
         if num_xs > 1 and num_xs != num_ys:
             raise DvcException(
@@ -237,6 +242,14 @@ class VegaConverter(Converter):
         else:
             common_prefix_len = 0
 
+        props_update["anchors_y_definitions"] = [
+            {
+                FILENAME: _get_short_y_file(y_file, common_prefix_len),
+                FIELD: y_field,
+            }
+            for y_file, y_field in ys
+        ]
+
         for i, (y_file, y_field) in enumerate(ys):
             if num_xs > 1:
                 x_file, x_field = xs[i]
@@ -249,15 +262,16 @@ class VegaConverter(Converter):
                     source_field=y_field,
                 )
 
-            if x_field == INDEX_FIELD and x_file is None:
-                _update_from_index(datapoints, INDEX_FIELD)
+            if x_field == INDEX and x_file is None:
+                _update_from_index(datapoints, INDEX)
             else:
                 x_datapoints = file2datapoints.get(x_file, [])
                 try:
                     _update_from_field(
                         datapoints,
-                        field=x_field,
+                        field="dvc_inferred_x_value" if multiple_x_fields else x_field,
                         source_datapoints=x_datapoints,
+                        source_field=x_field,
                     )
                 except IndexError:
                     raise DvcException(  # noqa: B904
@@ -266,15 +280,12 @@ class VegaConverter(Converter):
                         "They have to have same length."
                     )
 
-            y_file_short = y_file[common_prefix_len:].strip("/\\")
             _update_all(
                 datapoints,
                 update_dict={
-                    VERSION_FIELD: {
-                        "revision": revision,
-                        FILENAME_FIELD: y_file_short,
-                        "field": y_field,
-                    }
+                    REVISION: revision,
+                    FILENAME: _get_short_y_file(y_file, common_prefix_len),
+                    FIELD: y_field,
                 },
             )
 
@@ -295,10 +306,10 @@ class VegaConverter(Converter):
         generated datapoints and updated properties. `x`, `y` values and labels
         are inferred and always provided.
         """
-        self._infer_x_y()
+        inferred_properties = self._infer_x_y()
 
         datapoints = self._find_datapoints()
-        properties = {**self.properties, **self.inferred_properties}
+        properties = {**self.properties, **inferred_properties}
 
         properties["y_label"] = self.infer_y_label(properties)
         properties["x_label"] = self.infer_x_label(properties)
@@ -306,6 +317,10 @@ class VegaConverter(Converter):
         return datapoints, properties
 
 
+def _get_short_y_file(y_file, common_prefix_len):
+    return y_file[common_prefix_len:].strip("/\\")
+
+
 def _update_from_field(
     target_datapoints: List[Dict],
     field: str,

dvc/render/match.py

@@ -80,7 +80,7 @@ def match_defs_renderers(  # noqa: C901, PLR0912
     for plot_id, group in plots_data.group_definitions().items():
         plot_datapoints: List[Dict] = []
         props = _squash_plots_properties(group)
-        final_props: Dict = {}
+        first_props: Dict = {}
 
         def_errors: Dict[str, Exception] = {}
         src_errors: DefaultDict[str, Dict[str, Exception]] = defaultdict(dict)
@@ -90,6 +90,7 @@ def match_defs_renderers(  # noqa: C901, PLR0912
         if templates_dir is not None:
             props["template_dir"] = templates_dir
 
+        revs = []
         for rev, inner_id, plot_definition in group:
             plot_sources = infer_data_sources(inner_id, plot_definition)
             definitions_data = plots_data.get_definition_data(plot_sources, rev)
@@ -109,19 +110,24 @@ def match_defs_renderers(  # noqa: C901, PLR0912
 
             try:
                 dps, rev_props = converter.flat_datapoints(rev)
+                if dps and rev not in revs:
+                    revs.append(rev)
             except Exception as e:  # noqa: BLE001
                 logger.warning("In %r, %s", rev, str(e).lower())
                 def_errors[rev] = e
                 continue
 
-            if not final_props and rev_props:
-                final_props = rev_props
+            if not first_props and rev_props:
+                first_props = rev_props
             plot_datapoints.extend(dps)
 
-        if "title" not in final_props:
-            final_props["title"] = renderer_id
+        if "title" not in first_props:
+            first_props["title"] = renderer_id
+
+        if revs:
+            first_props["revs_with_datapoints"] = revs
 
         if renderer_cls is not None:
-            renderer = renderer_cls(plot_datapoints, renderer_id, **final_props)
+            renderer = renderer_cls(plot_datapoints, renderer_id, **first_props)
             renderers.append(RendererWithErrors(renderer, dict(src_errors), def_errors))
     return renderers

dvc/repo/commit.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Set
 
 from dvc import prompt
 
@@ -7,6 +7,7 @@ from .scm_context import scm_context
 
 if TYPE_CHECKING:
     from . import Repo
+    from .index import IndexView
 
 
 def _prepare_message(stage, changes):
@@ -88,14 +89,7 @@ def commit_2_to_3(repo: "Repo", dry: bool = False):
         outs_filter=lambda o: o.hash_name == "md5-dos2unix",
         recursive=True,
     )
-    migrated = set()
-    for out in view.outs:
-        dvcfile = out.stage.dvcfile.relpath
-        if isinstance(out.stage.dvcfile, ProjectFile):
-            lockfile = out.stage.dvcfile._lockfile.relpath
-            migrated.add(f"{dvcfile} ({lockfile})")
-        else:
-            migrated.add(dvcfile)
+    migrated = _migrateable_dvcfiles(view)
     if not migrated:
         ui.write("No DVC files in the repo to migrate to the 3.0 format.")
         return
@@ -110,11 +104,41 @@ def commit_2_to_3(repo: "Repo", dry: bool = False):
             for out in stage.filter_outs(filter_info)
             if outs_filter is not None and outs_filter(out)
         }
+        modified = False
         if outs:
             for out in outs:
                 out.update_legacy_hash_name(force=True)
+            modified = True
+        deps = {dep for dep in stage.deps if not stage.is_import and dep.is_in_repo}
+        if deps:
+            for dep in deps:
+                dep.update_legacy_hash_name(force=True)
+            modified = True
+        if modified:
             stage.save(allow_missing=True)
             stage.commit(allow_missing=True, relink=True)
             if not isinstance(stage.dvcfile, ProjectFile):
                 ui.write(f"Updating DVC file '{stage.dvcfile.relpath}'")
             stage.dump(update_pipeline=False)
+
+
+def _migrateable_dvcfiles(view: "IndexView") -> Set[str]:
+    from dvc.dvcfile import ProjectFile
+
+    migrated = set()
+    for stage, filter_info in view._stage_infos:
+        outs_filter = view._outs_filter
+        dvcfile = stage.dvcfile.relpath
+        assert outs_filter
+        if any(outs_filter(out) for out in stage.filter_outs(filter_info)) or (
+            not stage.is_import
+            and any(
+                dep.is_in_repo and dep.hash_name == "md5-dos2unix" for dep in stage.deps
+            )
+        ):
+            if isinstance(stage.dvcfile, ProjectFile):
+                lockfile = stage.dvcfile._lockfile.relpath
+                migrated.add(f"{dvcfile} ({lockfile})")
+            else:
+                migrated.add(dvcfile)
+    return migrated

dvc/repo/imp_db.py

@@ -27,6 +27,7 @@ def imp_db(  # noqa: PLR0913
     output_format: str = "csv",
     out: Optional[str] = None,
     force: bool = False,
+    connection: Optional[str] = None,
 ):
     ui.warn("WARNING: import-db is an experimental feature.")
     ui.warn(
@@ -50,7 +51,7 @@ def imp_db(  # noqa: PLR0913
         db.update({"model": model, "version": version, "project_dir": project_dir})
     else:
         out = out or f"results.{output_format}"
-        db["query"] = sql
+        db.update({"query": sql, "connection": connection})
 
     out = resolve_output(url or ".", out, force=force)
     path, wdir, out = resolve_paths(self, out, always_local=True)

dvc/scm.py

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

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

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

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

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

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

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

dvc/testing/benchmarks/fixtures.py

@@ -5,7 +5,6 @@ from subprocess import check_output
 from typing import Dict, Optional
 
 import pytest
-import virtualenv
 from dulwich.porcelain import clone
 from funcy import first
 from packaging import version
@@ -24,6 +23,8 @@ class VirtualEnv:
         self.bin = self.path / ("Scripts" if os.name == "nt" else "bin")
 
     def create(self) -> None:
+        import virtualenv
+
         virtualenv.cli_run([os.fspath(self.path)])
 
     def run(self, cmd: str, *args: str, env: Optional[Dict[str, str]] = None) -> None:
@@ -192,7 +193,7 @@ def make_dataset(request, bench_config, tmp_dir, dvc_bench_git_repo):
 
         path = tmp_dir / "dataset"
         root = dvc_bench_git_repo
-        src = root / "data" / bench_config.size / "dataset"
+        src = root / "data" / bench_config.dataset / "dataset"
         src_dvc = src.with_suffix(".dvc")
 
         dvc = Repo(root)

dvc/testing/benchmarks/plugin.py

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

dvc/ui/_rich_progress.py

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

dvc/utils/db.py

@@ -1,252 +0,0 @@
-import os
-from contextlib import contextmanager
-from importlib.util import find_spec
-from typing import (
-    TYPE_CHECKING,
-    List,
-    Optional,
-    Sequence,
-    Union,
-)
-
-from funcy import cut_prefix, identity
-
-from dvc.exceptions import DvcException
-from dvc.log import logger
-
-from . import packaging
-
-if TYPE_CHECKING:
-    from agate import Table
-    from dbt.config.profile import Profile
-    from dbt.contracts.results import RunResult
-    from rich.status import Status
-
-logger = logger.getChild(__name__)
-
-
-class DbtInternalError(DvcException):
-    pass
-
-
-def _check_dbt(action: Optional[str]):
-    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")
-    packaging.check_required_version(pkg="dbt-core")
-
-
-@contextmanager
-def check_dbt(action: Optional[str]):
-    _check_dbt(action)
-    yield
-
-
-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 + ") }}"
-
-
-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
-
-
-@check_dbt("model")
-def get_model(
-    name: str,
-    package: Optional[str] = None,
-    version: Optional[int] = None,
-    profile: Optional[str] = None,
-    target: Optional[str] = None,
-) -> "Table":
-    model = _ref(name, package, version=version)
-    q = f"select * from {model}"  # noqa: S608
-    return _dbt_show(
-        inline=q,
-        profile=profile,
-        target=target,
-    )
-
-
-def _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 _global_dbt_flags(
-    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_profile_or(
-    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 os.path.isfile(os.path.join(project_dir, "dbt_project.yml")):
-        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)
-
-
-@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
-
-
-@contextmanager
-def _suppress_status(status: Optional["Status"]):
-    if status:
-        status.stop()
-
-    yield
-    if status:
-        status.start()
-
-
-@check_dbt("query")
-def execute_sql(
-    sql: str,
-    profiles_dir: str,
-    project_dir: Optional[str],
-    profile: Optional[str],
-    target: Optional[str] = None,
-    status: Optional["Status"] = None,
-) -> "Table":
-    from dbt.adapters import factory as adapters_factory
-    from dbt.adapters.sql import SQLAdapter
-
-    flags = _global_dbt_flags(profiles_dir, os.getcwd(), target=target)
-    update_status = status.update if status else identity
-
-    with flags, adapters_factory.adapter_management():
-        # likely invalid connection profile or no adapter
-        with _handle_profile_parsing_error(), _suppress_status(status):
-            profile_obj = _get_profile_or(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, SQLAdapter)
-        with adapter.connection_named("debug"):
-            update_status("Testing Connection")
-            try:
-                adapter.debug_query()
-            except Exception as exc:
-                if status:
-                    status.stop()
-                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
-
-        with adapter.connection_named("execute"):
-            update_status("Executing query")
-            exec_resp = adapter.execute(sql, fetch=True)
-            _, table = exec_resp
-            return table

pyproject.toml

@@ -29,15 +29,16 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
 ]
 dynamic = ["version"]
+
 dependencies = [
     "colorama>=0.3.9",
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=2.22.1,<2.23.0",
+    "dvc-data>=2.22.5,<2.23.0",
     "dvc-http>=2.29.0",
-    "dvc-render>=0.3.1,<1",
-    "dvc-studio-client>=0.13.0,<1",
+    "dvc-render>=1.0.0,<2",
+    "dvc-studio-client>=0.17.1,<1",
     "dvc-task>=0.3.0,<1",
     "flatten_dict<1,>=0.4.1",
     # https://github.com/iterative/dvc/issues/9654
@@ -58,7 +59,7 @@ dependencies = [
     "requests>=2.22",
     "rich>=12",
     "ruamel.yaml>=0.17.11",
-    "scmrepo>=1.4.1,<2",
+    "scmrepo>=1.5.0,<2",
     "shortuuid>=0.5",
     "shtab<2,>=1.3.4",
     "tabulate>=0.8.7",
@@ -76,7 +77,8 @@ gdrive = ["dvc-gdrive==2.20"]
 gs = ["dvc-gs==2.22.1"]
 hdfs = ["dvc-hdfs==2.19"]
 lint = [
-    "mypy==1.7.0",
+    "mypy==1.7.1",
+    "pandas-stubs",
     "types-colorama",
     "types-psutil",
     "types-pyinstaller",
@@ -85,7 +87,7 @@ lint = [
     "types-toml",
     "types-tqdm",
 ]
-oss = ["dvc-oss>=2.20.0"]
+oss = ["dvc-oss>=2.22.0"]
 s3 = ["dvc-s3==2.23.0"]
 ssh = ["dvc-ssh>=2.22.1,<3"]
 ssh_gssapi = ["dvc-ssh[gssapi]>=2.22.1,<3"]
@@ -111,6 +113,8 @@ tests = [
     'pywin32>=225; sys_platform == "win32"', # optional test dependency
     'tzdata; sys_platform == "win32"',  # for testing with celery
     "dbt-sqlite", # optional dependency for testing `import-db`
+    "sqlalchemy>=1,<3", # optional dependency for `import-db`
+    "pandas>=1",
 ]
 webdav = ["dvc-webdav==2.19.1"]
 webhdfs = ["dvc-webhdfs==2.19"]
@@ -127,9 +131,6 @@ dvc = "dvc.cli:main"
 [project.entry-points."fsspec.specs"]
 dvc = "dvc.api:DVCFileSystem"
 
-[project.entry-points."pytest11"]
-dvc-testing = "dvc.testing.plugin"
-
 [project.entry-points."pyinstaller40"]
 hook-dirs = "dvc.__pyinstaller:get_hook_dirs"
 tests = "dvc.__pyinstaller:get_PyInstaller_tests"

tests/func/test_cli.py

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

tests/func/test_import_db.py

@@ -1,13 +1,25 @@
 """import_db tests without involving any databases."""
-
+import pytest
 from agate import Table
 
+from dvc.database.serializer import AgateSerializer
+
 data1 = Table([(i, i) for i in range(1, 10)], ["id", "value"])
 data2 = Table([(i, i) for i in range(10, 20)], ["id", "value"])
 
+serializers = AgateSerializer(data1), AgateSerializer(data2)
 
-def test_sql(mocker, tmp_dir, dvc):
-    mocker.patch("dvc.utils.db.execute_sql", side_effect=[data1, data2])
+
+@pytest.fixture
+def adapter(mocker):
+    m = mocker.patch("dvc.database.get_adapter")
+    adapter = mocker.MagicMock()
+    adapter.query.return_value.__enter__.side_effect = serializers
+    m.return_value.__enter__.return_value = adapter
+    return m
+
+
+def test_sql(adapter, tmp_dir, dvc):
     stage = dvc.imp_db(
         sql="select * from model", profile="profile", output_format="json"
     )
@@ -44,8 +56,46 @@ def test_sql(mocker, tmp_dir, dvc):
     }
 
 
+def test_sql_conn_string(adapter, tmp_dir, dvc):
+    with dvc.config.edit(level="local") as conf:
+        conf["db"] = {"conn": {"url": "conn"}}
+
+    stage = dvc.imp_db(
+        sql="select * from model", connection="conn", output_format="json"
+    )
+    db = {"file_format": "json", "connection": "conn", "query": "select * from model"}
+    assert (tmp_dir / stage.relpath).parse() == {
+        "md5": "f0dcca8e8411907adb25aa547d7c432c",
+        "frozen": True,
+        "deps": [{"db": db}],
+        "outs": [
+            {
+                "hash": "md5",
+                "md5": "56226d443067b195c90b427db557e1f2",
+                "path": "results.json",
+                "size": 243,
+            }
+        ],
+    }
+
+    dvc.update(stage.addressing)
+    assert (tmp_dir / stage.relpath).parse() == {
+        "md5": "8c0441b206d8ee7c7b9e3d54871c3cbb",
+        "frozen": True,
+        "deps": [{"db": db}],
+        "outs": [
+            {
+                "hash": "md5",
+                "md5": "08d806ee7e1f309c4561750805f02276",
+                "path": "results.json",
+                "size": 290,
+            }
+        ],
+    }
+
+
 def test_model(mocker, tmp_dir, dvc):
-    mocker.patch("dvc.utils.db.get_model", side_effect=[data1, data2])
+    mocker.patch("dvc.database.get_model", side_effect=serializers)
     stage = dvc.imp_db(model="model", output_format="json")
 
     assert (tmp_dir / stage.relpath).parse() == {

tests/func/test_import_db_e2e.py

@@ -12,10 +12,6 @@ from agate import Table
 
 from dvc.types import StrPath
 
-dbt_connections = pytest.importorskip(
-    "dbt.adapters.sql.connections", reason="dbt-core not installed"
-)
-SQLConnectionManager = dbt_connections.SQLConnectionManager
 pytest.importorskip("dbt.adapters.sqlite", reason="dbt-sqlite not installed")
 
 
@@ -52,6 +48,14 @@ def seed_db(db_path):
         yield inner
 
 
+@pytest.fixture
+def db_config(dvc, db_path):
+    # only needed for `sql_conn` tests
+    with dvc.config.edit(level="local") as conf:
+        conf["db"] = {"conn": {"url": f"sqlite:///{db_path.fs_path}"}}
+    return "conn"
+
+
 @pytest.fixture
 def dbt_profile(tmp_dir, scm, db_path):
     (tmp_dir / "profiles.yml").dump(
@@ -90,11 +94,14 @@ def dbt_model(dbt_project):
     return "model"
 
 
-@pytest.fixture(params=("sql", "model", "external_model"))
+@pytest.fixture(params=("sql", "model", "external_model", "sql_conn"))
 def import_db_parameters(request: pytest.FixtureRequest):
     if request.param == "sql":
         profile = request.getfixturevalue("dbt_profile")
         return {"sql": "select * from model", "profile": profile}
+    if request.param == "sql_conn":
+        conn = request.getfixturevalue("db_config")
+        return {"sql": "select * from model", "connection": conn}
 
     dbt_project = request.getfixturevalue("dbt_project")
     return {

tests/func/test_odb.py

@@ -165,7 +165,7 @@ def test_default_cache_type(dvc):
 @pytest.mark.skipif(os.name == "nt", reason="Not supported for Windows.")
 @pytest.mark.parametrize("group", [False, True])
 def test_shared_cache(tmp_dir, dvc, group):
-    from dvc.fs import system
+    from dvc_data.hashfile.db.local import umask
 
     if group:
         with dvc.config.edit() as conf:
@@ -178,7 +178,7 @@ def test_shared_cache(tmp_dir, dvc, group):
     tmp_dir.dvc_gen({"file": "file content", "dir": {"file2": "file 2 content"}})
 
     file_mode = oct(0o444)
-    dir_mode = oct(0o2775 if group else (0o777 & ~system.umask))
+    dir_mode = oct(0o2775 if group else (0o777 & ~umask))
     for root, dnames, fnames in os.walk(cache_dir):
         for dname in dnames:
             path = os.path.join(root, dname)

tests/func/test_scm.py

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

tests/integration/plots/test_plots.py

@@ -11,7 +11,7 @@ from bs4 import BeautifulSoup
 from funcy import first
 
 from dvc.cli import main
-from dvc.render import REVISION_FIELD, VERSION_FIELD
+from dvc.render import ANCHOR_DEFINITIONS, FILENAME, REVISION
 
 JSON_OUT = "vis_data"
 
@@ -102,6 +102,9 @@ def verify_vega(
     html_result,
     json_result,
     split_json_result,
+    title,
+    x_label,
+    y_label,
 ):
     if isinstance(versions, str):
         versions = [versions]
@@ -111,27 +114,60 @@ def verify_vega(
         assert j[0]["type"] == "vega"
         assert set(j[0]["revisions"]) == set(versions)
 
-    assert json_result[0]["datapoints"] == split_json_result[0]["datapoints"]
-    assert set(versions) == set(json_result[0]["datapoints"].keys())
+    assert (
+        json_result[0]["content"]["data"]["values"]
+        == split_json_result[0][ANCHOR_DEFINITIONS]["<DVC_METRIC_DATA>"]
+    )
+
+    assert set(versions) == set(json_result[0]["revisions"])
 
     assert json_result[0]["content"]["data"]["values"]
     assert html_result["data"]["values"]
-    assert split_json_result[0]["content"]["data"]["values"] == "<DVC_METRIC_DATA>"
 
-    def _assert_templates_equal(html_template, filled_template, split_template):
-        # besides data, json and split json should be equal
-        path = ["data", "values"]
+    content_str = json.dumps(split_json_result[0]["content"])
+    assert "<DVC_METRIC_DATA>" in content_str
+    assert "<DVC_METRIC_X_LABEL>" in content_str
+    assert "<DVC_METRIC_Y_LABEL>" in content_str
+
+    def _assert_templates_equal(
+        html_template, filled_template, split_template, title, x_label, y_label
+    ):
+        # besides split anchors, json and split json should be equal
+        paths = [["data", "values"], ["encoding", "color"]]
         tmp1 = deepcopy(html_template)
         tmp2 = deepcopy(filled_template)
-        tmp3 = deepcopy(split_template)
-        dpath.set(tmp1, path, {})
-        dpath.set(tmp2, path, {})
-        dpath.set(tmp3, path, {})
+        tmp3 = json.loads(
+            json.dumps(split_template)
+            .replace('"<DVC_METRIC_PLOT_HEIGHT>"', "300")
+            .replace('"<DVC_METRIC_PLOT_WIDTH>"', "300")
+            .replace("<DVC_METRIC_TITLE>", title)
+            .replace("<DVC_METRIC_X_LABEL>", x_label)
+            .replace("<DVC_METRIC_Y_LABEL>", y_label)
+            .replace(
+                '"<DVC_METRIC_ZOOM_AND_PAN>"',
+                json.dumps(
+                    {
+                        "name": "grid",
+                        "select": "interval",
+                        "bind": "scales",
+                    }
+                ),
+            )
+        )
+        for path in paths:
+            dpath.set(tmp1, path, {})
+            dpath.set(tmp2, path, {})
+            dpath.set(tmp3, path, {})
 
         assert tmp1 == tmp2 == tmp3
 
     _assert_templates_equal(
-        html_result, json_result[0]["content"], split_json_result[0]["content"]
+        html_result,
+        json_result[0]["content"],
+        split_json_result[0]["content"],
+        title,
+        x_label,
+        y_label,
     )
 
 
@@ -193,17 +229,13 @@ def test_repo_with_plots(tmp_dir, scm, dvc, capsys, run_copy_metrics, repo_with_
     ] == _update_datapoints(
         linear_v1,
         {
-            VERSION_FIELD: {
-                "revision": "workspace",
-                "filename": "linear.json",
-                "field": "y",
-            },
+            REVISION: "workspace",
         },
     )
     assert html_result["linear.json"]["data"]["values"] == _update_datapoints(
         linear_v1,
         {
-            REVISION_FIELD: "workspace",
+            REVISION: "workspace",
         },
     )
     assert json_data["confusion.json"][0]["content"]["data"][
@@ -211,27 +243,29 @@ def test_repo_with_plots(tmp_dir, scm, dvc, capsys, run_copy_metrics, repo_with_
     ] == _update_datapoints(
         confusion_v1,
         {
-            VERSION_FIELD: {
-                "revision": "workspace",
-                "filename": "confusion.json",
-                "field": "actual",
-            },
+            REVISION: "workspace",
         },
     )
     assert html_result["confusion.json"]["data"]["values"] == _update_datapoints(
         confusion_v1,
         {
-            REVISION_FIELD: "workspace",
+            REVISION: "workspace",
         },
     )
     verify_image(tmp_dir, "workspace", "image.png", image_v1, html_path, json_data)
 
-    for plot in ["linear.json", "confusion.json"]:
+    for plot, title, x_label, y_label in [
+        ("linear.json", "linear", "x", "y"),
+        ("confusion.json", "confusion matrix", "predicted", "actual"),
+    ]:
         verify_vega(
             "workspace",
             html_result[plot],
             json_data[plot],
             split_json_data[plot],
+            title,
+            x_label,
+            y_label,
         )
 
     verify_vega_props("confusion.json", json_data, **confusion_props)
@@ -250,12 +284,18 @@ def test_repo_with_plots(tmp_dir, scm, dvc, capsys, run_copy_metrics, repo_with_
     verify_image(tmp_dir, "workspace", "image.png", image_v2, html_path, json_data)
     verify_image(tmp_dir, "HEAD", "image.png", image_v1, html_path, json_data)
 
-    for plot in ["linear.json", "confusion.json"]:
+    for plot, title, x_label, y_label in [
+        ("linear.json", "linear", "x", "y"),
+        ("confusion.json", "confusion matrix", "predicted", "actual"),
+    ]:
         verify_vega(
             ["HEAD", "workspace"],
             html_result[plot],
             json_data[plot],
             split_json_data[plot],
+            title,
+            x_label,
+            y_label,
         )
     verify_vega_props("confusion.json", json_data, **confusion_props)
     path = tmp_dir / "subdir"
@@ -277,31 +317,23 @@ def test_repo_with_plots(tmp_dir, scm, dvc, capsys, run_copy_metrics, repo_with_
         ] == _update_datapoints(
             linear_v2,
             {
-                VERSION_FIELD: {
-                    "revision": "workspace",
-                    "filename": "../linear.json",
-                    "field": "y",
-                },
+                REVISION: "workspace",
             },
         ) + _update_datapoints(
             linear_v1,
             {
-                VERSION_FIELD: {
-                    "revision": "HEAD",
-                    "filename": "../linear.json",
-                    "field": "y",
-                },
+                REVISION: "HEAD",
             },
         )
         assert html_result["../linear.json"]["data"]["values"] == _update_datapoints(
             linear_v2,
             {
-                REVISION_FIELD: "workspace",
+                REVISION: "workspace",
             },
         ) + _update_datapoints(
             linear_v1,
             {
-                REVISION_FIELD: "HEAD",
+                REVISION: "HEAD",
             },
         )
         assert json_data["../confusion.json"][0]["content"]["data"][
@@ -309,43 +341,38 @@ def test_repo_with_plots(tmp_dir, scm, dvc, capsys, run_copy_metrics, repo_with_
         ] == _update_datapoints(
             confusion_v2,
             {
-                VERSION_FIELD: {
-                    "revision": "workspace",
-                    "filename": "../confusion.json",
-                    "field": "actual",
-                },
+                REVISION: "workspace",
             },
         ) + _update_datapoints(
             confusion_v1,
             {
-                VERSION_FIELD: {
-                    "revision": "HEAD",
-                    "filename": "../confusion.json",
-                    "field": "actual",
-                },
+                REVISION: "HEAD",
             },
         )
         assert html_result["../confusion.json"]["data"]["values"] == _update_datapoints(
             confusion_v2,
             {
-                REVISION_FIELD: "workspace",
+                REVISION: "workspace",
             },
         ) + _update_datapoints(
             confusion_v1,
             {
-                REVISION_FIELD: "HEAD",
+                REVISION: "HEAD",
             },
         )
 
-        for plot in [
-            "../linear.json",
-            "../confusion.json",
+        for plot, title, x_label, y_label in [
+            ("../linear.json", "linear", "x", "y"),
+            ("../confusion.json", "confusion matrix", "predicted", "actual"),
         ]:
             verify_vega(
                 ["HEAD", "workspace"],
                 html_result[plot],
                 json_data[plot],
                 split_json_data[plot],
+                title,
+                x_label,
+                y_label,
             )
         verify_image(
             path,
@@ -434,7 +461,7 @@ def test_repo_with_config_plots(tmp_dir, capsys, repo_with_config_plots):
     repo_state = repo_with_config_plots()
     plots = next(repo_state)
 
-    html_path, _, __ = call(capsys)
+    html_path, _, split_json_result = call(capsys)
 
     assert os.path.exists(html_path)
     html_result = extract_vega_specs(
@@ -444,20 +471,28 @@ def test_repo_with_config_plots(tmp_dir, capsys, repo_with_config_plots):
             "confusion_train_vs_test",
         ],
     )
+
     ble = _update_datapoints(
         plots["data"]["linear_train.json"],
         {
-            REVISION_FIELD: "linear_train.json",
+            REVISION: "workspace",
+            FILENAME: "linear_train.json",
         },
     ) + _update_datapoints(
         plots["data"]["linear_test.json"],
         {
-            REVISION_FIELD: "linear_test.json",
+            REVISION: "workspace",
+            FILENAME: "linear_test.json",
         },
     )
 
     assert html_result["linear_train_vs_test"]["data"]["values"] == ble
-    # TODO check json results once vscode is able to handle flexible plots
+    assert (
+        split_json_result["data"]["linear_train_vs_test"][0][ANCHOR_DEFINITIONS][
+            "<DVC_METRIC_DATA>"
+        ]
+        == ble
+    )
 
 
 @pytest.mark.vscode
@@ -466,7 +501,7 @@ def test_repo_with_dvclive_plots(tmp_dir, capsys, repo_with_dvclive_plots):
 
     for s in ("show", "diff"):
         _, json_result, split_json_result = call(capsys, subcommand=s)
-        expected_result = {
+        expected_result: Dict[str, Dict[str, list[str]]] = {
             "data": {
                 "dvclive/plots/metrics/metric.tsv": [],
             },

tests/unit/command/test_studio.py

@@ -0,0 +1,81 @@
+from dvc_studio_client.auth import AuthenticationExpired
+
+from dvc.cli import main
+from dvc.utils.studio import STUDIO_URL
+
+
+def test_studio_login_token_check_failed(mocker):
+    mocker.patch(
+        "dvc_studio_client.auth.get_access_token", side_effect=AuthenticationExpired
+    )
+
+    assert main(["studio", "login"]) == 1
+
+
+def test_studio_login_success(mocker, dvc):
+    mocker.patch(
+        "dvc_studio_client.auth.get_access_token",
+        return_value=("token_name", "isat_access_token"),
+    )
+
+    assert main(["studio", "login"]) == 0
+
+    config = dvc.config.load_one("global")
+    assert config["studio"]["token"] == "isat_access_token"
+    assert config["studio"]["url"] == STUDIO_URL
+
+
+def test_studio_login_arguments(mocker):
+    mock = mocker.patch(
+        "dvc_studio_client.auth.get_access_token",
+        return_value=("token_name", "isat_access_token"),
+    )
+
+    assert (
+        main(
+            [
+                "studio",
+                "login",
+                "--name",
+                "token_name",
+                "--hostname",
+                "https://example.com",
+                "--scopes",
+                "experiments",
+                "--no-open",
+            ]
+        )
+        == 0
+    )
+
+    mock.assert_called_with(
+        token_name="token_name",
+        hostname="https://example.com",
+        scopes="experiments",
+        use_device_code=True,
+        client_name="DVC",
+    )
+
+
+def test_studio_logout(dvc):
+    with dvc.config.edit("global") as conf:
+        conf["studio"]["token"] = "isat_access_token"
+
+    assert main(["studio", "logout"]) == 0
+    config = dvc.config.load_one("global")
+    assert "token" not in config["studio"]
+
+    assert main(["studio", "logout"]) == 1
+
+
+def test_studio_token(dvc, capsys):
+    with dvc.config.edit("global") as conf:
+        conf["studio"]["token"] = "isat_access_token"
+
+    assert main(["studio", "token"]) == 0
+    assert capsys.readouterr().out == "isat_access_token\n"
+
+    with dvc.config.edit("global") as conf:
+        del conf["studio"]["token"]
+
+    assert main(["studio", "token"]) == 1

tests/unit/render/test_convert.py

@@ -1,109 +1,87 @@
-from dvc.render import REVISION_FIELD, REVISIONS_KEY, SRC_FIELD, TYPE_KEY, VERSION_FIELD
+import json
+
+import pytest
+
+from dvc.render import ANCHOR_DEFINITIONS, FILENAME, REVISION, REVISIONS, SRC, TYPE_KEY
 from dvc.render.convert import to_json
 
 
 def test_to_json_vega(mocker):
     vega_renderer = mocker.MagicMock()
     vega_renderer.TYPE = "vega"
+    vega_renderer.get_revs.return_value = ["bar", "foo"]
     vega_renderer.get_filled_template.return_value = {"this": "is vega"}
-    vega_renderer.datapoints = [
-        {
-            "x": 1,
-            "y": 2,
-            VERSION_FIELD: {"revision": "foo"},
-            "filename": "foo.json",
-        },
-        {
-            "x": 2,
-            "y": 1,
-            VERSION_FIELD: {"revision": "bar"},
-            "filename": "foo.json",
-        },
-    ]
     result = to_json(vega_renderer)
     assert result[0] == {
         TYPE_KEY: vega_renderer.TYPE,
-        REVISIONS_KEY: ["bar", "foo"],
+        REVISIONS: ["bar", "foo"],
         "content": {"this": "is vega"},
-        "datapoints": {
-            "foo": [
-                {
-                    "x": 1,
-                    "y": 2,
-                    "filename": "foo.json",
-                    VERSION_FIELD: {"revision": "foo"},
-                },
-            ],
-            "bar": [
-                {
-                    "x": 2,
-                    "y": 1,
-                    "filename": "foo.json",
-                    VERSION_FIELD: {"revision": "bar"},
-                },
-            ],
-        },
     }
     vega_renderer.get_filled_template.assert_called()
 
 
+@pytest.mark.vscode
 def test_to_json_vega_split(mocker):
+    revs = ["bar", "foo"]
+    content = json.dumps(
+        {
+            "this": "is split vega",
+            "encoding": {"color": "<DVC_METRIC_COLOR>"},
+            "data": {"values": "<DVC_METRIC_DATA>"},
+        }
+    )
+    anchor_definitions = {
+        "<DVC_METRIC_COLOR>": {
+            "field": "rev",
+            "scale": {
+                "domain": revs,
+                "range": ["#ff0000", "#00ff00"],
+            },
+        },
+        "<DVC_METRIC_DATA>": [
+            {
+                "x": 1,
+                "y": 2,
+                REVISION: "foo",
+                FILENAME: "foo.json",
+            },
+            {
+                "x": 2,
+                "y": 1,
+                REVISION: "bar",
+                FILENAME: "foo.json",
+            },
+        ],
+    }
+
     vega_renderer = mocker.MagicMock()
     vega_renderer.TYPE = "vega"
-    vega_renderer.get_filled_template.return_value = {"this": "is split vega"}
-    vega_renderer.datapoints = [
-        {
-            "x": 1,
-            "y": 2,
-            VERSION_FIELD: {"revision": "foo"},
-            "filename": "foo.json",
-        },
-        {
-            "x": 2,
-            "y": 1,
-            VERSION_FIELD: {"revision": "bar"},
-            "filename": "foo.json",
-        },
-    ]
+    vega_renderer.get_partial_filled_template.return_value = (
+        content,
+        {ANCHOR_DEFINITIONS: anchor_definitions},
+    )
+    vega_renderer.get_revs.return_value = ["bar", "foo"]
+
     result = to_json(vega_renderer, split=True)
     assert result[0] == {
+        ANCHOR_DEFINITIONS: anchor_definitions,
         TYPE_KEY: vega_renderer.TYPE,
-        REVISIONS_KEY: ["bar", "foo"],
-        "content": {"this": "is split vega"},
-        "datapoints": {
-            "foo": [
-                {
-                    "x": 1,
-                    "y": 2,
-                    "filename": "foo.json",
-                    VERSION_FIELD: {"revision": "foo"},
-                }
-            ],
-            "bar": [
-                {
-                    "x": 2,
-                    "y": 1,
-                    "filename": "foo.json",
-                    VERSION_FIELD: {"revision": "bar"},
-                }
-            ],
-        },
+        REVISIONS: revs,
+        "content": content,
     }
-    vega_renderer.get_filled_template.assert_called_with(
-        as_string=False, skip_anchors=["data"]
-    )
+    vega_renderer.get_partial_filled_template.assert_called_once()
 
 
 def test_to_json_image(mocker):
     image_renderer = mocker.MagicMock()
     image_renderer.TYPE = "image"
     image_renderer.datapoints = [
-        {SRC_FIELD: "contentfoo", REVISION_FIELD: "foo"},
-        {SRC_FIELD: "contentbar", REVISION_FIELD: "bar"},
+        {SRC: "contentfoo", REVISION: "foo"},
+        {SRC: "contentbar", REVISION: "bar"},
     ]
     result = to_json(image_renderer)
     assert result[0] == {
-        "url": image_renderer.datapoints[0].get(SRC_FIELD),
-        REVISIONS_KEY: [image_renderer.datapoints[0].get(REVISION_FIELD)],
+        "url": image_renderer.datapoints[0].get(SRC),
+        REVISIONS: [image_renderer.datapoints[0].get(REVISION)],
         TYPE_KEY: image_renderer.TYPE,
     }

tests/unit/render/test_image_converter.py

@@ -1,4 +1,4 @@
-from dvc.render import REVISION_FIELD, SRC_FIELD
+from dvc.render import FILENAME, REVISION, SRC
 from dvc.render.converter.image import ImageConverter
 
 
@@ -8,9 +8,9 @@ def test_image_converter_no_out():
     datapoints, _ = converter.flat_datapoints("r")
 
     assert datapoints[0] == {
-        REVISION_FIELD: "r",
-        "filename": "image.png",
-        SRC_FIELD: converter._encode_image(b"content"),
+        REVISION: "r",
+        FILENAME: "image.png",
+        SRC: converter._encode_image(b"content"),
     }
 
 
@@ -21,9 +21,9 @@ def test_image_converter_with_out(tmp_dir):
     datapoints, _ = converter.flat_datapoints("r")
 
     assert datapoints[0] == {
-        REVISION_FIELD: "r",
-        "filename": "image.png",
-        SRC_FIELD: str(tmp_dir / "foo" / "r_image.png"),
+        REVISION: "r",
+        FILENAME: "image.png",
+        SRC: str(tmp_dir / "foo" / "r_image.png"),
     }
 
     assert (tmp_dir / "foo" / "r_image.png").read_bytes() == b"content"
@@ -37,9 +37,9 @@ def test_image_converter_with_slash_in_revision(tmp_dir):
     datapoints, _ = converter.flat_datapoints("feature/r")
 
     assert datapoints[0] == {
-        REVISION_FIELD: "feature/r",
-        "filename": "image.png",
-        SRC_FIELD: str(tmp_dir / "foo" / "feature_r_image.png"),
+        REVISION: "feature/r",
+        FILENAME: "image.png",
+        SRC: str(tmp_dir / "foo" / "feature_r_image.png"),
     }
 
     assert (tmp_dir / "foo" / "feature_r_image.png").read_bytes() == b"content"

tests/unit/render/test_match.py

@@ -1,9 +1,13 @@
 import pytest
 from funcy import set_in
 
-from dvc.render import VERSION_FIELD
+from dvc.render import FIELD, FILENAME, REVISION
 from dvc.render.converter.vega import VegaConverter
-from dvc.render.match import PlotsData, _squash_plots_properties, match_defs_renderers
+from dvc.render.match import (
+    PlotsData,
+    _squash_plots_properties,
+    match_defs_renderers,
+)
 
 
 @pytest.mark.parametrize(
@@ -158,25 +162,23 @@ def test_match_renderers(M):
     renderer = renderer_with_errors[0]
     assert renderer.datapoints == [
         {
-            VERSION_FIELD: {
-                "revision": "v1",
-                "filename": "file.json",
-                "field": "y",
-            },
+            REVISION: "v1",
+            FILENAME: "file.json",
+            FIELD: "y",
             "x": 1,
             "y": 1,
         },
         {
-            VERSION_FIELD: {
-                "revision": "v1",
-                "filename": "file.json",
-                "field": "y",
-            },
+            REVISION: "v1",
+            FILENAME: "file.json",
+            FIELD: "y",
             "x": 2,
             "y": 2,
         },
     ]
     assert renderer.properties == {
+        "anchors_y_definitions": [{FILENAME: "file.json", FIELD: "y"}],
+        "revs_with_datapoints": ["v1"],
         "title": "plot_id_1",
         "x": "x",
         "y": "y",

tests/unit/render/test_vega_converter.py

@@ -3,7 +3,7 @@ from collections import OrderedDict
 import pytest
 
 from dvc.exceptions import DvcException
-from dvc.render import VERSION_FIELD
+from dvc.render import FIELD, FILENAME, REVISION
 from dvc.render.converter.vega import FieldNotFoundError, VegaConverter, _lists
 
 
@@ -35,23 +35,25 @@ def test_finding_lists(dictionary, expected_result):
                 {
                     "v": 1,
                     "step": 0,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                 },
                 {
                     "v": 2,
                     "step": 1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                 },
             ],
-            {"x": "step", "y": "v", "x_label": "step", "y_label": "v"},
+            {
+                "anchors_y_definitions": [{FILENAME: "f", FIELD: "v"}],
+                "x": "step",
+                "y": "v",
+                "x_label": "step",
+                "y_label": "v",
+            },
             id="default_x_y",
         ),
         pytest.param(
@@ -61,23 +63,25 @@ def test_finding_lists(dictionary, expected_result):
                 {
                     "v": 1,
                     "v2": 0.1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "v": 2,
                     "v2": 0.2,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
             ],
-            {"x": "v", "y": "v2", "x_label": "v", "y_label": "v2"},
+            {
+                "anchors_y_definitions": [{FILENAME: "f", FIELD: "v2"}],
+                "x": "v",
+                "y": "v2",
+                "x_label": "v",
+                "y_label": "v2",
+            },
             id="choose_x_y",
         ),
         pytest.param(
@@ -99,23 +103,25 @@ def test_finding_lists(dictionary, expected_result):
                 {
                     "v": 1,
                     "v2": 0.1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "v": 2,
                     "v2": 0.2,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
             ],
-            {"x": "v", "y": "v2", "x_label": "x", "y_label": "y"},
+            {
+                "anchors_y_definitions": [{FILENAME: "f", FIELD: "v2"}],
+                "x": "v",
+                "y": "v2",
+                "x_label": "x",
+                "y_label": "y",
+            },
             id="find_in_nested_structure",
         ),
         pytest.param(
@@ -123,44 +129,36 @@ def test_finding_lists(dictionary, expected_result):
             {"y": {"f": ["v", "v2"]}},
             [
                 {
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                     "dvc_inferred_y_value": 1,
                     "v": 1,
                     "v2": 0.1,
                     "step": 0,
                 },
                 {
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                     "dvc_inferred_y_value": 2,
                     "v": 2,
                     "v2": 0.2,
                     "step": 1,
                 },
                 {
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                     "dvc_inferred_y_value": 0.1,
                     "v2": 0.1,
                     "v": 1,
                     "step": 0,
                 },
                 {
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                     "v": 2,
                     "v2": 0.2,
                     "dvc_inferred_y_value": 0.2,
@@ -168,6 +166,10 @@ def test_finding_lists(dictionary, expected_result):
                 },
             ],
             {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v"},
+                    {FILENAME: "f", FIELD: "v2"},
+                ],
                 "x": "step",
                 "y": "dvc_inferred_y_value",
                 "y_label": "y",
@@ -189,47 +191,43 @@ def test_finding_lists(dictionary, expected_result):
                     "z": 3,
                     "v": 1,
                     "step": 0,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                 },
                 {
                     "dvc_inferred_y_value": 2,
                     "z": 4,
                     "step": 1,
                     "v": 2,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                 },
                 {
                     "dvc_inferred_y_value": 3,
                     "v": 1,
                     "z": 3,
                     "step": 0,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "z",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "z",
                 },
                 {
                     "dvc_inferred_y_value": 4,
                     "v": 2,
                     "z": 4,
                     "step": 1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "z",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "z",
                 },
             ],
             {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v"},
+                    {FILENAME: "f", FIELD: "z"},
+                ],
                 "x": "step",
                 "y": "dvc_inferred_y_value",
                 "y_label": "y",
@@ -247,32 +245,35 @@ def test_finding_lists(dictionary, expected_result):
                 {
                     "v": 1,
                     "v2": 0.1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "v": 2,
                     "v2": 0.2,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "v": 3,
                     "v2": 0.3,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f2",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f2",
+                    FIELD: "v2",
                 },
             ],
-            {"x": "v", "y": "v2", "x_label": "v", "y_label": "v2"},
+            {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v2"},
+                    {FILENAME: "f2", FIELD: "v2"},
+                ],
+                "x": "v",
+                "y": "v2",
+                "x_label": "v",
+                "y_label": "v2",
+            },
             id="multi_file_json",
         ),
         pytest.param(
@@ -284,47 +285,43 @@ def test_finding_lists(dictionary, expected_result):
                     "v": 1,
                     "v2": 0.1,
                     "step": 0,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                 },
                 {
                     "dvc_inferred_y_value": 2,
                     "v": 2,
                     "v2": 0.2,
                     "step": 1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
                 },
                 {
                     "dvc_inferred_y_value": 0.1,
                     "v": 1,
                     "v2": 0.1,
                     "step": 0,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "dvc_inferred_y_value": 0.2,
                     "v": 2,
                     "v2": 0.2,
                     "step": 1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
             ],
             {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v"},
+                    {FILENAME: "f", FIELD: "v2"},
+                ],
                 "x": "step",
                 "y": "dvc_inferred_y_value",
                 "x_label": "step",
@@ -344,35 +341,34 @@ def test_finding_lists(dictionary, expected_result):
                     "v": 1,
                     "v2": 0.1,
                     "v3": 0.01,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "dvc_inferred_y_value": 0.01,
                     "v": 1,
                     "v2": 0.1,
                     "v3": 0.01,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v3",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v3",
                 },
                 {
                     "dvc_inferred_y_value": 0.1,
                     "v": 1,
                     "v2": 0.1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f2",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f2",
+                    FIELD: "v2",
                 },
             ],
             {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v2"},
+                    {FILENAME: "f", FIELD: "v3"},
+                    {FILENAME: "f2", FIELD: "v2"},
+                ],
                 "x": "v",
                 "y": "dvc_inferred_y_value",
                 "x_label": "v",
@@ -390,23 +386,23 @@ def test_finding_lists(dictionary, expected_result):
                 {
                     "v": 1,
                     "v2": 0.1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
                 },
                 {
                     "v": 1,
                     "v2": 0.1,
-                    VERSION_FIELD: {
-                        "revision": "r",
-                        "filename": "f2",
-                        "field": "v2",
-                    },
+                    REVISION: "r",
+                    FILENAME: "f2",
+                    FIELD: "v2",
                 },
             ],
             {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v2"},
+                    {FILENAME: "f2", FIELD: "v2"},
+                ],
                 "x": "v",
                 "y": "v2",
                 "x_label": "v",
@@ -414,6 +410,107 @@ def test_finding_lists(dictionary, expected_result):
             },
             id="multi_file_y_same_prefix",
         ),
+        pytest.param(
+            {
+                "f": {"metric": [{"x1": 1, "v": 0.1}]},
+                "f2": {"metric": [{"x2": 100, "v": 0.1}]},
+            },
+            {"y": {"f": ["v"], "f2": ["v"]}, "x": {"f": "x1", "f2": "x2"}},
+            [
+                {
+                    "x1": 1,
+                    "v": 0.1,
+                    "dvc_inferred_x_value": 1,
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
+                },
+                {
+                    "x2": 100,
+                    "v": 0.1,
+                    "dvc_inferred_x_value": 100,
+                    REVISION: "r",
+                    FILENAME: "f2",
+                    FIELD: "v",
+                },
+            ],
+            {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v"},
+                    {FILENAME: "f2", FIELD: "v"},
+                ],
+                "x": "dvc_inferred_x_value",
+                "y": "v",
+                "x_label": "x",
+                "y_label": "v",
+            },
+            id="multiple_x_fields",
+        ),
+        pytest.param(
+            {
+                "f": {
+                    "metric": [
+                        {"v": 1, "v2": 0.1, "x1": 100},
+                        {"v": 2, "v2": 0.2, "x1": 1000},
+                    ]
+                },
+                "f2": {"metric": [{"x2": -2}, {"x2": -4}]},
+            },
+            {"y": ["v", "v2"], "x": {"f": "x1", "f2": "x2"}},
+            [
+                {
+                    "dvc_inferred_x_value": 100,
+                    "dvc_inferred_y_value": 1,
+                    "v": 1,
+                    "v2": 0.1,
+                    "x1": 100,
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
+                },
+                {
+                    "dvc_inferred_x_value": 1000,
+                    "dvc_inferred_y_value": 2,
+                    "v": 2,
+                    "v2": 0.2,
+                    "x1": 1000,
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v",
+                },
+                {
+                    "dvc_inferred_x_value": -2,
+                    "dvc_inferred_y_value": 0.1,
+                    "v": 1,
+                    "v2": 0.1,
+                    "x1": 100,
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
+                },
+                {
+                    "dvc_inferred_x_value": -4,
+                    "dvc_inferred_y_value": 0.2,
+                    "v": 2,
+                    "v2": 0.2,
+                    "x1": 1000,
+                    REVISION: "r",
+                    FILENAME: "f",
+                    FIELD: "v2",
+                },
+            ],
+            {
+                "anchors_y_definitions": [
+                    {FILENAME: "f", FIELD: "v"},
+                    {FILENAME: "f", FIELD: "v2"},
+                ],
+                "x": "dvc_inferred_x_value",
+                "y": "dvc_inferred_y_value",
+                "x_label": "x",
+                "y_label": "y",
+            },
+            id="y_list_x_dict",
+        ),
     ],
 )
 def test_convert(