feat(sqllab): Add /sqllab endpoint to the v1 api (#24983)

superset/security/manager.py

@@ -236,6 +236,7 @@ class SupersetSecurityManager(  # pylint: disable=too-many-public-methods
         ("can_execute_sql_query", "SQLLab"),
         ("can_estimate_query_cost", "SQL Lab"),
         ("can_export_csv", "SQLLab"),
+        ("can_read", "SQLLab"),
         ("can_sqllab_history", "Superset"),
         ("can_sqllab", "Superset"),
         ("can_test_conn", "Superset"),  # Deprecated permission remove on 3.0.0

superset/sqllab/api.py

@@ -20,7 +20,7 @@ from urllib import parse
 
 import simplejson as json
 from flask import request, Response
-from flask_appbuilder.api import expose, protect, rison
+from flask_appbuilder.api import expose, protect, rison, safe
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from marshmallow import ValidationError
 
@@ -47,6 +47,7 @@ from superset.sqllab.schemas import (
     ExecutePayloadSchema,
     QueryExecutionResponseSchema,
     sql_lab_get_results_schema,
+    SQLLabBootstrapSchema,
 )
 from superset.sqllab.sql_json_executer import (
     ASynchronousSqlJsonExecutor,
@@ -54,6 +55,7 @@ from superset.sqllab.sql_json_executer import (
     SynchronousSqlJsonExecutor,
 )
 from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext
+from superset.sqllab.utils import bootstrap_sqllab_data
 from superset.sqllab.validators import CanAccessQueryValidatorImpl
 from superset.superset_typing import FlaskResponse
 from superset.utils import core as utils
@@ -83,8 +85,55 @@ class SqlLabRestApi(BaseSupersetApi):
         EstimateQueryCostSchema,
         ExecutePayloadSchema,
         QueryExecutionResponseSchema,
+        SQLLabBootstrapSchema,
     )
 
+    @expose("/", methods=("GET",))
+    @protect()
+    @safe
+    @statsd_metrics
+    @event_logger.log_this_with_context(
+        action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.get",
+        log_to_statsd=False,
+    )
+    def get(self) -> Response:
+        """Get the bootstrap data for SqlLab
+        ---
+        get:
+          summary: Get the bootstrap data for SqlLab page
+          description: >-
+            Assembles SQLLab bootstrap data (active_tab, databases, queries,
+            tab_state_ids) in a single endpoint. The data can be assembled
+            from the current user's id.
+          responses:
+            200:
+              description: Returns the initial bootstrap data for SqlLab
+              content:
+                application/json:
+                  schema:
+                    $ref: '#/components/schemas/SQLLabBootstrapSchema'
+            400:
+              $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
+            403:
+              $ref: '#/components/responses/403'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        user_id = utils.get_user_id()
+        # TODO: Replace with a command class once fully migrated to SPA
+        result = bootstrap_sqllab_data(user_id)
+
+        return json_success(
+            json.dumps(
+                {"result": result},
+                default=utils.json_iso_dttm_ser,
+                ignore_nan=True,
+            ),
+            200,
+        )
+
     @expose("/estimate/", methods=("POST",))
     @protect()
     @statsd_metrics

superset/sqllab/schemas.py

@@ -16,6 +16,8 @@
 # under the License.
 from marshmallow import fields, Schema
 
+from superset.databases.schemas import ImportV1DatabaseSchema
+
 sql_lab_get_results_schema = {
     "type": "object",
     "properties": {
@@ -95,3 +97,50 @@ class QueryExecutionResponseSchema(Schema):
     expanded_columns = fields.List(fields.Dict())
     query = fields.Nested(QueryResultSchema)
     query_id = fields.Integer()
+
+
+class TableSchema(Schema):
+    database_id = fields.Integer()
+    description = fields.String()
+    expanded = fields.Boolean()
+    id = fields.Integer()
+    schema = fields.String()
+    tab_state_id = fields.Integer()
+    table = fields.String()
+
+
+class TabStateSchema(Schema):
+    active = fields.Boolean()
+    autorun = fields.Boolean()
+    database_id = fields.Integer()
+    extra_json = fields.Dict()
+    hide_left_bar = fields.Boolean()
+    id = fields.String()
+    label = fields.String()
+    latest_query = fields.Nested(QueryResultSchema)
+    query_limit = fields.Integer()
+    saved_query = fields.Dict(
+        allow_none=True,
+        metadata={"id": "SavedQuery id"},
+    )
+    schema = fields.String()
+    sql = fields.String()
+    table_schemas = fields.List(fields.Nested(TableSchema))
+    user_id = fields.Integer()
+
+
+class SQLLabBootstrapSchema(Schema):
+    active_tab = fields.Nested(TabStateSchema)
+    databases = fields.Dict(
+        keys=fields.String(
+            metadata={"description": "Database id"},
+        ),
+        values=fields.Nested(ImportV1DatabaseSchema),
+    )
+    queries = fields.Dict(
+        keys=fields.String(
+            metadata={"description": "Query id"},
+        ),
+        values=fields.Nested(QueryResultSchema),
+    )
+    tab_state_ids = fields.List(fields.String())

tests/integration_tests/core_tests.py

@@ -1090,54 +1090,6 @@ class TestCore(SupersetTestCase, InsertChartMixin):
             data = self.get_resp(url)
             self.assertTrue(html_string in data)
 
-    @mock.patch.dict(
-        "superset.extensions.feature_flag_manager._feature_flags",
-        {"SQLLAB_BACKEND_PERSISTENCE": True},
-        clear=True,
-    )
-    def test_sqllab_backend_persistence_payload(self):
-        username = "admin"
-        self.login(username)
-        user_id = security_manager.find_user(username).id
-
-        # create a tab
-        data = {
-            "queryEditor": json.dumps(
-                {
-                    "title": "Untitled Query 1",
-                    "dbId": 1,
-                    "schema": None,
-                    "autorun": False,
-                    "sql": "SELECT ...",
-                    "queryLimit": 1000,
-                }
-            )
-        }
-        resp = self.get_json_resp("/tabstateview/", data=data)
-        tab_state_id = resp["id"]
-
-        # run a query in the created tab
-        self.run_sql(
-            "SELECT name FROM birth_names",
-            "client_id_1",
-            username=username,
-            raise_on_error=True,
-            sql_editor_id=str(tab_state_id),
-        )
-        # run an orphan query (no tab)
-        self.run_sql(
-            "SELECT name FROM birth_names",
-            "client_id_2",
-            username=username,
-            raise_on_error=True,
-        )
-
-        # we should have only 1 query returned, since the second one is not
-        # associated with any tabs
-        # TODO: replaces this spec by api/v1/sqllab spec later
-        payload = bootstrap_sqllab_data(user_id)
-        self.assertEqual(len(payload["queries"]), 1)
-
     @mock.patch.dict(
         "superset.extensions.feature_flag_manager._feature_flags",
         {"SQLLAB_BACKEND_PERSISTENCE": True},

tests/integration_tests/sql_lab/api_tests.py

@@ -28,6 +28,7 @@ import prison
 from sqlalchemy.sql import func
 from unittest import mock
 
+from flask_appbuilder.security.sqla.models import Role
 from tests.integration_tests.test_app import app
 from superset import db, sql_lab
 from superset.common.db_query_status import QueryStatus
@@ -37,11 +38,95 @@ from superset.utils import core as utils
 from superset.models.sql_lab import Query
 
 from tests.integration_tests.base_tests import SupersetTestCase
+from tests.integration_tests.fixtures.users import create_gamma_sqllab_no_data
 
 QUERIES_FIXTURE_COUNT = 10
 
 
 class TestSqlLabApi(SupersetTestCase):
+    @pytest.mark.usefixtures("create_gamma_sqllab_no_data")
+    @mock.patch.dict(
+        "superset.extensions.feature_flag_manager._feature_flags",
+        {"SQLLAB_BACKEND_PERSISTENCE": False},
+        clear=True,
+    )
+    def test_get_from_empty_bootsrap_data(self):
+        self.login(username="gamma_sqllab_no_data")
+        resp = self.client.get("/api/v1/sqllab/")
+        assert resp.status_code == 200
+        data = json.loads(resp.data.decode("utf-8"))
+        result = data.get("result")
+        assert result["active_tab"] == None
+        assert result["queries"] == {}
+        assert result["tab_state_ids"] == []
+        self.assertEqual(len(result["databases"]), 0)
+
+    @mock.patch.dict(
+        "superset.extensions.feature_flag_manager._feature_flags",
+        {"SQLLAB_BACKEND_PERSISTENCE": True},
+        clear=True,
+    )
+    def test_get_from_bootstrap_data_with_queries(self):
+        username = "admin"
+        self.login(username)
+
+        # create a tab
+        data = {
+            "queryEditor": json.dumps(
+                {
+                    "title": "Untitled Query 1",
+                    "dbId": 1,
+                    "schema": None,
+                    "autorun": False,
+                    "sql": "SELECT ...",
+                    "queryLimit": 1000,
+                }
+            )
+        }
+        resp = self.get_json_resp("/tabstateview/", data=data)
+        tab_state_id = resp["id"]
+
+        # run a query in the created tab
+        self.run_sql(
+            "SELECT name FROM birth_names",
+            "client_id_1",
+            username=username,
+            raise_on_error=True,
+            sql_editor_id=str(tab_state_id),
+        )
+        # run an orphan query (no tab)
+        self.run_sql(
+            "SELECT name FROM birth_names",
+            "client_id_2",
+            username=username,
+            raise_on_error=True,
+        )
+
+        # we should have only 1 query returned, since the second one is not
+        # associated with any tabs
+        resp = self.get_json_resp("/api/v1/sqllab/")
+        result = resp["result"]
+        self.assertEqual(len(result["queries"]), 1)
+
+    def test_get_access_denied(self):
+        new_role = Role(name="Dummy Role", permissions=[])
+        db.session.add(new_role)
+        db.session.commit()
+        unauth_user = self.create_user(
+            "unauth_user1",
+            "password",
+            "Dummy Role",
+            email=f"unauth_user1@superset.org",
+        )
+        self.login(username="unauth_user1", password="password")
+        rv = self.client.get("/api/v1/sqllab/")
+
+        assert rv.status_code == 403
+
+        db.session.delete(unauth_user)
+        db.session.delete(new_role)
+        db.session.commit()
+
     def test_estimate_required_params(self):
         self.login()