fix: improve get_db_engine_spec_for_backend (#21171)

* fix: improve get_db_engine_spec_for_backend

* Fix tests

* Fix docs

* fix lint

* fix fallback

* Fix engine validation

* Fix test
This commit is contained in:
Beto Dealmeida
2022-08-29 13:42:42 -05:00
committed by GitHub
parent 710a8ce5c0
commit 8772e2cdb3
13 changed files with 309 additions and 130 deletions

View File

@@ -59,7 +59,7 @@ def test_get_metrics(mocker: MockFixture) -> None:
},
]
database.get_db_engine_spec_for_backend = mocker.MagicMock( # type: ignore
database.get_db_engine_spec = mocker.MagicMock( # type: ignore
return_value=CustomSqliteEngineSpec
)
assert database.get_metrics("table") == [
@@ -70,3 +70,78 @@ def test_get_metrics(mocker: MockFixture) -> None:
"verbose_name": "COUNT(DISTINCT user_id)",
},
]
def test_get_db_engine_spec(mocker: MockFixture) -> None:
"""
Tests for ``get_db_engine_spec``.
"""
from superset.db_engine_specs import BaseEngineSpec
from superset.models.core import Database
# pylint: disable=abstract-method
class PostgresDBEngineSpec(BaseEngineSpec):
"""
A DB engine spec with drivers and a default driver.
"""
engine = "postgresql"
engine_aliases = {"postgres"}
drivers = {
"psycopg2": "The default Postgres driver",
"asyncpg": "An async Postgres driver",
}
default_driver = "psycopg2"
# pylint: disable=abstract-method
class OldDBEngineSpec(BaseEngineSpec):
"""
And old DB engine spec without drivers nor a default driver.
"""
engine = "mysql"
load_engine_specs = mocker.patch("superset.db_engine_specs.load_engine_specs")
load_engine_specs.return_value = [
PostgresDBEngineSpec,
OldDBEngineSpec,
]
assert (
Database(database_name="db", sqlalchemy_uri="postgresql://").db_engine_spec
== PostgresDBEngineSpec
)
assert (
Database(
database_name="db", sqlalchemy_uri="postgresql+psycopg2://"
).db_engine_spec
== PostgresDBEngineSpec
)
assert (
Database(
database_name="db", sqlalchemy_uri="postgresql+asyncpg://"
).db_engine_spec
== PostgresDBEngineSpec
)
assert (
Database(
database_name="db", sqlalchemy_uri="postgresql+fancynewdriver://"
).db_engine_spec
== PostgresDBEngineSpec
)
assert (
Database(database_name="db", sqlalchemy_uri="mysql://").db_engine_spec
== OldDBEngineSpec
)
assert (
Database(
database_name="db", sqlalchemy_uri="mysql+mysqlconnector://"
).db_engine_spec
== OldDBEngineSpec
)
assert (
Database(
database_name="db", sqlalchemy_uri="mysql+fancynewdriver://"
).db_engine_spec
== OldDBEngineSpec
)