perf: Refactor Dashboard.datasets_trimmed_for_slices (#15648)

Co-authored-by: John Bodley <john.bodley@airbnb.com>
This commit is contained in:
John Bodley 2021-07-13 20:00:57 -07:00 committed by GitHub
parent 42a1061bfa
commit b3e699b767
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 37 additions and 10 deletions

View File

@ -18,8 +18,9 @@ from __future__ import annotations
import json import json
import logging import logging
from collections import defaultdict
from functools import partial from functools import partial
from typing import Any, Callable, Dict, List, Set, Union from typing import Any, Callable, Dict, List, Set, Tuple, Type, Union
import sqlalchemy as sqla import sqlalchemy as sqla
from flask_appbuilder import Model from flask_appbuilder import Model
@ -171,8 +172,21 @@ class Dashboard( # pylint: disable=too-many-instance-attributes
@property @property
def datasources(self) -> Set[BaseDatasource]: def datasources(self) -> Set[BaseDatasource]:
# pylint: disable=using-constant-test # Verbose but efficient database enumeration of dashboard datasources.
return {slc.datasource for slc in self.slices if slc.datasource} datasources_by_cls_model: Dict[Type["BaseDatasource"], Set[int]] = defaultdict(
set
)
for slc in self.slices:
datasources_by_cls_model[slc.cls_model].add(slc.datasource_id)
return {
datasource
for cls_model, datasource_ids in datasources_by_cls_model.items()
for datasource in db.session.query(cls_model)
.filter(cls_model.id.in_(datasource_ids))
.all()
}
@property @property
def charts(self) -> List[BaseDatasource]: def charts(self) -> List[BaseDatasource]:
@ -246,13 +260,26 @@ class Dashboard( # pylint: disable=too-many-instance-attributes
unless=lambda: not is_feature_enabled("DASHBOARD_CACHE"), unless=lambda: not is_feature_enabled("DASHBOARD_CACHE"),
) )
def datasets_trimmed_for_slices(self) -> List[Dict[str, Any]]: def datasets_trimmed_for_slices(self) -> List[Dict[str, Any]]:
datasource_slices = utils.indexed(self.slices, "datasource") # Verbose but efficient database enumeration of dashboard datasources.
return [ slices_by_datasource: Dict[
# Filter out unneeded fields from the datasource payload Tuple[Type["BaseDatasource"], int], Set[Slice]
datasource.data_for_slices(slices) ] = defaultdict(set)
for datasource, slices in datasource_slices.items()
if datasource for slc in self.slices:
] slices_by_datasource[(slc.cls_model, slc.datasource_id)].add(slc)
result: List[Dict[str, Any]] = []
for (cls_model, datasource_id), slices in slices_by_datasource.items():
datasource = (
db.session.query(cls_model).filter_by(id=datasource_id).one_or_none()
)
if datasource:
# Filter out unneeded fields from the datasource payload
result.append(datasource.data_for_slices(slices))
return result
@property # type: ignore @property # type: ignore
def params(self) -> str: # type: ignore def params(self) -> str: # type: ignore