Skip to content

Commit

Permalink
pylint
Browse files Browse the repository at this point in the history
  • Loading branch information
mayurnewase authored and AAfghahi committed Dec 15, 2022
1 parent 7390186 commit a366362
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 0 deletions.
6 changes: 6 additions & 0 deletions superset/common/query_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
QueryContextProcessor,
)
from superset.common.query_object import QueryObject
from superset.models.slice import Slice

if TYPE_CHECKING:
from superset.connectors.base.models import BaseDatasource
Expand All @@ -46,6 +47,7 @@ class QueryContext:
enforce_numerical_metrics: ClassVar[bool] = True

datasource: BaseDatasource
slice_: Optional[Slice] = None
queries: List[QueryObject]
form_data: Optional[Dict[str, Any]]
result_type: ChartDataResultType
Expand All @@ -64,6 +66,7 @@ def __init__(
*,
datasource: BaseDatasource,
queries: List[QueryObject],
slice_: Optional[Slice],
form_data: Optional[Dict[str, Any]],
result_type: ChartDataResultType,
result_format: ChartDataResultFormat,
Expand All @@ -72,6 +75,7 @@ def __init__(
cache_values: Dict[str, Any],
) -> None:
self.datasource = datasource
self.slice_ = slice_
self.result_type = result_type
self.result_format = result_format
self.queries = queries
Expand All @@ -98,6 +102,8 @@ def get_payload(
def get_cache_timeout(self) -> Optional[int]:
if self.custom_cache_timeout is not None:
return self.custom_cache_timeout
if self.slice_ and self.slice_.cache_timeout is not None:
return self.slice_.cache_timeout
if self.datasource.cache_timeout is not None:
return self.datasource.cache_timeout
if hasattr(self.datasource, "database"):
Expand Down
11 changes: 11 additions & 0 deletions superset/common/query_context_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,12 @@
from typing import Any, Dict, List, Optional, TYPE_CHECKING

from superset import app, db
from superset.charts.dao import ChartDAO
from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
from superset.common.query_context import QueryContext
from superset.common.query_object_factory import QueryObjectFactory
from superset.connectors.connector_registry import ConnectorRegistry
from superset.models.slice import Slice
from superset.utils.core import DatasourceDict

if TYPE_CHECKING:
Expand Down Expand Up @@ -55,6 +57,11 @@ def create(
datasource_model_instance = None
if datasource:
datasource_model_instance = self._convert_to_model(datasource)

slice_ = None
if form_data and form_data.get("slice_id") is not None:
slice_ = self._get_slice(form_data.get("slice_id"))

result_type = result_type or ChartDataResultType.FULL
result_format = result_format or ChartDataResultFormat.JSON
queries_ = [
Expand All @@ -72,6 +79,7 @@ def create(
return QueryContext(
datasource=datasource_model_instance,
queries=queries_,
slice_=slice_,
form_data=form_data,
result_type=result_type,
result_format=result_format,
Expand All @@ -85,3 +93,6 @@ def _convert_to_model(self, datasource: DatasourceDict) -> BaseDatasource:
return ConnectorRegistry.get_datasource(
str(datasource["type"]), int(datasource["id"]), db.session
)

def _get_slice(self, slice_id: Any) -> Optional[Slice]:
return ChartDAO.find_by_id(slice_id)

0 comments on commit a366362

Please sign in to comment.