mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
database: drop dbt implementation (#10222)
* remove dbt implementation in the database * lazy load client * fixup * fixup * fix server side cursors * use exec_driver_sql for testing connection
This commit is contained in:
skshetry
+16
-121
@@ -1,9 +1,7 @@
|
||||
import argparse
|
||||
|
||||
from funcy import compact, merge
|
||||
|
||||
from dvc.cli import completion
|
||||
from dvc.cli.command import CmdBase
|
||||
from dvc.cli.command import CmdBase, CmdBaseNoRepo
|
||||
from dvc.cli.utils import append_doc_link
|
||||
from dvc.log import logger
|
||||
from dvc.ui import ui
|
||||
@@ -11,94 +9,35 @@ from dvc.ui import ui
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
class CmdTestDb(CmdBase):
|
||||
class CmdTestDb(CmdBaseNoRepo):
|
||||
def run(self):
|
||||
from dvc.database import get_client
|
||||
from dvc.database.dbt_utils import DBT_PROJECT_FILE, is_dbt_project
|
||||
from dvc.dependency.db import _get_dbt_config
|
||||
from dvc.config import Config
|
||||
from dvc.database import client
|
||||
from dvc.exceptions import DvcException
|
||||
|
||||
connection = self.args.conn
|
||||
|
||||
db_config = self.repo.config.get("db", {})
|
||||
config = db_config.get(connection, {})
|
||||
if connection and not config:
|
||||
db_config = Config.from_cwd().get("db", {})
|
||||
if connection not in db_config:
|
||||
raise DvcException(f"connection {connection} not found in config")
|
||||
|
||||
cli_config = compact(
|
||||
{
|
||||
"url": self.args.url,
|
||||
"username": self.args.username,
|
||||
"password": self.args.password,
|
||||
}
|
||||
)
|
||||
conn_config = merge(config, cli_config)
|
||||
|
||||
cli_dbt_config = 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:
|
||||
raise DvcException(
|
||||
f"no config set and {DBT_PROJECT_FILE} is missing; "
|
||||
"provide arguments or set a configuration"
|
||||
)
|
||||
|
||||
adapter = get_client(conn_config, project_dir=project_dir, **dbt_config)
|
||||
with adapter as db:
|
||||
config = db_config.get(connection, {})
|
||||
if self.args.url:
|
||||
config["url"] = self.args.url
|
||||
if self.args.username:
|
||||
config["username"] = self.args.username
|
||||
if self.args.password:
|
||||
config["password"] = self.args.password
|
||||
with client(config) as db:
|
||||
ui.write(f"Testing with {db}", styled=True)
|
||||
|
||||
creds = getattr(db, "creds", {})
|
||||
for k, v in creds.items():
|
||||
ui.write("\t", f"{k}:", v, styled=True)
|
||||
|
||||
if creds:
|
||||
ui.write()
|
||||
|
||||
db.test_connection()
|
||||
|
||||
ui.write("Connection successful", styled=True)
|
||||
|
||||
|
||||
class CmdImportDb(CmdBase):
|
||||
def run(self):
|
||||
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,
|
||||
@@ -116,43 +55,9 @@ def add_parser(subparsers, parent_parser):
|
||||
help=IMPORT_HELP,
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
import_parser.add_argument(
|
||||
"--url",
|
||||
help=argparse.SUPPRESS,
|
||||
# help="Location of dbt repository",
|
||||
)
|
||||
import_parser.add_argument(
|
||||
"--rev",
|
||||
nargs="?",
|
||||
help=argparse.SUPPRESS,
|
||||
# help="Git revision (e.g. SHA, branch, tag)",
|
||||
metavar="<commit>",
|
||||
)
|
||||
import_parser.add_argument(
|
||||
"--project-dir",
|
||||
nargs="?",
|
||||
help=argparse.SUPPRESS,
|
||||
# help="Subdirectory to the dbt project location",
|
||||
)
|
||||
|
||||
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=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(
|
||||
"--output-format",
|
||||
default="csv",
|
||||
@@ -177,16 +82,9 @@ def add_parser(subparsers, parent_parser):
|
||||
)
|
||||
import_parser.add_argument(
|
||||
"--conn",
|
||||
nargs="?",
|
||||
required=True,
|
||||
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)
|
||||
|
||||
@@ -197,11 +95,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")
|
||||
test_db_parser.add_argument("--dbt-conn", action="store_true", default=False)
|
||||
test_db_parser.add_argument("--conn", required=True)
|
||||
test_db_parser.add_argument("--url")
|
||||
test_db_parser.add_argument("--password")
|
||||
test_db_parser.add_argument("--username")
|
||||
test_db_parser.add_argument("--profile")
|
||||
test_db_parser.add_argument("--target")
|
||||
test_db_parser.set_defaults(func=CmdTestDb)
|
||||
|
||||
@@ -321,8 +321,6 @@ SCHEMA = {
|
||||
"feature": FeatureSchema(
|
||||
{
|
||||
Optional("machine", default=False): Bool,
|
||||
"dbt_profile": str,
|
||||
"dbt_target": str,
|
||||
},
|
||||
),
|
||||
"plots": {
|
||||
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.engine import make_url as _make_url
|
||||
from sqlalchemy.exc import NoSuchModuleError
|
||||
|
||||
from dvc import env
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.log import logger
|
||||
from dvc.types import StrOrBytesPath
|
||||
from dvc.utils import env2bool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.engine import URL, Connectable, Engine
|
||||
from sqlalchemy.sql.expression import Selectable
|
||||
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
def noop(_):
|
||||
pass
|
||||
|
||||
|
||||
def make_url(url: Union["URL", str], **kwargs: Any) -> "URL":
|
||||
return _make_url(url).set(**kwargs)
|
||||
|
||||
|
||||
def url_from_config(config: Union[str, "URL", Dict[str, str]]) -> "URL":
|
||||
if isinstance(config, dict):
|
||||
return make_url(**config)
|
||||
return make_url(config)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Serializer:
|
||||
sql: "Union[str, Selectable]"
|
||||
con: "Union[str, Connectable]"
|
||||
chunksize: int = 10_000
|
||||
|
||||
def to_csv(self, file: StrOrBytesPath, progress=noop):
|
||||
import pandas as pd
|
||||
|
||||
with open(file, mode="wb") as f:
|
||||
idfs = pd.read_sql(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: StrOrBytesPath, progress=noop): # noqa: ARG002
|
||||
import pandas as pd
|
||||
|
||||
path = os.fsdecode(file)
|
||||
df = pd.read_sql(self.sql, self.con)
|
||||
df.to_json(path, orient="records")
|
||||
|
||||
def export(self, file: StrOrBytesPath, format: str = "csv", progress=noop): # noqa: A002
|
||||
if format == "json":
|
||||
return self.to_json(file, progress=progress)
|
||||
return self.to_csv(file, progress=progress)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Client:
|
||||
engine: "Engine"
|
||||
|
||||
def test_connection(self, onerror: Optional[Callable[[], Any]] = None) -> None:
|
||||
try:
|
||||
with self.engine.connect() as conn:
|
||||
conn.exec_driver_sql("select 1")
|
||||
except Exception as exc:
|
||||
if callable(onerror):
|
||||
onerror()
|
||||
logger.exception(
|
||||
"Could not connect to the database. "
|
||||
"Check your database credentials and try again.",
|
||||
exc_info=False,
|
||||
)
|
||||
raise DvcException("The database returned the following error") from exc
|
||||
|
||||
def export(
|
||||
self,
|
||||
sql: "Union[str, Selectable]",
|
||||
file: StrOrBytesPath,
|
||||
format: str = "csv", # noqa: A002
|
||||
progress=noop,
|
||||
) -> None:
|
||||
con = self.engine.connect()
|
||||
if format == "csv":
|
||||
con = con.execution_options(stream_results=True) # use server-side cursors
|
||||
|
||||
with con:
|
||||
serializer = Serializer(sql, con)
|
||||
return serializer.export(file, format=format, progress=progress)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def handle_error(url: "URL"):
|
||||
try:
|
||||
yield
|
||||
except (ModuleNotFoundError, NoSuchModuleError) as e:
|
||||
# TODO: write installation instructions
|
||||
driver = url.drivername
|
||||
raise DvcException(f"Could not load database driver for {driver!r}") from e
|
||||
|
||||
|
||||
@contextmanager
|
||||
def client(
|
||||
url_or_config: Union[str, "URL", Dict[str, str]], **engine_kwargs: Any
|
||||
) -> Iterator[Client]:
|
||||
url = url_from_config(url_or_config)
|
||||
echo = env2bool(env.DVC_SQLALCHEMY_ECHO, False)
|
||||
engine_kwargs.setdefault("echo", echo)
|
||||
|
||||
with handle_error(url):
|
||||
engine = create_engine(url, **engine_kwargs)
|
||||
|
||||
try:
|
||||
yield Client(engine)
|
||||
finally:
|
||||
engine.dispose()
|
||||
@@ -1,22 +0,0 @@
|
||||
from typing import Any, ContextManager, Optional, Union
|
||||
|
||||
from . import dbt_query, sqla
|
||||
from .dbt_models import get_model
|
||||
from .serializer import export
|
||||
|
||||
Client = Union[sqla.SQLAlchemyClient, dbt_query.DbtClient]
|
||||
|
||||
|
||||
def get_client(
|
||||
config,
|
||||
project_dir: Optional[str] = None,
|
||||
profile: Optional[str] = None,
|
||||
target: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> "ContextManager[Client]":
|
||||
if config:
|
||||
return sqla.client(config, **kwargs)
|
||||
return dbt_query.client(project_dir=project_dir, profile=profile, target=target)
|
||||
|
||||
|
||||
__all__ = ["export", "get_client", "get_model", "Client"]
|
||||
@@ -1,40 +0,0 @@
|
||||
from typing import List, Optional
|
||||
|
||||
from dvc.log import logger
|
||||
|
||||
from .dbt_runner import dbt_show
|
||||
from .dbt_utils import check_dbt
|
||||
from .serializer import AgateSerializer
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
def ref(
|
||||
name: str,
|
||||
package: Optional[str] = None,
|
||||
version: Optional[int] = None,
|
||||
) -> str:
|
||||
parts: List[str] = []
|
||||
if package:
|
||||
parts.append(repr(package))
|
||||
|
||||
parts.append(repr(name))
|
||||
if version:
|
||||
parts.append(f"{version=}")
|
||||
|
||||
inner = ",".join(parts)
|
||||
return "{{ ref(" + inner + ") }}"
|
||||
|
||||
|
||||
@check_dbt("model")
|
||||
def get_model(
|
||||
name: str,
|
||||
package: Optional[str] = None,
|
||||
version: Optional[int] = None,
|
||||
profile: Optional[str] = None,
|
||||
target: Optional[str] = None,
|
||||
) -> AgateSerializer:
|
||||
model = ref(name, package, version=version)
|
||||
q = f"select * from {model}" # noqa: S608
|
||||
table = dbt_show(inline=q, profile=profile, target=target)
|
||||
return AgateSerializer(table)
|
||||
@@ -1,102 +0,0 @@
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
ContextManager,
|
||||
Dict,
|
||||
Iterator,
|
||||
Optional,
|
||||
Union,
|
||||
)
|
||||
|
||||
from attrs import define, field
|
||||
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.log import logger
|
||||
|
||||
from .dbt_utils import check_dbt
|
||||
from .serializer import AgateSerializer, PandasSQLSerializer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from dbt.adapters.base.impl import BaseAdapter
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
@define
|
||||
class DbtClient:
|
||||
adapter: "BaseAdapter" = field(repr=lambda o: type(o).__qualname__)
|
||||
creds: Dict[str, Any] = field(repr=False)
|
||||
|
||||
@contextmanager
|
||||
def query(self, sql: str) -> Iterator[Union[AgateSerializer, PandasSQLSerializer]]:
|
||||
from dbt.adapters.sql import SQLAdapter
|
||||
|
||||
with self.adapter.connection_named("execute"):
|
||||
if isinstance(self.adapter, SQLAdapter):
|
||||
# try to use connection to stream results
|
||||
# otherwise fallback to fetch all at once
|
||||
try:
|
||||
connection = self.adapter.connections.get_thread_connection().handle
|
||||
except Exception: # noqa: BLE001
|
||||
logger.debug("", exc_info=True)
|
||||
else:
|
||||
# pandas raises warning that it does not support
|
||||
# dbapi2-based raw connection, but it does seem to work.
|
||||
yield PandasSQLSerializer(sql, connection)
|
||||
return
|
||||
_, table = self.adapter.execute(sql, fetch=True)
|
||||
yield AgateSerializer(table)
|
||||
|
||||
def 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:
|
||||
self.adapter.debug_query()
|
||||
except Exception as exc:
|
||||
if callable(onerror):
|
||||
onerror()
|
||||
|
||||
logger.exception(
|
||||
"Could not connect to the database. "
|
||||
"Check your database credentials and try again.",
|
||||
exc_info=False,
|
||||
)
|
||||
raise DvcException("The database returned the following error") from exc
|
||||
|
||||
|
||||
@contextmanager
|
||||
@check_dbt("query")
|
||||
def client(
|
||||
project_dir: Optional[str] = None,
|
||||
profiles_dir: Optional[str] = None,
|
||||
profile: Optional[str] = None,
|
||||
target: Optional[str] = None,
|
||||
) -> Iterator["DbtClient"]:
|
||||
from dbt.adapters import factory as adapters_factory
|
||||
from dbt.adapters.base.impl import BaseAdapter
|
||||
|
||||
from .dbt_utils import get_or_build_profile, get_profiles_dir, init_dbt
|
||||
|
||||
profiles_dir = profiles_dir or get_profiles_dir(project_dir)
|
||||
flags = init_dbt(profiles_dir, os.getcwd(), target=target)
|
||||
|
||||
with flags, adapters_factory.adapter_management():
|
||||
# likely invalid connection profile or no adapter
|
||||
profile_obj = get_or_build_profile(project_dir, profile, target)
|
||||
|
||||
adapters_factory.register_adapter(profile_obj) # type: ignore[arg-type]
|
||||
adapter = adapters_factory.get_adapter(profile_obj) # type: ignore[arg-type]
|
||||
assert isinstance(adapter, BaseAdapter)
|
||||
try:
|
||||
creds = dict(profile_obj.credentials.connection_info())
|
||||
except: # noqa: E722
|
||||
creds = {}
|
||||
|
||||
yield DbtClient(adapter, creds)
|
||||
@@ -1,61 +0,0 @@
|
||||
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
|
||||
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.log import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agate import Table
|
||||
from dbt.contracts.results import RunResult
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
class DbtInternalError(DvcException):
|
||||
pass
|
||||
|
||||
|
||||
def _kw_to_cmd_args(**kwargs: Union[None, bool, str, int]) -> List[str]:
|
||||
args: List[str] = []
|
||||
for key, value in kwargs.items():
|
||||
key = key.replace("_", "-")
|
||||
if value is None:
|
||||
continue # skip creating a flag in this case
|
||||
if value is True:
|
||||
args.append(f"--{key}")
|
||||
elif value is False:
|
||||
args.append(f"--no-{key}")
|
||||
else:
|
||||
args.extend([f"--{key}", str(value)])
|
||||
return args
|
||||
|
||||
|
||||
def _dbt_invoke(*posargs: str, quiet: bool = True, **kw: Union[None, bool, str, int]):
|
||||
from dbt.cli.main import dbtRunner
|
||||
|
||||
args = _kw_to_cmd_args(quiet=quiet or None) # global options
|
||||
args.extend([*posargs, *_kw_to_cmd_args(**kw)])
|
||||
|
||||
runner = dbtRunner()
|
||||
result = runner.invoke(args)
|
||||
if result.success:
|
||||
return result.result
|
||||
raise DbtInternalError(f"failed to run dbt {posargs[0]}") from result.exception
|
||||
|
||||
|
||||
def dbt_show(
|
||||
inline: Optional[str] = None,
|
||||
limit: int = -1,
|
||||
profile: Optional[str] = None,
|
||||
target: Optional[str] = None,
|
||||
) -> "Table":
|
||||
from dbt.contracts.results import RunExecutionResult
|
||||
|
||||
result = _dbt_invoke(
|
||||
"show", inline=inline, limit=limit, profile=profile, target=target
|
||||
)
|
||||
assert isinstance(result, RunExecutionResult)
|
||||
|
||||
run_results: Sequence["RunResult"] = result.results
|
||||
run_result, *_ = run_results
|
||||
assert run_result.agate_table is not None
|
||||
return run_result.agate_table
|
||||
@@ -1,100 +0,0 @@
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from importlib.util import find_spec
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from funcy import cut_prefix
|
||||
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.log import logger
|
||||
from dvc.utils import packaging
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from dbt.config.profile import Profile
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
DBT_PROJECT_FILE = "dbt_project.yml"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def check_dbt(action: Optional[str] = None):
|
||||
if not (find_spec("dbt") and find_spec("dbt.cli")):
|
||||
action = f" {action}" if action else ""
|
||||
raise DvcException(f"Could not run{action}. dbt-core is not installed")
|
||||
yield packaging.check_required_version(pkg="dbt-core")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def init_dbt(
|
||||
profiles_dir: str,
|
||||
project_dir: str,
|
||||
target: Optional[str] = None,
|
||||
):
|
||||
from argparse import Namespace
|
||||
|
||||
from dbt import flags
|
||||
|
||||
prev = flags.get_flags()
|
||||
try:
|
||||
args = Namespace(
|
||||
use_colors=True,
|
||||
project_dir=project_dir,
|
||||
profiles_dir=profiles_dir,
|
||||
target=target,
|
||||
)
|
||||
flags.set_from_args(args, None)
|
||||
yield flags.get_flags()
|
||||
finally:
|
||||
flags.set_flags(prev)
|
||||
|
||||
|
||||
def get_profiles_dir(project_dir: Optional[str] = None) -> str:
|
||||
from dbt.cli.resolvers import default_profiles_dir
|
||||
|
||||
if profiles_dir := os.getenv("DBT_PROFILES_DIR"):
|
||||
return profiles_dir
|
||||
if project_dir and os.path.isfile(os.path.join(project_dir, "profiles.yml")):
|
||||
return project_dir
|
||||
return os.fspath(default_profiles_dir())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _handle_profile_parsing_error():
|
||||
from dbt.exceptions import DbtRuntimeError
|
||||
|
||||
try:
|
||||
yield
|
||||
except DbtRuntimeError as e:
|
||||
cause = e.__cause__ is not None and e.__cause__.__context__
|
||||
if isinstance(cause, ModuleNotFoundError) and (
|
||||
adapter := cut_prefix(cause.name, "dbt.adapters.")
|
||||
):
|
||||
# DbtRuntimeError is very noisy, so send it to debug
|
||||
logger.debug("", exc_info=True)
|
||||
raise DvcException(f"dbt-{adapter} dependency is missing") from cause
|
||||
raise DvcException("failed to read connection profiles") from e
|
||||
|
||||
|
||||
def is_dbt_project(project_dir: str):
|
||||
return os.path.isfile(os.path.join(project_dir, DBT_PROJECT_FILE))
|
||||
|
||||
|
||||
@_handle_profile_parsing_error()
|
||||
def get_or_build_profile(
|
||||
project_dir: Optional[str], profile: Optional[str], target: Optional[str]
|
||||
) -> "Profile":
|
||||
from dbt.config.profile import Profile
|
||||
from dbt.config.renderer import ProfileRenderer
|
||||
from dbt.config.runtime import load_profile
|
||||
|
||||
if project_dir and is_dbt_project(project_dir):
|
||||
return load_profile(
|
||||
project_dir, {}, profile_name_override=profile, target_override=target
|
||||
)
|
||||
|
||||
if not profile:
|
||||
raise DvcException("No profile specified to query from.")
|
||||
|
||||
renderer = ProfileRenderer({})
|
||||
return Profile.render(renderer, profile, target_override=target)
|
||||
@@ -1,54 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import sqlalchemy as sa
|
||||
from agate import Table
|
||||
|
||||
|
||||
def noop(_):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class PandasSQLSerializer:
|
||||
sql: "Union[str, sa.sql.Selectable]"
|
||||
con: "sa.engine.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)
|
||||
@@ -1,85 +0,0 @@
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, ContextManager, Dict, Iterator, Optional, Union
|
||||
|
||||
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
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
def make_url(url: Union[URL, str], **kwargs: Any) -> URL:
|
||||
return _make_url(url).set(**kwargs)
|
||||
|
||||
|
||||
def url_from_config(config: Union[str, URL, Dict[str, str]]) -> URL:
|
||||
if isinstance(config, (str, URL)):
|
||||
return make_url(config)
|
||||
return make_url(**config)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SQLAlchemyClient:
|
||||
engine: Engine
|
||||
|
||||
@contextmanager
|
||||
def query(self, sql: Union[str, Selectable]) -> Iterator[PandasSQLSerializer]:
|
||||
with self.engine.connect().execution_options(stream_results=True) as 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:
|
||||
with self.engine.connect() as conn:
|
||||
conn.execute(text("select 1"))
|
||||
except Exception as exc:
|
||||
if callable(onerror):
|
||||
onerror()
|
||||
logger.exception(
|
||||
"Could not connect to the database. "
|
||||
"Check your database credentials and try again.",
|
||||
exc_info=False,
|
||||
)
|
||||
raise DvcException("The database returned the following error") from exc
|
||||
|
||||
|
||||
@contextmanager
|
||||
def handle_error(url: URL):
|
||||
try:
|
||||
yield
|
||||
except (ModuleNotFoundError, NoSuchModuleError) as e:
|
||||
# TODO: write installation instructions
|
||||
raise DvcException(
|
||||
f"Could not load database driver for {url.drivername!r}"
|
||||
) from e
|
||||
|
||||
|
||||
@contextmanager
|
||||
def client(
|
||||
url_or_config: Union[str, URL, Dict[str, str]], **engine_kwargs: Any
|
||||
) -> Iterator[SQLAlchemyClient]:
|
||||
url = url_from_config(url_or_config)
|
||||
with handle_error(url):
|
||||
echo = env2bool(env.DVC_SQLALCHEMY_ECHO, False)
|
||||
engine_kwargs.setdefault("echo", echo)
|
||||
engine = create_engine(url, **engine_kwargs)
|
||||
try:
|
||||
yield SQLAlchemyClient(engine)
|
||||
finally:
|
||||
engine.dispose()
|
||||
@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Mapping, Set
|
||||
from dvc.output import ARTIFACT_SCHEMA, DIR_FILES_SCHEMA, Output
|
||||
|
||||
from .base import Dependency
|
||||
from .db import DB_SCHEMA, PARAM_DB, DbDependency, DbtDependency
|
||||
from .db import DbDependency
|
||||
from .param import ParamsDependency
|
||||
from .repo import RepoDependency
|
||||
|
||||
@@ -15,7 +15,7 @@ from .repo import RepoDependency
|
||||
SCHEMA: Mapping[str, Any] = {
|
||||
**ARTIFACT_SCHEMA,
|
||||
**RepoDependency.REPO_SCHEMA,
|
||||
**DB_SCHEMA,
|
||||
**DbDependency.DB_SCHEMA,
|
||||
Output.PARAM_FILES: [DIR_FILES_SCHEMA],
|
||||
Output.PARAM_FS_CONFIG: dict,
|
||||
}
|
||||
@@ -23,16 +23,13 @@ SCHEMA: Mapping[str, Any] = {
|
||||
|
||||
def _get(stage, p, info, **kwargs):
|
||||
d = info or {}
|
||||
db = d.get(PARAM_DB, {})
|
||||
params = d.pop(ParamsDependency.PARAM_PARAMS, None)
|
||||
repo = d.pop(RepoDependency.PARAM_REPO, None)
|
||||
|
||||
if params:
|
||||
return ParamsDependency(stage, p, params)
|
||||
if DbtDependency.PARAM_MODEL in db:
|
||||
return DbtDependency(repo, stage, info)
|
||||
if db:
|
||||
return DbDependency(stage, info)
|
||||
if DbDependency.PARAM_DB in d:
|
||||
return DbDependency(stage, d)
|
||||
|
||||
assert p
|
||||
if repo:
|
||||
|
||||
+35
-250
@@ -1,58 +1,20 @@
|
||||
import os
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Union
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional
|
||||
|
||||
from funcy import compact, log_durations
|
||||
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.log import logger
|
||||
from dvc.scm import SCM
|
||||
|
||||
from .base import Dependency
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from rich.status import Status
|
||||
|
||||
from dvc.output import Output
|
||||
from dvc.stage import Stage
|
||||
|
||||
logger = logger.getChild(__name__)
|
||||
|
||||
|
||||
PARAM_DB = "db"
|
||||
PARAM_PROFILE = "profile"
|
||||
PARAM_FILE_FORMAT = "file_format"
|
||||
|
||||
|
||||
def _get_dbt_config(config: Dict) -> Dict:
|
||||
conf = config.get("feature", {})
|
||||
pref = "dbt_"
|
||||
return {k.lstrip(pref): v for k, v in conf.items() if k.startswith(pref)}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def log_status(
|
||||
msg, status: Optional["Status"] = None, log=logger.debug
|
||||
) -> Iterator["Status"]:
|
||||
from funcy import log_durations
|
||||
|
||||
from dvc.ui import ui
|
||||
|
||||
with log_durations(log, msg), status or ui.status(msg) as st:
|
||||
st.update(msg)
|
||||
yield st
|
||||
|
||||
|
||||
@contextmanager
|
||||
def chdir(path):
|
||||
wdir = os.getcwd()
|
||||
os.chdir(path)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
os.chdir(wdir)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def download_progress(to: "Output") -> Iterator[Callable[[int], Any]]:
|
||||
from dvc.ui import ui
|
||||
@@ -80,14 +42,27 @@ class AbstractDependency(Dependency):
|
||||
|
||||
class DbDependency(AbstractDependency):
|
||||
PARAM_CONNECTION = "connection"
|
||||
PARAM_DBT_CONNECTION = "dbt_connection"
|
||||
PARAM_DB = "db"
|
||||
PARAM_QUERY = "query"
|
||||
PARAM_TABLE = "table"
|
||||
QUERY_SCHEMA = {PARAM_QUERY: str, PARAM_CONNECTION: str, PARAM_TABLE: str}
|
||||
PARAM_FILE_FORMAT = "file_format"
|
||||
DB_SCHEMA = {
|
||||
PARAM_DB: {
|
||||
PARAM_QUERY: str,
|
||||
PARAM_CONNECTION: str,
|
||||
PARAM_FILE_FORMAT: str,
|
||||
PARAM_TABLE: str,
|
||||
}
|
||||
}
|
||||
|
||||
def __init__(self, stage: "Stage", info, *args, **kwargs):
|
||||
super().__init__(stage, info, *args, **kwargs)
|
||||
self.target = None
|
||||
self.db_info: Dict[str, str] = self.info.get(self.PARAM_DB, {})
|
||||
self.connection = self.db_info.get(self.PARAM_CONNECTION)
|
||||
|
||||
@property
|
||||
def sql(self) -> Optional[str]:
|
||||
return self.db_info.get(self.PARAM_QUERY) or self.db_info.get(self.PARAM_TABLE)
|
||||
|
||||
def __repr__(self):
|
||||
return "{}:{}".format(
|
||||
@@ -97,10 +72,7 @@ class DbDependency(AbstractDependency):
|
||||
def __str__(self):
|
||||
from dvc.utils.humanize import truncate_text
|
||||
|
||||
db_info = self.info.get(PARAM_DB, {})
|
||||
query = db_info.get(self.PARAM_QUERY)
|
||||
table = db_info.get(self.PARAM_TABLE)
|
||||
return truncate_text(table or query or "[no table or query]", 50)
|
||||
return truncate_text(self.sql or "", 50)
|
||||
|
||||
def workspace_status(self):
|
||||
return False # no workspace to check
|
||||
@@ -112,8 +84,8 @@ class DbDependency(AbstractDependency):
|
||||
"""nothing to save."""
|
||||
|
||||
def dumpd(self, **kwargs):
|
||||
db_info = compact(self.info.get(PARAM_DB, {}))
|
||||
return {PARAM_DB: db_info} if db_info else {}
|
||||
db_info = compact(self.db_info)
|
||||
return {self.PARAM_DB: db_info} if db_info else {}
|
||||
|
||||
def update(self, rev=None):
|
||||
"""nothing to update."""
|
||||
@@ -125,209 +97,22 @@ class DbDependency(AbstractDependency):
|
||||
file_format: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
from dvc.database import export, get_client
|
||||
from dvc.database import client
|
||||
from dvc.ui import ui
|
||||
|
||||
db_info = self.info.get(PARAM_DB, {})
|
||||
query = db_info.get(self.PARAM_QUERY)
|
||||
table = db_info.get(self.PARAM_TABLE)
|
||||
if not (query or table):
|
||||
sql = self.sql
|
||||
if not sql:
|
||||
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")
|
||||
db_config = self.repo.config.get(self.PARAM_DB, {})
|
||||
config = db_config.get(self.connection)
|
||||
if not config:
|
||||
raise DvcException(f"connection {self.connection} not found in config")
|
||||
|
||||
connection = db_info.get(self.PARAM_CONNECTION)
|
||||
|
||||
db_config = self.repo.config.get("db", {})
|
||||
config = db_config.get(connection)
|
||||
if connection and not config:
|
||||
raise DvcException(f"connection {connection} not found in config")
|
||||
|
||||
project_dir = self.repo.root_dir
|
||||
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:
|
||||
logger.debug("using adapter: %s", db)
|
||||
with log_status("Testing connection") as status:
|
||||
file_format = file_format or self.db_info.get(self.PARAM_FILE_FORMAT, "csv")
|
||||
with client(config) as db:
|
||||
msg = "Testing connection"
|
||||
with log_durations(logger.debug, msg), ui.status(msg) as status:
|
||||
db.test_connection(onerror=status.stop)
|
||||
|
||||
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)
|
||||
status.stop()
|
||||
with download_progress(to) as prog:
|
||||
return export(serializer, file, format=file_format, progress=prog)
|
||||
|
||||
|
||||
class DbtDependency(AbstractDependency):
|
||||
PARAM_MODEL = "model"
|
||||
PARAM_VERSION = "version"
|
||||
PARAM_PROJECT_DIR = "project_dir"
|
||||
DBT_SCHEMA = {
|
||||
PARAM_MODEL: str,
|
||||
PARAM_VERSION: str,
|
||||
PARAM_PROJECT_DIR: str,
|
||||
}
|
||||
|
||||
def __init__(self, def_repo: Dict[str, Any], stage: "Stage", info, *args, **kwargs):
|
||||
self.def_repo = def_repo or {}
|
||||
self.target = None
|
||||
super().__init__(stage, info, *args, **kwargs)
|
||||
|
||||
def __repr__(self):
|
||||
return "{}:{}".format(
|
||||
self.__class__.__name__,
|
||||
"".join(f"{k}=={v}" for k, v in {**self.def_repo, **self.info}.items()),
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
from .repo import RepoDependency
|
||||
|
||||
repo = self.def_repo.get(RepoDependency.PARAM_URL)
|
||||
rev = self.def_repo.get(RepoDependency.PARAM_REV)
|
||||
|
||||
db_info = self.info.get(PARAM_DB, {})
|
||||
db = db_info.get(self.PARAM_MODEL, "")
|
||||
project_dir = db_info.get(self.PARAM_PROJECT_DIR, "")
|
||||
repo_info = ""
|
||||
if repo:
|
||||
repo_info += repo
|
||||
if rev:
|
||||
repo_info += f"@{rev}"
|
||||
if project_dir:
|
||||
repo_info += f":/{project_dir}"
|
||||
return db + (f"({repo_info})" if repo_info else "")
|
||||
|
||||
@property
|
||||
def locked_rev(self):
|
||||
from .repo import RepoDependency
|
||||
|
||||
return self.def_repo.get(RepoDependency.PARAM_REV_LOCK)
|
||||
|
||||
@property
|
||||
def rev(self):
|
||||
from .repo import RepoDependency
|
||||
|
||||
return self.def_repo.get(RepoDependency.PARAM_REV)
|
||||
|
||||
def workspace_status(self):
|
||||
if not self.def_repo:
|
||||
return
|
||||
|
||||
current = self._get_clone(self.locked_rev or self.rev).get_rev()
|
||||
updated = self._get_clone(self.rev).get_rev()
|
||||
if current != updated:
|
||||
return {str(self): "update available"}
|
||||
return {}
|
||||
|
||||
def status(self):
|
||||
return self.workspace_status()
|
||||
|
||||
def save(self):
|
||||
from .repo import RepoDependency
|
||||
|
||||
if not self.def_repo:
|
||||
return
|
||||
|
||||
rev = self._get_clone(self.locked_rev or self.rev).get_rev()
|
||||
if self.def_repo.get(RepoDependency.PARAM_REV_LOCK) is None:
|
||||
self.def_repo[RepoDependency.PARAM_REV_LOCK] = rev
|
||||
|
||||
def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]:
|
||||
from .repo import RepoDependency
|
||||
|
||||
def_repo = {}
|
||||
if self.def_repo:
|
||||
def_repo = RepoDependency._dump_def_repo(self.def_repo)
|
||||
|
||||
db_info = compact(self.info.get(PARAM_DB, {}))
|
||||
return compact({RepoDependency.PARAM_REPO: def_repo, PARAM_DB: db_info})
|
||||
|
||||
def _get_clone(self, rev):
|
||||
from dvc.repo.open_repo import _cached_clone
|
||||
|
||||
from .repo import RepoDependency
|
||||
|
||||
url = self.def_repo.get(RepoDependency.PARAM_URL)
|
||||
repo_root = self.repo.root_dir if self.repo else os.getcwd()
|
||||
return SCM(_cached_clone(url, rev) if url else repo_root)
|
||||
|
||||
def update(self, rev=None):
|
||||
from .repo import RepoDependency
|
||||
|
||||
if not self.def_repo:
|
||||
return
|
||||
|
||||
if rev:
|
||||
self.def_repo[RepoDependency.PARAM_REV] = rev
|
||||
else:
|
||||
rev = self.rev
|
||||
self.def_repo[RepoDependency.PARAM_REV_LOCK] = self._get_clone(rev).get_rev()
|
||||
|
||||
def download(
|
||||
self,
|
||||
to: "Output",
|
||||
jobs: Optional[int] = None, # noqa: ARG002
|
||||
file_format: Optional[str] = None,
|
||||
) -> None:
|
||||
from .repo import RepoDependency
|
||||
|
||||
project_dir = self.info.get(PARAM_DB, {}).get(self.PARAM_PROJECT_DIR, "")
|
||||
if self.def_repo:
|
||||
ctx = repo = self._get_clone(self.locked_rev or self.rev)
|
||||
self.def_repo[RepoDependency.PARAM_REV_LOCK] = repo.get_rev()
|
||||
root = wdir = repo.root_dir
|
||||
else:
|
||||
ctx = nullcontext()
|
||||
root = self.repo.root_dir
|
||||
wdir = self.stage.wdir
|
||||
|
||||
project_path = os.path.join(wdir, project_dir) if project_dir else root
|
||||
with ctx, chdir(project_path):
|
||||
self._download_db(to, file_format=file_format)
|
||||
|
||||
def _download_db(
|
||||
self,
|
||||
to: "Output",
|
||||
version: Optional[int] = None,
|
||||
file_format: Optional[str] = None,
|
||||
) -> None:
|
||||
from dvc.database import export, get_model
|
||||
|
||||
db_info = self.info.get(PARAM_DB, {})
|
||||
model = db_info.get(self.PARAM_MODEL)
|
||||
version = version or db_info.get(self.PARAM_VERSION)
|
||||
file_format = file_format or db_info.get(PARAM_FILE_FORMAT) or "csv"
|
||||
assert file_format
|
||||
if not model:
|
||||
raise DvcException("Cannot download, no model specified")
|
||||
|
||||
dbt_config = _get_dbt_config(self.repo.config)
|
||||
profile = db_info.get(PARAM_PROFILE) or dbt_config.get(PARAM_PROFILE)
|
||||
target = self.target or db_info.get("target") or dbt_config.get("target")
|
||||
|
||||
with log_status("Downloading model"):
|
||||
serializer = get_model(
|
||||
model, version=version, profile=profile, target=target
|
||||
)
|
||||
# NOTE: we keep everything in memory, and then export it out later.
|
||||
with download_progress(to) as progress:
|
||||
export(serializer, to.fs_path, format=file_format, progress=progress)
|
||||
|
||||
|
||||
DB_SCHEMA = {
|
||||
PARAM_DB: {
|
||||
PARAM_PROFILE: str,
|
||||
PARAM_FILE_FORMAT: str,
|
||||
**DbDependency.QUERY_SCHEMA,
|
||||
**DbtDependency.DBT_SCHEMA,
|
||||
},
|
||||
}
|
||||
with download_progress(to) as progress:
|
||||
db.export(sql, to.fs_path, format=file_format, progress=progress)
|
||||
|
||||
+17
-30
@@ -1,4 +1,6 @@
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
from typing import TYPE_CHECKING, Dict, Optional
|
||||
|
||||
from funcy import compact
|
||||
|
||||
from dvc.exceptions import OutputDuplicationError
|
||||
from dvc.repo.scm_context import scm_context
|
||||
@@ -13,18 +15,11 @@ from . import locked
|
||||
|
||||
@locked
|
||||
@scm_context
|
||||
def imp_db( # noqa: PLR0913
|
||||
def imp_db(
|
||||
self: "Repo",
|
||||
url: Optional[str] = None,
|
||||
rev: Optional[str] = None,
|
||||
project_dir: Optional[str] = None,
|
||||
sql: Optional[str] = None,
|
||||
table: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
version: Optional[int] = None,
|
||||
frozen: bool = True,
|
||||
profile: Optional[str] = None,
|
||||
target: Optional[str] = None,
|
||||
output_format: str = "csv",
|
||||
out: Optional[str] = None,
|
||||
force: bool = False,
|
||||
@@ -35,41 +30,33 @@ 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}
|
||||
if rev:
|
||||
erepo["rev"] = rev
|
||||
|
||||
assert sql or table
|
||||
assert output_format in ("csv", "json")
|
||||
|
||||
db: Dict[str, Any] = {"file_format": output_format}
|
||||
if profile:
|
||||
db["profile"] = profile
|
||||
db: Dict[str, str] = compact(
|
||||
{
|
||||
"connection": connection,
|
||||
"file_format": output_format,
|
||||
"query": sql,
|
||||
"table": table,
|
||||
}
|
||||
)
|
||||
|
||||
if model:
|
||||
db.update({"model": model, "version": version, "project_dir": project_dir})
|
||||
else:
|
||||
db.update({"connection": connection, "query": sql, "table": table})
|
||||
|
||||
file_name = model or table or "results"
|
||||
file_name = table or "results"
|
||||
out = out or f"{file_name}.{output_format}"
|
||||
out = resolve_output(url or ".", out, force=force)
|
||||
out = resolve_output(".", out, force=force)
|
||||
|
||||
path, wdir, out = resolve_paths(self, out, always_local=True)
|
||||
stage = self.stage.create(
|
||||
single_stage=True,
|
||||
validate=False,
|
||||
fname=path,
|
||||
deps=[None],
|
||||
wdir=wdir,
|
||||
deps=[url],
|
||||
outs=[out],
|
||||
erepo=erepo,
|
||||
fs_config=None,
|
||||
db=db,
|
||||
)
|
||||
|
||||
stage.deps[0].target = target
|
||||
try:
|
||||
self.check_graph(stages={stage})
|
||||
except OutputDuplicationError as exc:
|
||||
|
||||
@@ -286,17 +286,17 @@ class Stage(params.StageParams):
|
||||
if not self.is_import:
|
||||
return False
|
||||
|
||||
from dvc.dependency import DbDependency, DbtDependency
|
||||
from dvc.dependency import DbDependency
|
||||
|
||||
return isinstance(self.deps[0], (DbDependency, DbtDependency))
|
||||
return isinstance(self.deps[0], DbDependency)
|
||||
|
||||
@property
|
||||
def is_versioned_import(self) -> bool:
|
||||
from dvc.dependency import DbDependency, DbtDependency
|
||||
from dvc.dependency import DbDependency
|
||||
|
||||
return (
|
||||
self.is_import
|
||||
and not isinstance(self.deps[0], (DbDependency, DbtDependency))
|
||||
and not isinstance(self.deps[0], DbDependency)
|
||||
and self.deps[0].fs.version_aware
|
||||
)
|
||||
|
||||
@@ -462,7 +462,7 @@ class Stage(params.StageParams):
|
||||
if not (self.is_repo_import or self.is_import):
|
||||
raise StageUpdateError(self.relpath)
|
||||
|
||||
# always force update DbDep/DbtDep since we don't know if it's changed
|
||||
# always force update DbDep since we don't know if it's changed
|
||||
force = self.is_db_import
|
||||
update_import(
|
||||
self,
|
||||
|
||||
+1
-2
@@ -112,9 +112,8 @@ tests = [
|
||||
"pytest-xdist>=3.2",
|
||||
'pywin32>=225; sys_platform == "win32"', # optional test dependency
|
||||
'tzdata; sys_platform == "win32"', # for testing with celery
|
||||
"dbt-sqlite", # optional dependency for testing `import-db`
|
||||
"sqlalchemy>=1,<3", # optional dependency for `import-db`
|
||||
"pandas>=1",
|
||||
"pandas>=1", # optional dependency for `import-db`
|
||||
]
|
||||
webdav = ["dvc-webdav>=3,<4"]
|
||||
webhdfs = ["dvc-webhdfs>=3.1,<4"]
|
||||
|
||||
+83
-132
@@ -1,165 +1,116 @@
|
||||
"""import_db tests without involving any databases."""
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from agate import Table
|
||||
|
||||
from dvc.database.serializer import AgateSerializer
|
||||
|
||||
data1 = Table([(i, i) for i in range(1, 10)], ["id", "value"])
|
||||
data2 = Table([(i, i) for i in range(10, 20)], ["id", "value"])
|
||||
|
||||
serializers = AgateSerializer(data1), AgateSerializer(data2)
|
||||
from funcy import compact
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(mocker):
|
||||
m = mocker.patch("dvc.database.get_client")
|
||||
client = mocker.MagicMock()
|
||||
client.query.return_value.__enter__.side_effect = serializers
|
||||
client.table.return_value.__enter__.side_effect = serializers
|
||||
m.return_value.__enter__.return_value = client
|
||||
return m
|
||||
def db_path(tmp_dir):
|
||||
return tmp_dir / "main.db"
|
||||
|
||||
|
||||
def test_sql_with_dbt(client, tmp_dir, dvc):
|
||||
stage = dvc.imp_db(
|
||||
sql="select * from model", profile="profile", output_format="json"
|
||||
)
|
||||
@pytest.fixture
|
||||
def seed_db(db_path):
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("CREATE TABLE model (id INTEGER PRIMARY KEY, value INTEGER)")
|
||||
|
||||
db = {"file_format": "json", "profile": "profile", "query": "select * from model"}
|
||||
def inner(values):
|
||||
conn.executemany("INSERT INTO model(value) VALUES(?)", [(i,) for i in values])
|
||||
conn.commit()
|
||||
|
||||
assert (tmp_dir / stage.relpath).parse() == {
|
||||
"md5": "dae475048115315c8c26a10538e4b2ec",
|
||||
"frozen": True,
|
||||
"deps": [{"db": db}],
|
||||
"outs": [
|
||||
{
|
||||
"hash": "md5",
|
||||
"md5": "56226d443067b195c90b427db557e1f2",
|
||||
"path": "results.json",
|
||||
"size": 243,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
dvc.update(stage.addressing)
|
||||
assert (tmp_dir / stage.relpath).parse() == {
|
||||
"md5": "79b42f9ddb2a3ed999eb494ccab470c6",
|
||||
"frozen": True,
|
||||
"deps": [{"db": db}],
|
||||
"outs": [
|
||||
{
|
||||
"hash": "md5",
|
||||
"md5": "08d806ee7e1f309c4561750805f02276",
|
||||
"path": "results.json",
|
||||
"size": 290,
|
||||
}
|
||||
],
|
||||
}
|
||||
with closing(conn):
|
||||
yield inner
|
||||
|
||||
|
||||
def test_sql_with_conn_string(client, tmp_dir, dvc):
|
||||
@pytest.fixture
|
||||
def db_connection(dvc, db_path):
|
||||
with dvc.config.edit(level="local") as conf:
|
||||
conf["db"] = {"conn": {"url": "conn"}}
|
||||
conf["db"] = {"conn": {"url": f"sqlite:///{db_path.fs_path}"}}
|
||||
return "conn"
|
||||
|
||||
stage = dvc.imp_db(
|
||||
sql="select * from model", connection="conn", output_format="json"
|
||||
|
||||
def load_data(file, output_format):
|
||||
if output_format == "json":
|
||||
return pd.read_json(file, orient="records")
|
||||
return pd.read_csv(file)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("output_format", ("csv", "json"))
|
||||
@pytest.mark.parametrize(
|
||||
"args,file_name",
|
||||
[
|
||||
({"sql": "select * from model"}, "results"),
|
||||
({"table": "model"}, "model"),
|
||||
],
|
||||
)
|
||||
def test(
|
||||
M,
|
||||
tmp_dir,
|
||||
scm,
|
||||
dvc,
|
||||
db_connection,
|
||||
seed_db,
|
||||
output_format,
|
||||
args,
|
||||
file_name,
|
||||
):
|
||||
seed_db(values=range(5))
|
||||
if output_format == "json":
|
||||
file_size = 96, 192
|
||||
md5 = "6039fe7565d212b339aaa446ca234e5d", "e1b8adf4d9eb9ab2b64d3ab0bb5f65ac"
|
||||
elif os.name == "nt":
|
||||
file_size = 35, 61
|
||||
md5 = "14c34db5ddd184345c06f74718539f04", "3bb836e6d43c9afa43a9d73b36bbbab4"
|
||||
else:
|
||||
file_size = 29, 50
|
||||
md5 = "6f7fc0d701d1ac13eec83d79fffaf427", "c04f712f8167496a2fb43f289f2b7e28"
|
||||
|
||||
db = compact(
|
||||
{
|
||||
"file_format": output_format,
|
||||
"connection": db_connection,
|
||||
"table": args.get("table"),
|
||||
"query": args.get("sql"),
|
||||
}
|
||||
)
|
||||
db = {"file_format": "json", "connection": "conn", "query": "select * from model"}
|
||||
stage = dvc.imp_db(**args, connection=db_connection, output_format=output_format)
|
||||
|
||||
output_file = f"{file_name}.{output_format}"
|
||||
df = load_data(output_file, output_format)
|
||||
assert df.values.tolist() == [[i + 1, i] for i in range(5)]
|
||||
assert (tmp_dir / stage.relpath).parse() == {
|
||||
"md5": "f0dcca8e8411907adb25aa547d7c432c",
|
||||
"md5": M.instance_of(str),
|
||||
"frozen": True,
|
||||
"deps": [{"db": db}],
|
||||
"outs": [
|
||||
{
|
||||
"md5": md5[0],
|
||||
"size": file_size[0],
|
||||
"hash": "md5",
|
||||
"md5": "56226d443067b195c90b427db557e1f2",
|
||||
"path": "results.json",
|
||||
"size": 243,
|
||||
"path": output_file,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
seed_db(values=range(5, 10))
|
||||
|
||||
dvc.update(stage.addressing)
|
||||
|
||||
df = load_data(output_file, output_format)
|
||||
assert df.values.tolist() == [[i + 1, i] for i in range(10)]
|
||||
assert (tmp_dir / stage.relpath).parse() == {
|
||||
"md5": "8c0441b206d8ee7c7b9e3d54871c3cbb",
|
||||
"md5": M.instance_of(str),
|
||||
"frozen": True,
|
||||
"deps": [{"db": db}],
|
||||
"outs": [
|
||||
{
|
||||
"md5": md5[1],
|
||||
"size": file_size[1],
|
||||
"hash": "md5",
|
||||
"md5": "08d806ee7e1f309c4561750805f02276",
|
||||
"path": "results.json",
|
||||
"size": 290,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def test_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")
|
||||
|
||||
assert (tmp_dir / stage.relpath).parse() == {
|
||||
"md5": "a8157fb5b28457e1910506e9d4d4fca3",
|
||||
"frozen": True,
|
||||
"deps": [{"db": {"file_format": "json", "model": "model"}}],
|
||||
"outs": [
|
||||
{
|
||||
"hash": "md5",
|
||||
"md5": "56226d443067b195c90b427db557e1f2",
|
||||
"path": "model.json",
|
||||
"size": 243,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
dvc.update(stage.addressing)
|
||||
assert (tmp_dir / stage.relpath).parse() == {
|
||||
"md5": "2f84287f109fc2b8755a9e64ad910b5d",
|
||||
"frozen": True,
|
||||
"deps": [{"db": {"file_format": "json", "model": "model"}}],
|
||||
"outs": [
|
||||
{
|
||||
"hash": "md5",
|
||||
"md5": "08d806ee7e1f309c4561750805f02276",
|
||||
"path": "model.json",
|
||||
"size": 290,
|
||||
"path": output_file,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
@@ -1,147 +0,0 @@
|
||||
"""E2E tests involving dvc/dbt and a ~~real~~ sqlite db.
|
||||
|
||||
Any other tests should likely be mocked tests in `test_import_db` rather than here.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
import pytest
|
||||
from agate import Table
|
||||
|
||||
from dvc.types import StrPath
|
||||
|
||||
pytest.importorskip("dbt.adapters.sqlite", reason="dbt-sqlite not installed")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def disable_tracking(monkeypatch):
|
||||
monkeypatch.setenv("DO_NOT_TRACK", "1")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_sqlite_adapter(mocker):
|
||||
# https://github.com/codeforkjeff/dbt-sqlite/issues/47
|
||||
mocker.patch(
|
||||
"dbt.adapters.sqlite.connections.SQLiteCredentials.unique_field", "sqlite_host"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db_path(tmp_dir):
|
||||
d = tmp_dir / "db"
|
||||
d.mkdir()
|
||||
return d / "main.db"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def seed_db(db_path):
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("CREATE TABLE model (id INTEGER PRIMARY KEY, value INTEGER)")
|
||||
|
||||
def inner(values):
|
||||
conn.executemany("INSERT INTO model(value) VALUES(?)", [(i,) for i in values])
|
||||
conn.commit()
|
||||
|
||||
with closing(conn):
|
||||
yield inner
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db_config(dvc, db_path):
|
||||
# only needed for `sql_conn` tests
|
||||
with dvc.config.edit(level="local") as conf:
|
||||
conf["db"] = {"conn": {"url": f"sqlite:///{db_path.fs_path}"}}
|
||||
return "conn"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dbt_profile(tmp_dir, scm, db_path):
|
||||
(tmp_dir / "profiles.yml").dump(
|
||||
{
|
||||
"sqlite_profile": {
|
||||
"target": "dev",
|
||||
"outputs": {
|
||||
"dev": {
|
||||
"type": "sqlite",
|
||||
"threads": 1,
|
||||
"database": "database",
|
||||
"schema": "main",
|
||||
"schemas_and_paths": {"main": os.fspath(db_path)},
|
||||
"schema_directory": os.fspath(db_path.parent),
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
scm.add_commit(["profiles.yml"], message="add profiles.yml")
|
||||
return "sqlite_profile"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dbt_project(tmp_dir, scm, dbt_profile):
|
||||
(tmp_dir / "dbt_project.yml").dump({"name": "project", "profile": dbt_profile})
|
||||
scm.add_commit(["dbt_project.yml"], message="add dbt_project")
|
||||
return tmp_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dbt_model(dbt_project):
|
||||
dbt_project.scm_gen(
|
||||
{"models": {"model.sql": "select * from models"}}, commit="add models"
|
||||
)
|
||||
return "model"
|
||||
|
||||
|
||||
@pytest.fixture(params=("sql", "model", "external_model", "sql_conn", "table"))
|
||||
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}
|
||||
if request.param == "table":
|
||||
conn = request.getfixturevalue("db_config")
|
||||
return {"table": "model", "connection": conn}
|
||||
|
||||
dbt_project = request.getfixturevalue("dbt_project")
|
||||
return {
|
||||
"model": request.getfixturevalue("dbt_model"),
|
||||
"url": None if request.param == "model" else os.fspath(dbt_project),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_name(import_db_parameters):
|
||||
if model := import_db_parameters.get("model"):
|
||||
return model
|
||||
if table := import_db_parameters.get("table"):
|
||||
return table
|
||||
return "results"
|
||||
|
||||
|
||||
def load_table(file: StrPath, typ: str) -> "Table":
|
||||
return getattr(Table, f"from_{typ}")(file)
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::ResourceWarning") # dbt leaks fileobj from logger
|
||||
@pytest.mark.parametrize("output_format", ("csv", "json"))
|
||||
def test_e2e(
|
||||
tmp_dir, scm, dvc, seed_db, import_db_parameters, file_name, output_format
|
||||
):
|
||||
seed_db(values=range(5))
|
||||
|
||||
stage = dvc.imp_db(**import_db_parameters, output_format=output_format)
|
||||
|
||||
output_file = tmp_dir / f"{file_name}.{output_format}"
|
||||
output = load_table(output_file, output_format)
|
||||
assert output.rows == [(i + 1, i) for i in range(5)]
|
||||
|
||||
seed_db(values=range(5, 10))
|
||||
|
||||
dvc.update(stage.addressing)
|
||||
|
||||
output = load_table(output_file, output_format)
|
||||
assert output.rows == [(i + 1, i) for i in range(10)]
|
||||
Reference in New Issue
Block a user