mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-06 23:09:55 +01:00
skshetry
8768b50aed
* experimental: support connection strings for databases using connectorx This adds an experimental support for connector strings using connectorx. User has to provide config for the connection strings with: Example: ```console dvc config --local db_connection.pg postgres://user:pass@localhost:5432/database ``` And then, they can import using: ```console dvc import-db --sql 'query' --conn pg ``` Check supported databases and how to configure in https://sfu-db.github.io/connector-x/databases.html. This `db_connection` config is subject to change. connectorx still has the same problem with dbt as before: no support for streaming, but it'll be a bit more simpler to configure with the connection strings. * add pandas to deps * sqlalchemy * dbt streaming sql * fix test * remove connectorx * remove check in set * use mode='wb' for serialization
129 lines
3.7 KiB
Python
129 lines
3.7 KiB
Python
"""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)
|
|
|
|
|
|
@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"
|
|
)
|
|
|
|
db = {"file_format": "json", "profile": "profile", "query": "select * from model"}
|
|
|
|
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,
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
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.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,
|
|
}
|
|
],
|
|
}
|