Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-10 07:18:27 +01:00
Code Issues Packages Projects Releases Wiki Activity

import-db: support importing a table (#10219)

Browse Source 
* import-db: support importing a table

* unhide the import-db command
This commit is contained in:
skshetry
2024-01-05 09:26:07 +05:45
committed by GitHub
parent f1de712710
commit 0100c14d6a
11 changed files with 178 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/commands/imp_db.py
+65 -25
View File
@@ -18,7 +18,7 @@ class CmdTestDb(CmdBase):
from dvc.dependency.db import _get_dbt_config
from dvc.exceptions import DvcException
connection = self.args.connection
connection = self.args.conn
db_config = self.repo.config.get("db", {})
config = db_config.get(connection, {})
@@ -34,11 +34,18 @@ class CmdTestDb(CmdBase):
)
conn_config = merge(config, cli_config)
cli_dbt_config = {"profile": self.args.profile, "target": self.args.target}
cli_dbt_config = compact(
{"profile": self.args.profile, "target": self.args.target}
)
dbt_config = merge(_get_dbt_config(self.repo.config), cli_dbt_config)
project_dir = self.repo.root_dir
if not (conn_config or dbt_config):
if not self.args.dbt_conn:
raise DvcException(
"no config set; provide arguments or set a configuration"
)
if is_dbt_project(project_dir):
ui.write("Using", DBT_PROJECT_FILE, "for testing", styled=True)
else:
@@ -65,61 +72,87 @@ class CmdTestDb(CmdBase):
class CmdImportDb(CmdBase):
def run(self):
if not (self.args.sql or self.args.model):
raise argparse.ArgumentTypeError("Either of --sql or --model is required.")
from dvc.exceptions import InvalidArgumentError
if self.args.table or self.args.sql:
arg = "--table" if self.args.table else "--sql"
options = {
"url": self.args.url,
"rev": self.args.rev,
"project_dir": self.args.project_dir,
}
opt = next((o for o, v in options.items() if v), None)
if opt:
raise InvalidArgumentError(f"argument {opt}: not allowed with {arg}")
if not self.args.conn and not self.args.dbt_conn:
raise InvalidArgumentError(f"{arg} requires --conn")
if self.args.model and self.args.conn:
raise InvalidArgumentError("argument --model: not allowed with --conn")
self.repo.imp_db(
url=self.args.url,
rev=self.args.rev,
project_dir=self.args.project_dir,
sql=self.args.sql,
table=self.args.table,
model=self.args.model,
profile=self.args.profile,
target=self.args.target,
output_format=self.args.output_format,
out=self.args.out,
force=self.args.force,
connection=self.args.connection,
connection=self.args.conn,
)
return 0
def add_parser(subparsers, parent_parser):
IMPORT_HELP = (
"Download file or directory tracked by DVC or by Git "
"into the workspace, and track it."
)
IMPORT_HELP = "Snapshot a table or a SQL query result to a CSV/JSON format"
import_parser = subparsers.add_parser(
"import-db",
parents=[parent_parser],
description=append_doc_link(IMPORT_HELP, "import"),
add_help=False,
description=append_doc_link(IMPORT_HELP, "import-db"),
help=IMPORT_HELP,
formatter_class=argparse.RawTextHelpFormatter,
)
import_parser.add_argument(
"--url", help="Location of DVC or Git repository to download from"
"--url",
help=argparse.SUPPRESS,
# help="Location of dbt repository",
)
import_parser.add_argument(
"--rev",
nargs="?",
help="Git revision (e.g. SHA, branch, tag)",
help=argparse.SUPPRESS,
# help="Git revision (e.g. SHA, branch, tag)",
metavar="<commit>",
)
import_parser.add_argument(
"--project-dir", nargs="?", help="Subdirectory to the dbt project location"
"--project-dir",
nargs="?",
help=argparse.SUPPRESS,
# help="Subdirectory to the dbt project location",
)
group = import_parser.add_mutually_exclusive_group()
group.add_argument(
"--sql",
help="SQL query",
)
group = import_parser.add_mutually_exclusive_group(required=True)
group.add_argument("--sql", help="SQL query to snapshot")
group.add_argument("--table", help="Table to snapshot")
group.add_argument(
"--model",
help="Model name to download",
help=argparse.SUPPRESS,
# help="Model name to download",
)
import_parser.add_argument(
"--profile",
help=argparse.SUPPRESS,
# help="Profile to use",
)
import_parser.add_argument(
"--target",
help=argparse.SUPPRESS,
# help="Target to use",
)
import_parser.add_argument("--profile", help="Profile to use")
import_parser.add_argument("--target", help="Target to use")
import_parser.add_argument(
"--output-format",
default="csv",
@@ -144,10 +177,16 @@ def add_parser(subparsers, parent_parser):
)
import_parser.add_argument(
"--conn",
dest="connection",
nargs="?",
help="Database connection to use, needs to be set in config",
)
import_parser.add_argument(
"--dbt-conn",
action="store_true",
default=False,
help=argparse.SUPPRESS,
# help="Use dbt connection",
)
import_parser.set_defaults(func=CmdImportDb)
@@ -158,7 +197,8 @@ def add_parser(subparsers, parent_parser):
description=append_doc_link(TEST_DB_HELP, "test-db"),
add_help=False,
)
test_db_parser.add_argument("--conn", dest="connection")
test_db_parser.add_argument("--conn")
test_db_parser.add_argument("--dbt-conn", action="store_true", default=False)
test_db_parser.add_argument("--url")
test_db_parser.add_argument("--password")
test_db_parser.add_argument("--username")
dvc/database/dbt_query.py
+17 -3
View File
@@ -1,6 +1,15 @@
import os
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
from typing import (
TYPE_CHECKING,
Any,
Callable,
ContextManager,
Dict,
Iterator,
Optional,
Union,
)
from attrs import define, field
@@ -41,6 +50,11 @@ class DbtClient:
_, table = self.adapter.execute(sql, fetch=True)
yield AgateSerializer(table)
def table(
self, name: str, limit: Optional[int] = None
) -> ContextManager[Union[AgateSerializer, PandasSQLSerializer]]:
raise NotImplementedError
def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None:
with self.adapter.connection_named("debug"):
try:
@@ -50,8 +64,8 @@ class DbtClient:
onerror()
logger.exception(
"dvc was unable to connect to the specified database. "
"Please check your database credentials and try again.",
"Could not connect to the database. "
"Check your database credentials and try again.",
exc_info=False,
)
raise DvcException("The database returned the following error") from exc
dvc/database/serializer.py
+2 -2
View File
@@ -12,8 +12,8 @@ def noop(_):
@dataclass
class PandasSQLSerializer:
sql: "Union[sa.TextClause, str]"
con: "sa.Connection"
sql: "Union[str, sa.sql.Selectable]"
con: "sa.engine.Connection"
chunksize: int = 10_000
def to_csv(self, file: str, progress=noop) -> None:
dvc/database/sqla.py
+16 -4
View File
@@ -1,14 +1,18 @@
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterator, Optional, Union
from typing import Any, Callable, ContextManager, Dict, Iterator, Optional, Union
from sqlalchemy import create_engine, text
from sqlalchemy import create_engine, select, text
from sqlalchemy.engine import URL, Engine
from sqlalchemy.engine import make_url as _make_url
from sqlalchemy.exc import NoSuchModuleError
from sqlalchemy.sql import table
from sqlalchemy.sql.expression import Selectable
from dvc import env
from dvc.exceptions import DvcException
from dvc.log import logger
from dvc.utils import env2bool
from .serializer import PandasSQLSerializer
@@ -30,9 +34,15 @@ class SQLAlchemyClient:
engine: Engine
@contextmanager
def query(self, sql: str) -> Iterator[PandasSQLSerializer]:
def query(self, sql: Union[str, Selectable]) -> Iterator[PandasSQLSerializer]:
with self.engine.connect().execution_options(stream_results=True) as conn:
yield PandasSQLSerializer(text(sql), conn)
yield PandasSQLSerializer(sql, conn)
def table(
self, name: str, limit: Optional[int] = None
) -> ContextManager[PandasSQLSerializer]:
stmt = select("*").select_from(table(name)).limit(limit)
return self.query(stmt)
def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None:
try:
@@ -66,6 +76,8 @@ def client(
) -> Iterator[SQLAlchemyClient]:
url = url_from_config(url_or_config)
with handle_error(url):
echo = env2bool(env.DVC_SQLALCHEMY_ECHO, False)
engine_kwargs.setdefault("echo", echo)
engine = create_engine(url, **engine_kwargs)
try:
yield SQLAlchemyClient(engine)
dvc/dependency/__init__.py
+3 -3
View File
@@ -29,10 +29,10 @@ def _get(stage, p, info, **kwargs):
if params:
return ParamsDependency(stage, p, params)
if DbDependency.PARAM_QUERY in db:
return DbDependency(stage, info)
if db:
if DbtDependency.PARAM_MODEL in db:
return DbtDependency(repo, stage, info)
if db:
return DbDependency(stage, info)
assert p
if repo:
dvc/dependency/db.py
+23 -14
View File
@@ -79,9 +79,11 @@ class AbstractDependency(Dependency):
class DbDependency(AbstractDependency):
PARAM_QUERY = "query"
PARAM_CONNECTION = "connection"
QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str}
PARAM_DBT_CONNECTION = "dbt_connection"
PARAM_QUERY = "query"
PARAM_TABLE = "table"
QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str, PARAM_TABLE: str}
def __init__(self, stage: "Stage", info, *args, **kwargs):
super().__init__(stage, info, *args, **kwargs)
@@ -96,8 +98,9 @@ class DbDependency(AbstractDependency):
from dvc.utils.humanize import truncate_text
db_info = self.info.get(PARAM_DB, {})
query = db_info.get(self.PARAM_QUERY, "[query]")
return truncate_text(query, 50)
query = db_info.get(self.PARAM_QUERY)
table = db_info.get(self.PARAM_TABLE)
return truncate_text(table or query or "[no table or query]", 50)
def workspace_status(self):
return False # no workspace to check
@@ -126,20 +129,29 @@ class DbDependency(AbstractDependency):
db_info = self.info.get(PARAM_DB, {})
query = db_info.get(self.PARAM_QUERY)
if not query:
raise DvcException("Cannot download: no query specified")
table = db_info.get(self.PARAM_TABLE)
if not (query or table):
raise DvcException("Cannot download: no query or table specified")
if query and table:
raise DvcException(
"Cannot download: query and table are mutually exclusive"
)
dbt_config = _get_dbt_config(self.repo.config)
profile = db_info.get(PARAM_PROFILE) or dbt_config.get(PARAM_PROFILE)
target = self.target or dbt_config.get("target")
connection = db_info.get(self.PARAM_CONNECTION)
db_config = self.repo.config.get("db", {})
config = db_config.get(connection)
if connection and not config:
raise DvcException(f"connection {connection} not found in config")
project_dir = self.repo.root_dir
file = to.fs_path
file_format = file_format or db_info.get(PARAM_FILE_FORMAT) or "csv"
assert file_format
with get_client(
config, project_dir=project_dir, profile=profile, target=target
) as db:
@@ -147,15 +159,12 @@ class DbDependency(AbstractDependency):
with log_status("Testing connection") as status:
db.test_connection(onerror=status.stop)
file_format = file_format or db_info.get(PARAM_FILE_FORMAT, "csv")
assert file_format
with log_status("Executing query") as status, db.query(query) as serializer:
status.stop()
query = db.query(query) if query else db.table(table)
with log_status("Executing query") as status, query as serializer:
logger.debug("using serializer: %s", serializer)
with download_progress(to) as progress:
return export(
serializer, to.fs_path, format=file_format, progress=progress
)
status.stop()
with download_progress(to) as prog:
return export(serializer, file, format=file_format, progress=prog)
class DbtDependency(AbstractDependency):
dvc/env.py
+1
View File
@@ -16,5 +16,6 @@ DVC_STUDIO_OFFLINE = "DVC_STUDIO_OFFLINE"
DVC_STUDIO_REPO_URL = "DVC_STUDIO_REPO_URL"
DVC_STUDIO_TOKEN = "DVC_STUDIO_TOKEN" # noqa: S105
DVC_STUDIO_URL = "DVC_STUDIO_URL"
DVC_SQLALCHEMY_ECHO = "DVC_SQLALCHEMY_ECHO"
DVC_SYSTEM_CONFIG_DIR = "DVC_SYSTEM_CONFIG_DIR"
DVC_UPDATER_ENDPOINT = "DVC_UPDATER_ENDPOINT"
dvc/repo/imp_db.py
+5 -3
View File
@@ -19,6 +19,7 @@ def imp_db( # noqa: PLR0913
rev: Optional[str] = None,
project_dir: Optional[str] = None,
sql: Optional[str] = None,
table: Optional[str] = None,
model: Optional[str] = None,
version: Optional[int] = None,
frozen: bool = True,
@@ -34,6 +35,7 @@ def imp_db( # noqa: PLR0913
"The functionality may change or break without notice, "
"which could lead to unexpected behavior."
)
assert model or sql or table
erepo = None
if model and url:
erepo = {"url": url}
@@ -47,12 +49,12 @@ def imp_db( # noqa: PLR0913
db["profile"] = profile
if model:
out = out or f"{model}.{output_format}"
db.update({"model": model, "version": version, "project_dir": project_dir})
else:
out = out or f"results.{output_format}"
db.update({"query": sql, "connection": connection})
db.update({"connection": connection, "query": sql, "table": table})
file_name = model or table or "results"
out = out or f"{file_name}.{output_format}"
out = resolve_output(url or ".", out, force=force)
path, wdir, out = resolve_paths(self, out, always_local=True)
stage = self.stage.create(
tests/func/test_import_db.py
+39 -2
View File
@@ -15,11 +15,12 @@ def client(mocker):
m = mocker.patch("dvc.database.get_client")
client = mocker.MagicMock()
client.query.return_value.__enter__.side_effect = serializers
client.table.return_value.__enter__.side_effect = serializers
m.return_value.__enter__.return_value = client
return m
def test_sql(client, tmp_dir, dvc):
def test_sql_with_dbt(client, tmp_dir, dvc):
stage = dvc.imp_db(
sql="select * from model", profile="profile", output_format="json"
)
@@ -56,7 +57,7 @@ def test_sql(client, tmp_dir, dvc):
}
def test_sql_conn_string(client, tmp_dir, dvc):
def test_sql_with_conn_string(client, tmp_dir, dvc):
with dvc.config.edit(level="local") as conf:
conf["db"] = {"conn": {"url": "conn"}}
@@ -94,6 +95,42 @@ def test_sql_conn_string(client, tmp_dir, dvc):
}
def test_table_with_conn_string(client, tmp_dir, dvc):
with dvc.config.edit(level="local") as conf:
conf["db"] = {"conn": {"url": "conn"}}
stage = dvc.imp_db(table="model", connection="conn", output_format="json")
db = {"file_format": "json", "connection": "conn", "table": "model"}
assert (tmp_dir / stage.relpath).parse() == {
"md5": "39dfcd61935e039a7bc3f158d7fcf62c",
"frozen": True,
"deps": [{"db": db}],
"outs": [
{
"hash": "md5",
"md5": "56226d443067b195c90b427db557e1f2",
"path": "model.json",
"size": 243,
}
],
}
dvc.update(stage.addressing)
assert (tmp_dir / stage.relpath).parse() == {
"md5": "20d77a700710bad7d094c460f5939b7c",
"frozen": True,
"deps": [{"db": db}],
"outs": [
{
"hash": "md5",
"md5": "08d806ee7e1f309c4561750805f02276",
"path": "model.json",
"size": 290,
}
],
}
def test_model(mocker, tmp_dir, dvc):
mocker.patch("dvc.database.get_model", side_effect=serializers)
stage = dvc.imp_db(model="model", output_format="json")
tests/func/test_import_db_e2e.py
+6 -1
View File
@@ -94,7 +94,7 @@ def dbt_model(dbt_project):
return "model"
@pytest.fixture(params=("sql", "model", "external_model", "sql_conn"))
@pytest.fixture(params=("sql", "model", "external_model", "sql_conn", "table"))
def import_db_parameters(request: pytest.FixtureRequest):
if request.param == "sql":
profile = request.getfixturevalue("dbt_profile")
@@ -102,6 +102,9 @@ def import_db_parameters(request: pytest.FixtureRequest):
if request.param == "sql_conn":
conn = request.getfixturevalue("db_config")
return {"sql": "select * from model", "connection": conn}
if request.param == "table":
conn = request.getfixturevalue("db_config")
return {"table": "model", "connection": conn}
dbt_project = request.getfixturevalue("dbt_project")
return {
@@ -114,6 +117,8 @@ def import_db_parameters(request: pytest.FixtureRequest):
def file_name(import_db_parameters):
if model := import_db_parameters.get("model"):
return model
if table := import_db_parameters.get("table"):
return table
return "results"
tests/unit/command/test_help.py
+1 -1
View File
@@ -28,7 +28,7 @@ def command_tuples():
# the no. of commands will usually go up,
# but if we ever remove commands and drop below, adjust the magic number accordingly
assert len(commands) >= 96
assert len(commands) >= 112
return sorted(commands)