Merge pull request #276 from RobinPapke/blended_dimension_choices

Update fetching choices for data blending

fireant/dataset/data_blending.py

@@ -1,12 +1,40 @@
 from fireant.dataset.fields import Field
 from fireant.dataset.klass import DataSet
-from fireant.queries.builder import DataSetBlenderQueryBuilder
+from fireant.queries.builder import (
+    DataSetBlenderQueryBuilder,
+    DimensionChoicesQueryBuilder,
+)
+from fireant.queries.builder.query_builder import validate_fields
 from fireant.utils import (
     immutable,
     deepcopy,
 )
 
 
+def _wrap_dataset_fields(dataset):
+    wrapped_fields = []
+    for field in dataset.fields:
+        wrapped_field = Field(
+              alias=field.alias,
+              definition=field,
+              data_type=field.data_type,
+              label=field.label,
+              hint_table=field.hint_table,
+              prefix=field.prefix,
+              suffix=field.suffix,
+              thousands=field.thousands,
+              precision=field.precision,
+              hyperlink_template=field.hyperlink_template,
+        )
+
+        if not field.definition.is_aggregate:
+            wrapped_field.choices = DimensionChoicesBlenderQueryBuilder(dataset, field)
+
+        wrapped_fields.append(wrapped_field)
+
+    return wrapped_fields
+
+
 class DataSetBlender:
     """
     The DataSetBlender class is the DataSet equivalent for implementing data blending, across distinct DataSet
@@ -36,24 +64,7 @@ def __init__(self, primary_dataset, secondary_dataset, field_map):
         #   1. DataSetBlender doesn't share a reference to a field with a DataSet
         #   2. When complex fields are added, the `definition` attribute will always have at least one field within
         #      its object graph
-        all_fields = [*secondary_dataset.fields, *primary_dataset.fields]
-        self.fields = DataSet.Fields(
-            [
-                Field(
-                    alias=field.alias,
-                    definition=field,
-                    data_type=field.data_type,
-                    label=field.label,
-                    hint_table=field.hint_table,
-                    prefix=field.prefix,
-                    suffix=field.suffix,
-                    thousands=field.thousands,
-                    precision=field.precision,
-                    hyperlink_template=field.hyperlink_template,
-                )
-                for field in all_fields
-            ]
-        )
+        self.fields = DataSet.Fields([*_wrap_dataset_fields(primary_dataset), *_wrap_dataset_fields(secondary_dataset)])
 
         # add query builder entry points
         self.query = DataSetBlenderQueryBuilder(self)
@@ -118,3 +129,13 @@ def on_dimensions(self):
             field_map[primary_ds_field] = secondary_ds_field
 
         return self.on(field_map)
+
+
+class DimensionChoicesBlenderQueryBuilder(DimensionChoicesQueryBuilder):
+    @immutable
+    def filter(self, *filters):
+        for filter_ in filters:
+            filter_.field = filter_.field.definition  # replace blender filter field with field of primary/secondary
+
+        validate_fields([filter_.field for filter_ in filters], self.dataset)
+        self._filters += [filter_ for filter_ in filters]

fireant/queries/builder/query_builder.py

@@ -30,28 +30,28 @@ def get_column_names(database, table):
     return {column_definition[0] for column_definition in column_definitions}
 
 
-def _strip_modifiers(fields):
-    for field in fields:
-        node = field
-        while hasattr(node, "dimension"):
-            node = node.dimension
-        yield node
-
-
-def _validate_fields(fields, dataset):
+def validate_fields(fields, dataset):
     fields = [find_field_in_modified_field(field) for field in fields]
 
     invalid = [field.alias for field in fields if field not in dataset.fields]
     if not invalid:
         return
 
     raise DataSetException(
-        "Only fields from dataset can be used in a dataset query. Found invalid fields: {}.".format(
-            ", ".join(invalid)
-        )
+          "Only fields from dataset can be used in a dataset query. Found invalid fields: {}.".format(
+                ", ".join(invalid)
+          )
     )
 
 
+def _strip_modifiers(fields):
+    for field in fields:
+        node = field
+        while hasattr(node, "dimension"):
+            node = node.dimension
+        yield node
+
+
 class QueryBuilder(object):
     """
     This is the base class for building dataset queries. This class provides an interface for building dataset queries
