def get_query_result(self, query_object: QueryObject) -> Dict[str, Any]: """Returns a pandas dataframe based on the query object""" # Here, we assume that all the queries will use the same datasource, which is # a valid assumption for current setting. In the long term, we may # support multiple queries from different data sources. timestamp_format = None if self.datasource.type == "table": dttm_col = self.datasource.get_column(query_object.granularity) if dttm_col: timestamp_format = dttm_col.python_date_format # The datasource here can be different backend but the interface is common result = self.datasource.query(query_object.to_dict()) df = result.df # Transform the timestamp we received from database to pandas supported # datetime format. If no python_date_format is specified, the pattern will # be considered as the default ISO date format # If the datetime format is unix, the parse will use the corresponding # parsing logic if not df.empty: normalize_dttm_col( df=df, timestamp_format=timestamp_format, offset=self.datasource.offset, time_shift=query_object.time_shift, ) if self.enforce_numerical_metrics: self.df_metrics_to_num(df, query_object) df.replace([np.inf, -np.inf], np.nan, inplace=True) df = query_object.exec_post_processing(df) return { "query": result.query, "status": result.status, "error_message": result.error_message, "df": df, }
def query_cache_key(self, query_obj: QueryObject, **kwargs: Any) -> Optional[str]: """ Returns a QueryObject cache key for objects in self.queries """ extra_cache_keys = self.datasource.get_extra_cache_keys(query_obj.to_dict()) cache_key = ( query_obj.cache_key( datasource=self.datasource.uid, extra_cache_keys=extra_cache_keys, rls=security_manager.get_rls_ids(self.datasource) if is_feature_enabled("ROW_LEVEL_SECURITY") and self.datasource.is_rls_supported else [], changed_on=self.datasource.changed_on, **kwargs, ) if query_obj else None ) return cache_key
def get_query_result(self, query_object: QueryObject) -> QueryResult: """Returns a pandas dataframe based on the query object""" query_context = self._query_context # Here, we assume that all the queries will use the same datasource, which is # a valid assumption for current setting. In the long term, we may # support multiple queries from different data sources. # The datasource here can be different backend but the interface is common result = query_context.datasource.query(query_object.to_dict()) query = result.query + ";\n\n" df = result.df # Transform the timestamp we received from database to pandas supported # datetime format. If no python_date_format is specified, the pattern will # be considered as the default ISO date format # If the datetime format is unix, the parse will use the corresponding # parsing logic if not df.empty: df = self.normalize_df(df, query_object) if query_object.time_offsets: time_offsets = self.processing_time_offsets(df, query_object) df = time_offsets["df"] queries = time_offsets["queries"] query += ";\n\n".join(queries) query += ";\n\n" # Re-raising QueryObjectValidationError try: df = query_object.exec_post_processing(df) except InvalidPostProcessingError as ex: raise QueryObjectValidationError from ex result.df = df result.query = query result.from_dttm = query_object.from_dttm result.to_dttm = query_object.to_dttm return result