fix(sqllab): Inefficient SqlaTable query (#24343) (#24344)

apache · Jun 12, 2023 · 6d9df43 · 6d9df43
1 parent 0e3f1f6
commit 6d9df43
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/superset/databases/commands/tables.py b/superset/databases/commands/tables.py
@@ -17,6 +17,8 @@
 import logging
 from typing import Any, cast
 
+from sqlalchemy.orm import lazyload, load_only
+
 from superset.commands.base import BaseCommand
 from superset.connectors.sqla.models import SqlaTable
 from superset.databases.commands.exceptions import (
@@ -74,10 +76,18 @@ def run(self) -> dict[str, Any]:
             extra_dict_by_name = {
                 table.name: table.extra_dict
                 for table in (
-                    db.session.query(SqlaTable).filter(
+                    db.session.query(SqlaTable)
+                    .filter(
                         SqlaTable.database_id == self._model.id,
                         SqlaTable.schema == self._schema_name,
                     )
+                    .options(
+                        load_only(
+                            SqlaTable.schema, SqlaTable.table_name, SqlaTable.extra
+                        ),
+                        lazyload(SqlaTable.columns),
+                        lazyload(SqlaTable.metrics),
+                    )
                 ).all()
             }
 

diff --git a/superset/views/core.py b/superset/views/core.py
@@ -39,6 +39,7 @@
 from flask_babel import gettext as __, lazy_gettext as _
 from sqlalchemy import and_, or_
 from sqlalchemy.exc import DBAPIError, NoSuchModuleError, SQLAlchemyError
+from sqlalchemy.orm import lazyload, load_only
 
 from superset import (
     app,
@@ -1014,10 +1015,16 @@ def tables(  # pylint: disable=no-self-use
         extra_dict_by_name = {
             table.name: table.extra_dict
             for table in (
-                db.session.query(SqlaTable).filter(
+                db.session.query(SqlaTable)
+                .filter(
                     SqlaTable.database_id == database.id,
                     SqlaTable.schema == schema_parsed,
                 )
+                .options(
+                    load_only(SqlaTable.schema, SqlaTable.table_name, SqlaTable.extra),
+                    lazyload(SqlaTable.columns),
+                    lazyload(SqlaTable.metrics),
+                )
             ).all()
         }