@@ -88,7 +88,7 @@ def dimension(self, *dimensions):
         :return:
             A copy of the query with the dimensions added.
         """
-        _validate_fields(dimensions, self.dataset)
+        validate_fields(dimensions, self.dataset)
         aliases = {dimension.alias for dimension in self._dimensions}
         self._dimensions += [
             dimension for dimension in dimensions if dimension.alias not in aliases
@@ -104,7 +104,7 @@ def filter(self, *filters):
         :return:
             A copy of the query with the filters added.
         """
-        _validate_fields([fltr.field for fltr in filters], self.dataset)
+        validate_fields([fltr.field for fltr in filters], self.dataset)
         self._filters += [f for f in filters]
 
     @immutable
@@ -117,7 +117,7 @@ def orderby(self, field: Field, orientation: Order = None):
         :return:
             A copy of the query with the order by added.
         """
-        _validate_fields([field], self.dataset)
+        validate_fields([field], self.dataset)
 
         if self._orders is None:
             self._orders = []
@@ -203,7 +203,7 @@ def reference(self, *references):
         :return:
             A copy of the query with the references added.
         """
-        _validate_fields([reference.field for reference in references], self.dataset)
+        validate_fields([reference.field for reference in references], self.dataset)
         self._references += references
 
 
@@ -232,8 +232,8 @@ def widget(self, *widgets):
         :return:
             A copy of the query with the widgets added.
         """
-        _validate_fields(
-            [field for widget in widgets for field in widget.metrics], self.dataset
+        validate_fields(
+              [field for widget in widgets for field in widget.metrics], self.dataset
         )
 
         self._widgets += widgets

fireant/tests/queries/test_dimension_choices.py

@@ -5,13 +5,15 @@
     patch,
 )
 
+from fireant import DataSetException
 from fireant.tests.dataset.matchers import (
     FieldMatcher,
     PypikaQueryMatcher,
 )
 from fireant.tests.dataset.mocks import (
     mock_dataset,
     mock_hint_dataset,
+    mock_dataset_blender,
 )
 
 
@@ -240,3 +242,43 @@ def test_query_choices_for_field(self, mock_fetch_data: Mock):
                                                                     'GROUP BY "$political_party" '
                                                                     'ORDER BY "$political_party"')],
                                                 FieldMatcher(mock_dataset.fields.political_party))
+
+# noinspection SqlDialectInspection,SqlNoDataSourceInspection
+@patch('fireant.queries.builder.dimension_choices_query_builder.fetch_data')
+class BlenderDimensionsChoicesFetchTests(TestCase):
+    def test_query_choices_for_primary_field_in_dataset_blender(self, mock_fetch_data: Mock):
+        mock_dataset_blender.fields.political_party \
+            .choices \
+            .filter(mock_dataset_blender.fields.political_party.isin(['d', 'r'])) \
+            .fetch()
+
+        mock_fetch_data.assert_called_once_with(ANY,
+                                                [PypikaQueryMatcher('SELECT "political_party" "$political_party" '
+                                                                    'FROM "politics"."politician" '
+                                                                    'WHERE "political_party" IN (\'d\',\'r\') '
+                                                                    'AND NOT "political_party" IS NULL '
+                                                                    'GROUP BY "$political_party" '
+                                                                    'ORDER BY "$political_party"')],
+                                                FieldMatcher(mock_dataset_blender.fields.political_party))
+
+    def test_query_choices_for_secondary_field_in_dataset_blender(self, mock_fetch_data: Mock):
+        mock_dataset_blender.fields.state \
+            .choices \
+            .filter(mock_dataset_blender.fields.state.isin(['Texas'])) \
+            .fetch()
+
+        mock_fetch_data.assert_called_once_with(ANY,
+                                                [PypikaQueryMatcher('SELECT "state" "$state" '
+                                                                    'FROM "politics"."politician_spend" '
+                                                                    'WHERE "state" IN (\'Texas\') '
+                                                                    'AND NOT "state" IS NULL '
+                                                                    'GROUP BY "$state" '
+                                                                    'ORDER BY "$state"')],
+                                                FieldMatcher(mock_dataset_blender.fields.state))
+
+    def test_query_choices_with_fields_from_different_datasets_raises_exception(self, mock_fetch_data: Mock):
+        with self.assertRaises(DataSetException):
+            mock_dataset_blender.fields.political_party \
+                .choices \
+                .filter(mock_dataset_blender.fields.state.isin(['Texas'])) \
+                .fetch()