feat: query REST API more fields and tests (#11482)

* feat: query REST API more fields and tests

* fix missing field

superset/models/sql_lab.py

@@ -156,6 +156,10 @@ class Query(Model, ExtraJSONMixin):
     def username(self) -> str:
         return self.user.username
 
+    @property
+    def sql_tables(self) -> List[Table]:
+        return list(ParsedQuery(self.sql).tables)
+
     def raise_for_access(self) -> None:
         """
         Raise an exception if the user cannot access the resource.

superset/queries/api.py

@@ -19,6 +19,7 @@ import logging
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 
 from superset.constants import RouteMethod
+from superset.databases.filters import DatabaseFilter
 from superset.models.sql_lab import Query
 from superset.queries.filters import QueryFilter
 from superset.queries.schemas import openapi_spec_methods_override
@@ -37,9 +38,18 @@ class QueryRestApi(BaseSupersetModelRestApi):
 
     class_permission_name = "QueryView"
     list_columns = [
-        "user.username",
+        "changed_on",
         "database.database_name",
+        "rows",
+        "schema",
+        "sql",
+        "sql_tables",
         "status",
+        "tab_name",
+        "user.first_name",
+        "user.id",
+        "user.last_name",
+        "user.username",
         "start_time",
         "end_time",
         "rows",
@@ -47,30 +57,30 @@ class QueryRestApi(BaseSupersetModelRestApi):
         "tracking_url",
     ]
     show_columns = [
+        "changed_on",
         "client_id",
-        "tmp_table_name",
-        "tmp_schema_name",
-        "status",
-        "tab_name",
-        "sql_editor_id",
         "database.id",
-        "schema",
+        "end_result_backend_time",
-        "sql",
+        "end_time",
-        "select_sql",
+        "error_message",
         "executed_sql",
         "limit",
+        "progress",
+        "results_key",
+        "rows",
+        "schema",
         "select_as_cta",
         "select_as_cta_used",
-        "progress",
+        "select_sql",
-        "rows",
+        "sql",
-        "error_message",
+        "sql_editor_id",
-        "results_key",
-        "start_time",
         "start_running_time",
-        "end_time",
+        "start_time",
-        "end_result_backend_time",
+        "status",
+        "tab_name",
+        "tmp_schema_name",
+        "tmp_table_name",
         "tracking_url",
-        "changed_on",
     ]
     base_filters = [["id", QueryFilter, lambda: []]]
     base_order = ("changed_on", "desc")
@@ -78,7 +88,21 @@ class QueryRestApi(BaseSupersetModelRestApi):
     openapi_spec_tag = "Queries"
     openapi_spec_methods = openapi_spec_methods_override
 
+    order_columns = [
+        "changed_on",
+        "database.database_name",
+        "rows",
+        "schema",
+        "sql",
+        "tab_name",
+        "user.first_name",
+    ]
+
     related_field_filters = {
         "created_by": RelatedFieldFilter("first_name", FilterRelatedOwners),
     }
-    allowed_rel_fields = {"user"}
+
+    search_columns = ["changed_on", "database", "sql", "status", "user"]
+
+    filter_rel_fields = {"database": [["id", DatabaseFilter, lambda: []]]}
+    allowed_rel_fields = {"database", "user"}

tests/queries/api_tests.py

@@ -16,22 +16,25 @@
 # under the License.
 # isort:skip_file
 """Unit tests for Superset"""
+from datetime import datetime, timedelta
 import json
 import random
 import string
-from typing import Dict, Any
 
+import pytest
 import prison
 from sqlalchemy.sql import func
 
 import tests.test_app
 from superset import db, security_manager
 from superset.models.core import Database
-from superset.utils.core import get_example_database
+from superset.utils.core import get_example_database, get_main_database, QueryStatus
 from superset.models.sql_lab import Query
 
 from tests.base_tests import SupersetTestCase
 
+QUERIES_FIXTURE_COUNT = 10
+
 
 class TestQueryApi(SupersetTestCase):
     def insert_query(
@@ -67,6 +70,45 @@ class TestQueryApi(SupersetTestCase):
         db.session.commit()
         return query
 
+    @pytest.fixture()
+    def create_queries(self):
+        with self.create_app().app_context():
+            queries = []
+            admin_id = self.get_user("admin").id
+            alpha_id = self.get_user("alpha").id
+            example_database_id = get_example_database().id
+            main_database_id = get_main_database().id
+            for cx in range(QUERIES_FIXTURE_COUNT - 1):
+                queries.append(
+                    self.insert_query(
+                        example_database_id,
+                        admin_id,
+                        self.get_random_string(),
+                        sql=f"SELECT col1, col2 from table{cx}",
+                        rows=cx,
+                        status=QueryStatus.SUCCESS
+                        if (cx % 2) == 0
+                        else QueryStatus.RUNNING,
+                    )
+                )
+            queries.append(
+                self.insert_query(
+                    main_database_id,
+                    alpha_id,
+                    self.get_random_string(),
+                    sql=f"SELECT col1, col2 from table{QUERIES_FIXTURE_COUNT}",
+                    rows=QUERIES_FIXTURE_COUNT,
+                    status=QueryStatus.SUCCESS,
+                )
+            )
+
+            yield queries
+
+            # rollback changes
+            for query in queries:
+                db.session.delete(query)
+            db.session.commit()
+
     @staticmethod
     def get_random_string(length: int = 10):
         letters = string.ascii_letters
@@ -74,7 +116,7 @@ class TestQueryApi(SupersetTestCase):
 
     def test_get_query(self):
         """
-            Query API: Test get query
+        Query API: Test get query
         """
         admin = self.get_user("admin")
         client_id = self.get_random_string()
@@ -131,7 +173,7 @@ class TestQueryApi(SupersetTestCase):
 
     def test_get_query_not_found(self):
         """
-            Query API: Test get query not found
+        Query API: Test get query not found
         """
         admin = self.get_user("admin")
         client_id = self.get_random_string()
@@ -144,7 +186,7 @@ class TestQueryApi(SupersetTestCase):
 
     def test_get_query_no_data_access(self):
         """
-            Query API: Test get dashboard without data access
+        Query API: Test get query without data access
         """
         gamma1 = self.create_user(
             "gamma_1", "password", "Gamma", email="gamma1@superset.org"
@@ -198,34 +240,131 @@ class TestQueryApi(SupersetTestCase):
         db.session.delete(gamma2)
         db.session.commit()
 
-    def test_get_query_filter(self):
+    @pytest.mark.usefixtures("create_queries")
+    def test_get_list_query(self):
         """
-            Query API: Test get queries filter
+        Query API: Test get list query
         """
-        admin = self.get_user("admin")
-        client_id = self.get_random_string()
-        query = self.insert_query(
-            get_example_database().id,
-            admin.id,
-            client_id,
-            sql="SELECT col1, col2 from table1",
-        )
-
         self.login(username="admin")
-        arguments = {"filters": [{"col": "sql", "opr": "sw", "value": "SELECT col1"}]}
+        uri = "api/v1/query/"
-        uri = f"api/v1/query/?q={prison.dumps(arguments)}"
         rv = self.client.get(uri)
         self.assertEqual(rv.status_code, 200)
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(data["count"], 1)
+        assert data["count"] == QUERIES_FIXTURE_COUNT
+        # check expected columns
+        assert sorted(list(data["result"][0].keys())) == [
+            "changed_on",
+            "database",
+            "end_time",
+            "rows",
+            "schema",
+            "sql",
+            "sql_tables",
+            "start_time",
+            "status",
+            "tab_name",
+            "tmp_table_name",
+            "tracking_url",
+            "user",
+        ]
+        assert sorted(list(data["result"][0]["user"].keys())) == [
+            "first_name",
+            "id",
+            "last_name",
+            "username",
+        ]
+        assert list(data["result"][0]["database"].keys()) == [
+            "database_name",
+        ]
 
-        # rollback changes
+    @pytest.mark.usefixtures("create_queries")
-        db.session.delete(query)
+    def test_get_list_query_filter_sql(self):
-        db.session.commit()
-
-    def test_get_queries_no_data_access(self):
         """
-            Query API: Test get queries no data access
+        Query API: Test get list query filter
+        """
+        self.login(username="admin")
+        arguments = {"filters": [{"col": "sql", "opr": "ct", "value": "table2"}]}
+        uri = f"api/v1/query/?q={prison.dumps(arguments)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 200
+        data = json.loads(rv.data.decode("utf-8"))
+        assert data["count"] == 1
+
+    @pytest.mark.usefixtures("create_queries")
+    def test_get_list_query_filter_database(self):
+        """
+        Query API: Test get list query filter database
+        """
+        self.login(username="admin")
+        database_id = get_main_database().id
+        arguments = {
+            "filters": [{"col": "database", "opr": "rel_o_m", "value": database_id}]
+        }
+        uri = f"api/v1/query/?q={prison.dumps(arguments)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 200
+        data = json.loads(rv.data.decode("utf-8"))
+        assert data["count"] == 1
+
+    @pytest.mark.usefixtures("create_queries")
+    def test_get_list_query_filter_user(self):
+        """
+        Query API: Test get list query filter user
+        """
+        self.login(username="admin")
+        alpha_id = self.get_user("alpha").id
+        arguments = {"filters": [{"col": "user", "opr": "rel_o_m", "value": alpha_id}]}
+        uri = f"api/v1/query/?q={prison.dumps(arguments)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 200
+        data = json.loads(rv.data.decode("utf-8"))
+        assert data["count"] == 1
+
+    @pytest.mark.usefixtures("create_queries")
+    def test_get_list_query_filter_changed_on(self):
+        """
+        Query API: Test get list query filter changed_on
+        """
+        self.login(username="admin")
+        now_time = datetime.now()
+        yesterday_time = now_time - timedelta(days=1)
+        arguments = {
+            "filters": [
+                {"col": "changed_on", "opr": "lt", "value": str(now_time)},
+                {"col": "changed_on", "opr": "gt", "value": str(yesterday_time)},
+            ]
+        }
+        uri = f"api/v1/query/?q={prison.dumps(arguments)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 200
+        data = json.loads(rv.data.decode("utf-8"))
+        assert data["count"] == QUERIES_FIXTURE_COUNT
+
+    @pytest.mark.usefixtures("create_queries")
+    def test_get_list_query_order(self):
+        """
+        Query API: Test get list query filter changed_on
+        """
+        self.login(username="admin")
+        order_columns = [
+            "changed_on",
+            "database.database_name",
+            "rows",
+            "schema",
+            "sql",
+            "tab_name",
+            "user.first_name",
+        ]
+
+        for order_column in order_columns:
+            arguments = {"order_column": order_column, "order_direction": "asc"}
+            uri = f"api/v1/query/?q={prison.dumps(arguments)}"
+            rv = self.client.get(uri)
+            assert rv.status_code == 200
+
+    def test_get_list_query_no_data_access(self):
+        """
+        Query API: Test get queries no data access
         """
         admin = self.get_user("admin")
         client_id = self.get_random_string()
@@ -240,9 +379,9 @@ class TestQueryApi(SupersetTestCase):
         arguments = {"filters": [{"col": "sql", "opr": "sw", "value": "SELECT col1"}]}
         uri = f"api/v1/query/?q={prison.dumps(arguments)}"
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(data["count"], 0)
+        assert data["count"] == 0
 
         # rollback changes
         db.session.delete(query)