address concern II

apache · Jun 15, 2022 · 64ecbfe · 64ecbfe
1 parent a7555ee
commit 64ecbfe
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 34 deletions.
diff --git a/superset/examples/helpers.py b/superset/examples/helpers.py
@@ -23,17 +23,16 @@
 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/"
 
 misc_dash_slices: Set[str] = set()  # slices assembled in a 'Misc Chart' dashboard
 
 
 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
@@ -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
@@ -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"
-        )