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)