feat(sip-68): Add DatasourceDAO class to manage querying different datasources easier

superset/dao/datasource/dao.py

@@ -0,0 +1,151 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import List, Optional, Set, Union
+
+from flask_babel import _
+from sqlalchemy import or_
+from sqlalchemy.orm import Session, subqueryload
+from sqlalchemy.orm.exc import NoResultFound
+
+from superset.connectors.sqla.models import SqlaTable, Table
+from superset.dao.base import BaseDAO
+from superset.datasets.commands.exceptions import DatasetNotFoundError
+from superset.datasets.models import Dataset
+from superset.models.core import Database
+from superset.models.sql_lab import Query, SavedQuery
+
+Datasource = Union[Dataset, SqlaTable, Table, Query, SavedQuery]
+
+
+class DatasourceDAO(BaseDAO):
+
+    sources = {
+        # using table -> SqlaTable for backward compatibility at the moment
+        "table": SqlaTable,
+        "query": Query,
+        "saved_query": SavedQuery,
+        "sl_dataset": Dataset,
+        "sl_table": Table,
+    }
+
+    @classmethod
+    def get_datasource(
+        cls, datasource_type: str, datasource_id: int, session: Session
+    ) -> Datasource:
+        if datasource_type not in cls.sources:
+            raise DatasetNotFoundError()
+
+        datasource = (
+            session.query(cls.sources[datasource_type])
+            .filter_by(id=datasource_id)
+            .one_or_none()
+        )
+
+        if not datasource:
+            raise DatasetNotFoundError()
+
+        return datasource
+
+    def get_all_datasources(self, session: Session) -> List[Datasource]:
+        datasources: List["Datasource"] = []
+        for source_class in DatasourceDAO.sources.values():
+            qry = session.query(source_class)
+            qry = source_class.default_query(qry)
+            datasources.extend(qry.all())
+        return datasources
+
+    def get_datasource_by_id(self, session: Session, datasource_id: int) -> Datasource:
+        """
+        Find a datasource instance based on the unique id.
+        :param session: Session to use
+        :param datasource_id: unique id of datasource
+        :return: Datasource corresponding to the id
+        :raises NoResultFound: if no datasource is found corresponding to the id
+        """
+        for datasource_class in DatasourceDAO.sources.values():
+            try:
+                return (
+                    session.query(datasource_class)
+                    .filter(datasource_class.id == datasource_id)
+                    .one()
+                )
+            except NoResultFound:
+                # proceed to next datasource type
+                pass
+        raise NoResultFound(_("Datasource id not found: %(id)s", id=datasource_id))
+
+    def get_datasource_by_name(  # pylint: disable=too-many-arguments
+        self,
+        session: Session,
+        datasource_type: str,
+        datasource_name: str,
+        schema: str,
+        database_name: str,
+    ) -> Optional[Datasource]:
+        datasource_class = DatasourceDAO.sources[datasource_type]
+        return datasource_class.get_datasource_by_name(
+            session, datasource_name, schema, database_name
+        )
+
+    def query_datasources_by_permissions(  # pylint: disable=invalid-name
+        self,
+        session: Session,
+        database: Database,
+        permissions: Set[str],
+        schema_perms: Set[str],
+    ) -> List[Datasource]:
+        # TODO(bogdan): add unit test
+        datasource_class = DatasourceDAO.sources[database.type]
+        return (
+            session.query(datasource_class)
+            .filter_by(database_id=database.id)
+            .filter(
+                or_(
+                    datasource_class.perm.in_(permissions),
+                    datasource_class.schema_perm.in_(schema_perms),
+                )
+            )
+            .all()
+        )
+
+    def get_eager_datasource(
+        self, session: Session, datasource_type: str, datasource_id: int
+    ) -> Datasource:
+        """Returns datasource with columns and metrics."""
+        datasource_class = DatasourceDAO.sources[datasource_type]
+        return (
+            session.query(datasource_class)
+            .options(
+                subqueryload(datasource_class.columns),
+                subqueryload(datasource_class.metrics),
+            )
+            .filter_by(id=datasource_id)
+            .one()
+        )
+
+    def query_datasources_by_name(
+        self,
+        session: Session,
+        database: Database,
+        datasource_name: str,
+        schema: Optional[str] = None,
+    ) -> List[Datasource]:
+        datasource_class = DatasourceDAO.sources[database.type]
+        return datasource_class.query_datasources_by_name(
+            session, database, datasource_name, schema=schema
+        )

tests/unit_tests/dao/datasource_test.py

@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+from sqlalchemy.orm.session import Session
+
+
+def create_test_data(session: Session) -> None:
+    from superset.columns.models import Column
+    from superset.connectors.sqla.models import SqlaTable, TableColumn
+    from superset.models.core import Database
+    from superset.models.sql_lab import Query, SavedQuery
+
+    engine = session.get_bind()
+    SqlaTable.metadata.create_all(engine)  # pylint: disable=no-member
+
+    db = Database(database_name="my_database", sqlalchemy_uri="sqlite://")
+
+    columns = [
+        TableColumn(column_name="a", type="INTEGER"),
+    ]
+
+    sqla_table = SqlaTable(
+        table_name="my_sqla_table",
+        columns=columns,
+        metrics=[],
+        database=db,
+    )
+
+    query_obj = Query(
+        client_id="foo",
+        database=db,
+        tab_name="test_tab",
+        sql_editor_id="test_editor_id",
+        sql="select * from bar",
+        select_sql="select * from bar",
+        executed_sql="select * from bar",
+        limit=100,
+        select_as_cta=False,
+        rows=100,
+        error_message="none",
+        results_key="abc",
+    )
+
+    saved_query = SavedQuery(database=db, sql="select * from foo")
+
+    session.add(saved_query)
+    session.add(query_obj)
+    session.add(db)
+    session.add(sqla_table)
+    session.flush()
+
+
+def test_get_datasource_sqlatable(app_context: None, session: Session) -> None:
+    from superset.connectors.sqla.models import SqlaTable
+    from superset.dao.datasource.dao import DatasourceDAO
+
+    create_test_data(session)
+
+    result = DatasourceDAO.get_datasource(
+        datasource_type="table", datasource_id=1, session=session
+    )
+
+    assert 1 == result.id
+    assert "my_sqla_table" == result.table_name
+    assert isinstance(result, SqlaTable)
+
+
+def test_get_datasource_query(app_context: None, session: Session) -> None:
+    from superset.dao.datasource.dao import DatasourceDAO
+    from superset.models.sql_lab import Query
+
+    create_test_data(session)
+
+    result = DatasourceDAO.get_datasource(
+        datasource_type="query", datasource_id=1, session=session
+    )
+
+    assert result.id == 1
+    assert isinstance(result, Query)
+
+
+def test_get_datasource_saved_query(app_context: None, session: Session) -> None:
+    from superset.dao.datasource.dao import DatasourceDAO
+    from superset.models.sql_lab import SavedQuery
+
+    create_test_data(session)
+
+    result = DatasourceDAO.get_datasource(
+        datasource_type="saved_query", datasource_id=1, session=session
+    )
+
+    assert result.id == 1
+    assert isinstance(result, SavedQuery)
+
+
+def test_get_datasource_sl_table(app_context: None, session: Session) -> None:
+    pass
+
+
+def test_get_datasource_sl_dataset(app_context: None, session: Session) -> None:
+    pass