Skip to content

Commit

Permalink
fix(sqllab): Inefficient SqlaTable query (#24343)
Browse files Browse the repository at this point in the history
Use a combination of lazyload and load_only to ensure sqlalchemy doesn't
eagerly join to additional tables, which was happening in a very
inefficient way for a large number of tables and resulting in this
endpoint never returning (in addition to high db load)
  • Loading branch information
giftig committed Jun 12, 2023
1 parent 0e3f1f6 commit bb5d308
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 2 deletions.
12 changes: 11 additions & 1 deletion superset/databases/commands/tables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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()
}

Expand Down
9 changes: 8 additions & 1 deletion superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
}

Expand Down

0 comments on commit bb5d308

Please sign in to comment.