From 98325ea46d0f874b83b4a7ea71c508f0d4e56e02 Mon Sep 17 00:00:00 2001
From: Mayur <mayurnewase111@gmail.com>
Date: Thu, 29 Sep 2022 17:21:54 +0530
Subject: [PATCH] pylint

---
 superset/common/query_context.py         |  6 ++++++
 superset/common/query_context_factory.py | 11 +++++++++++
 2 files changed, 17 insertions(+)

diff --git a/superset/common/query_context.py b/superset/common/query_context.py
index 4a91c6ad6db..3f366770993 100644
--- a/superset/common/query_context.py
+++ b/superset/common/query_context.py
@@ -27,6 +27,7 @@ from superset.common.query_context_processor import (
     QueryContextProcessor,
 )
 from superset.common.query_object import QueryObject
+from superset.models.slice import Slice
 
 if TYPE_CHECKING:
     from superset.connectors.base.models import BaseDatasource
@@ -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
@@ -64,6 +66,7 @@ class QueryContext:
         *,
         datasource: BaseDatasource,
         queries: List[QueryObject],
+        slice_: Optional[Slice],
         form_data: Optional[Dict[str, Any]],
         result_type: ChartDataResultType,
         result_format: ChartDataResultFormat,
@@ -72,6 +75,7 @@ class QueryContext:
         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
@@ -98,6 +102,8 @@ class QueryContext:
     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"):
diff --git a/superset/common/query_context_factory.py b/superset/common/query_context_factory.py
index 2056109bbff..11dba36be23 100644
--- a/superset/common/query_context_factory.py
+++ b/superset/common/query_context_factory.py
@@ -19,10 +19,12 @@ from __future__ import annotations
 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:
@@ -55,6 +57,11 @@ class QueryContextFactory:  # pylint: disable=too-few-public-methods
         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_ = [
@@ -72,6 +79,7 @@ class QueryContextFactory:  # pylint: disable=too-few-public-methods
         return QueryContext(
             datasource=datasource_model_instance,
             queries=queries_,
+            slice_=slice_,
             form_data=form_data,
             result_type=result_type,
             result_format=result_format,
@@ -85,3 +93,6 @@ class QueryContextFactory:  # pylint: disable=too-few-public-methods
         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)