From 38808a8cf9bfa06497516e650766550447ad0d4a Mon Sep 17 00:00:00 2001
From: Diego Medina <diegomedina24@gmail.com>
Date: Tue, 26 Apr 2022 14:22:03 -0300
Subject: [PATCH] Fix test

---
 superset/queries/schemas.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/superset/queries/schemas.py b/superset/queries/schemas.py
index 7dbe71d808641..a6403cdbd2338 100644
--- a/superset/queries/schemas.py
+++ b/superset/queries/schemas.py
@@ -14,10 +14,14 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from marshmallow import fields
+from typing import List
+
+from marshmallow import fields, Schema
 from marshmallow_sqlalchemy import SQLAlchemyAutoSchema
 
+from superset.dashboards.schemas import UserSchema
 from superset.models.sql_lab import Query
+from superset.sql_parse import Table
 
 openapi_spec_methods_override = {
     "get": {"get": {"description": "Get query detail information."}},
@@ -31,6 +35,10 @@
 }
 
 
+class DatabaseSchema(Schema):
+    database_name = fields.String()
+
+
 class QuerySchema(SQLAlchemyAutoSchema):
     """
     Schema for the ``Query`` model.
@@ -38,8 +46,14 @@ class QuerySchema(SQLAlchemyAutoSchema):
 
     start_time = fields.Float(attribute="start_time")
     end_time = fields.Float(attribute="end_time")
+    user = fields.Nested(UserSchema)
+    database = fields.Nested(DatabaseSchema)
+    sql_tables = fields.Method("get_sql_tables")
 
     class Meta:  # pylint: disable=too-few-public-methods
         model = Query
         load_instance = True
         include_relationships = True
+
+    def get_sql_tables(_, obj: Query) -> List[Table]:
+        return obj.sql_tables