def test_fetch_data_with_date_dimensions(self, reduce_mock): database = MagicMock() database.fetch_dataframes.return_value = [] dimensions = [number_field, date_field, boolean_field, text_field] fetch_data(database, self.test_queries, dimensions) database.fetch_dataframes.assert_called_with( 'SELECT * FROM "politics"."politician"', 'SELECT * FROM "politics"."hints"', parse_dates={'$date': {}})
def test_fetch_data_strips_rows_over_max_result_set_size( self, reduce_mock): database = MagicMock(max_result_set_size=2) database.fetch_dataframes.return_value = [ pd.DataFrame([{ "a": 1.0 }, { "a": 2.0 }, { "a": 3.0 }]), pd.DataFrame([{ "a": 1.0 }]), ] reduce_mock.side_effect = lambda *args: args[ 0] # let the reduce mock pass on the dataframes unchanged max_rows_returned, result = fetch_data(database, self.test_queries, self.test_dimensions) database.fetch_dataframes.assert_called_with( 'SELECT * FROM "politics"."politician"', 'SELECT * FROM "politics"."hints"', parse_dates={}) assert_frame_equal(result[0], pd.DataFrame([{"a": 1.0}, {"a": 2.0}])) assert_frame_equal(result[1], pd.DataFrame([{"a": 1.0}])) self.assertEqual(max_rows_returned, 3)
def fetch(self, hint=None, force_include=()) -> List[str]: """ Fetch the data for this query and transform it into the widgets. :param hint: For database vendors that support it, add a query hint to collect analytics on the queries triggered by fireant. :param force_include: A list of dimension values to include in the result set. This can be used to avoid having necessary results cut off due to the pagination. These results will be returned at the head of the results. :return: A list of dict (JSON) objects containing the widget configurations. """ query = add_hints(self.sql, hint)[0] dimension = self.dimensions[0] alias_definition = dimension.definition.as_( alias_selector(dimension.alias)) dimension_definition = dimension.definition if self.hint_table: alias_definition = alias_definition.replace_table( alias_definition.table, self.hint_table) dimension_definition = dimension.definition.replace_table( dimension_definition.table, self.hint_table) if force_include: include = self.dataset.database.to_char(dimension_definition).isin( [str(x) for x in force_include]) # Ensure that these values are included query = query.orderby(include, order=Order.desc) # Filter out NULL values from choices query = query.where(dimension_definition.notnull()) # Order by the dimension definition that the choices are for query = query.orderby(alias_definition) max_rows_returned, data = fetch_data(self.dataset.database, [query], self.dimensions) if len(data.index.names) > 1: display_alias = data.index.names[1] data.reset_index(display_alias, inplace=True) choices = data[display_alias] else: data["display"] = data.index.tolist() choices = data["display"] dimension_display = self.dimensions[-1] choices = choices.map( lambda raw: display_value(raw, dimension_display) or raw) return self._transform_for_return(choices, max_rows_returned=max_rows_returned)
def test_fetch_data(self, reduce_mock): database = MagicMock(max_result_set_size=1000) database.fetch_dataframes.return_value = [ self.test_result_a, self.test_result_b, ] mocked_result = pd.DataFrame([{"a": 1.0}]) reduce_mock.return_value = mocked_result max_rows_returned, result = fetch_data(database, self.test_queries, self.test_dimensions) pandas.testing.assert_frame_equal(mocked_result, result) database.fetch_dataframes.assert_called_with( 'SELECT * FROM "politics"."politician"', 'SELECT * FROM "politics"."hints"', parse_dates={}) reduce_mock.assert_called_once_with( [self.test_result_a, self.test_result_b], (), self.test_dimensions, ())
def fetch(self, hint=None): queries = add_hints(self.sql, hint) max_rows_returned, data = fetch_data(self.dataset.database, queries, self.dimensions) data = self._get_latest_data_from_df(data) return self._transform_for_return(data, max_rows_returned=max_rows_returned)