diff --git a/superset/examples/helpers.py b/superset/examples/helpers.py index 59e3943a3c18f..8c2ad29f49102 100644 --- a/superset/examples/helpers.py +++ b/superset/examples/helpers.py @@ -23,9 +23,8 @@ from urllib import request from superset import app, db -from superset.datasource.dao import DatasourceDAO +from superset.connectors.sqla.models import SqlaTable from superset.models.slice import Slice -from superset.utils.core import DatasourceType BASE_URL = "https://github.com/apache-superset/examples-data/blob/master/" @@ -33,7 +32,7 @@ def get_table_connector_registry() -> Any: - return DatasourceDAO.sources[DatasourceType.TABLE] + return SqlaTable def get_examples_folder() -> str: diff --git a/tests/integration_tests/dashboard_utils.py b/tests/integration_tests/dashboard_utils.py index d93e783f7a3a6..115d3269f2e50 100644 --- a/tests/integration_tests/dashboard_utils.py +++ b/tests/integration_tests/dashboard_utils.py @@ -23,7 +23,6 @@ from superset import db from superset.connectors.sqla.models import SqlaTable -from superset.datasource.dao import DatasourceDAO from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.slice import Slice @@ -36,9 +35,8 @@ def get_table( schema: Optional[str] = None, ): schema = schema or get_example_default_schema() - table_source = DatasourceDAO.sources[DatasourceType.TABLE] return ( - db.session.query(table_source) + db.session.query(SqlaTable) .filter_by(database_id=database.id, schema=schema, table_name=table_name) .one_or_none() ) @@ -55,8 +53,7 @@ def create_table_metadata( table = get_table(table_name, database, schema) if not table: - table_source = DatasourceDAO.sources[DatasourceType.TABLE] - table = table_source(schema=schema, table_name=table_name) + table = SqlaTable(schema=schema, table_name=table_name) if fetch_values_predicate: table.fetch_values_predicate = fetch_values_predicate table.database = database diff --git a/tests/integration_tests/datasource_tests.py b/tests/integration_tests/datasource_tests.py index e98faa720e140..478c9e081a4d9 100644 --- a/tests/integration_tests/datasource_tests.py +++ b/tests/integration_tests/datasource_tests.py @@ -26,7 +26,6 @@ from superset.connectors.sqla.models import SqlaTable from superset.dao.exceptions import DatasourceNotFound, DatasourceTypeNotSupportedError from superset.datasets.commands.exceptions import DatasetNotFoundError -from superset.datasource.dao import DatasourceDAO from superset.exceptions import SupersetGenericDBErrorException from superset.models.core import Database from superset.utils.core import DatasourceType, get_example_default_schema @@ -258,9 +257,10 @@ def test_external_metadata_error_return_400(self, mock_get_datasource): pytest.raises( SupersetGenericDBErrorException, - lambda: DatasourceDAO.get_datasource( - db.session, DatasourceType.TABLE, tbl.id - ).external_metadata(), + lambda: db.session.query(SqlaTable) + .filter_by(id=tbl.id) + .one_or_none() + .external_metadata(), ) resp = self.client.get(url) @@ -387,27 +387,6 @@ def my_check(datasource): app.config["DATASET_HEALTH_CHECK"] = my_check self.login(username="admin") tbl = self.get_table(name="birth_names") - datasource = DatasourceDAO.get_datasource( - db.session, DatasourceType.TABLE, tbl.id - ) + datasource = db.session.query(SqlaTable).filter_by(id=tbl.id).one_or_none() assert datasource.health_check_message == "Warning message!" app.config["DATASET_HEALTH_CHECK"] = None - - def test_get_datasource_failed(self): - pytest.raises( - DatasourceNotFound, - lambda: DatasourceDAO.get_datasource( - db.session, DatasourceType.TABLE, 9999999 - ), - ) - - self.login(username="admin") - resp = self.get_json_resp("/datasource/get/druid/500000/", raise_on_error=False) - self.assertEqual(resp.get("error"), "'druid' is not a valid DatasourceType") - - resp = self.get_json_resp( - "/datasource/get/invalid-datasource-type/500000/", raise_on_error=False - ) - self.assertEqual( - resp.get("error"), "'invalid-datasource-type' is not a valid DatasourceType" - )