Add unmoving at reference location counts to api, and remove some sam…

…pling from top level queries

flowmachine/flowmachine/core/server/query_schemas/active_at_reference_location.py

@@ -4,7 +4,6 @@
 
 from marshmallow import fields, post_load
 from marshmallow.validate import OneOf
-from marshmallow_oneofschema import OneOfSchema
 
 from flowmachine.features.subscriber.active_at_reference_location import (
     ActiveAtReferenceLocation,
@@ -16,19 +15,10 @@
 
 __all__ = ["ActiveAtReferenceLocationSchema", "ActiveAtReferenceLocationExposed"]
 
-from .daily_location import DailyLocationSchema
-from .modal_location import ModalLocationSchema
+from .reference_location import ReferenceLocationSchema
 from .unique_locations import UniqueLocationsSchema
 
 
-class ReferenceLocationSchema(OneOfSchema):
-    type_field = "query_kind"
-    type_schemas = {
-        "daily_location": DailyLocationSchema,
-        "modal_location": ModalLocationSchema,
-    }
-
-
 class ActiveAtReferenceLocationSchema(BaseQueryWithSamplingSchema):
     # query_kind parameter is required here for claims validation
     query_kind = fields.String(validate=OneOf(["active_at_reference_location"]))

flowmachine/flowmachine/core/server/query_schemas/consecutive_trips.py

@@ -2,25 +2,22 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marshmallow import fields, post_load
+from marshmallow import fields, post_load, Schema
 from marshmallow.validate import OneOf
 
 from flowmachine.features.location.redacted_consecutive_trips import (
     RedactedConsecutiveTrips,
 )
 from flowmachine.features.utilities.subscriber_locations import SubscriberLocations
 from flowmachine.features.location.consecutive_trips import ConsecutiveTrips
+from . import BaseExposedQuery
 from .custom_fields import SubscriberSubset
 from .aggregation_unit import AggregationUnit, get_spatial_unit_obj
-from .base_query_with_sampling import (
-    BaseQueryWithSamplingSchema,
-    BaseExposedQueryWithSampling,
-)
 
 __all__ = ["ConsecutiveTripsSchema", "ConsecutiveTripsExposed"]
 
 
-class ConsecutiveTripsSchema(BaseQueryWithSamplingSchema):
+class ConsecutiveTripsSchema(Schema):
     # query_kind parameter is required here for claims validation
     query_kind = fields.String(validate=OneOf(["consecutive_trips"]))
     start_date = fields.Date(required=True)
@@ -33,7 +30,7 @@ def make_query_object(self, params, **kwargs):
         return ConsecutiveTripsExposed(**params)
 
 
-class ConsecutiveTripsExposed(BaseExposedQueryWithSampling):
+class ConsecutiveTripsExposed(BaseExposedQuery):
     def __init__(
         self,
         start_date,

flowmachine/flowmachine/core/server/query_schemas/flowmachine_query.py

@@ -34,6 +34,9 @@
 from .total_network_objects import TotalNetworkObjectsSchema
 from .dfs_metric_total_amount import DFSTotalMetricAmountSchema
 from .unique_visitor_counts import UniqueVisitorCountsSchema
+from .unmoving_at_reference_location_counts import (
+    UnmovingAtReferenceLocationCountsSchema,
+)
 from .unmoving_counts import UnmovingCountsSchema
 
 
@@ -59,6 +62,7 @@ class FlowmachineQuerySchema(OneOfSchema):
         "unique_visitor_counts": UniqueVisitorCountsSchema,
         "consecutive_trips": ConsecutiveTripsSchema,
         "unmoving_counts": UnmovingCountsSchema,
+        "unmoving_at_reference_location_counts": UnmovingAtReferenceLocationCountsSchema,
     }
 
 

flowmachine/flowmachine/core/server/query_schemas/reference_location.py

@@ -0,0 +1,16 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from marshmallow_oneofschema import OneOfSchema
+
+from flowmachine.core.server.query_schemas.daily_location import DailyLocationSchema
+from flowmachine.core.server.query_schemas.modal_location import ModalLocationSchema
+
+
+class ReferenceLocationSchema(OneOfSchema):
+    type_field = "query_kind"
+    type_schemas = {
+        "daily_location": DailyLocationSchema,
+        "modal_location": ModalLocationSchema,
+    }

flowmachine/flowmachine/core/server/query_schemas/unique_location_counts.py

@@ -2,21 +2,18 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marshmallow import fields, post_load
-from marshmallow.validate import OneOf, Length
+from marshmallow import fields, post_load, Schema
+from marshmallow.validate import OneOf
 
 from flowmachine.features import UniqueLocationCounts
+from . import BaseExposedQuery
 from .custom_fields import SubscriberSubset
 from .aggregation_unit import AggregationUnit, get_spatial_unit_obj
-from .base_query_with_sampling import (
-    BaseQueryWithSamplingSchema,
-    BaseExposedQueryWithSampling,
-)
 
 __all__ = ["UniqueLocationCountsSchema", "UniqueLocationCountsExposed"]
 
 
-class UniqueLocationCountsSchema(BaseQueryWithSamplingSchema):
+class UniqueLocationCountsSchema(Schema):
     query_kind = fields.String(validate=OneOf(["unique_location_counts"]))
     start_date = fields.Date(required=True)
     end_date = fields.Date(required=True)
@@ -28,7 +25,7 @@ def make_query_object(self, params, **kwargs):
         return UniqueLocationCountsExposed(**params)
 
 
-class UniqueLocationCountsExposed(BaseExposedQueryWithSampling):
+class UniqueLocationCountsExposed(BaseExposedQuery):
     def __init__(
         self,
         *,
@@ -44,10 +41,9 @@ def __init__(
         self.end_date = end_date
         self.aggregation_unit = aggregation_unit
         self.subscriber_subset = subscriber_subset
-        self.sampling = sampling
 
     @property
-    def _unsampled_query_obj(self):
+    def _flowmachine_query_obj(self):
         """
         Return the underlying flowmachine unique_location_counts object.
 

flowmachine/flowmachine/core/server/query_schemas/unmoving_at_reference_location_counts.py

@@ -0,0 +1,64 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from marshmallow import fields, post_load, Schema
+from marshmallow.validate import OneOf
+
+from flowmachine.features.location.redacted_unmoving_at_reference_location_counts import (
+    RedactedUnmovingAtReferenceLocationCounts,
+)
+from flowmachine.features.location.unmoving_at_reference_location_counts import (
+    UnmovingAtReferenceLocationCounts,
+)
+from flowmachine.features.subscriber.unmoving_at_reference_location import (
+    UnmovingAtReferenceLocation,
+)
+from . import BaseExposedQuery
+from .reference_location import ReferenceLocationSchema
+from .base_query_with_sampling import BaseQueryWithSamplingSchema
+from .unique_locations import UniqueLocationsSchema
+
+__all__ = [
+    "UnmovingAtReferenceLocationCountsSchema",
+    "UnmovingAtReferenceLocationCountsExposed",
+]
+
+
+class UnmovingAtReferenceLocationCountsSchema(Schema):
+    # query_kind parameter is required here for claims validation
+    query_kind = fields.String(
+        validate=OneOf(["unmoving_at_reference_location_counts"])
+    )
+    locations = fields.Nested(UniqueLocationsSchema)
+    reference_locations = fields.Nested(ReferenceLocationSchema)
+
+    @post_load
+    def make_query_object(self, params, **kwargs):
+        return UnmovingAtReferenceLocationCountsExposed(**params)
+
+
+class UnmovingAtReferenceLocationCountsExposed(BaseExposedQuery):
+    def __init__(self, locations, reference_locations):
+        # Note: all input parameters need to be defined as attributes on `self`
+        # so that marshmallow can serialise the object correctly.
+        self.locations = locations
+        self.reference_locations = reference_locations
+
+    @property
+    def _flowmachine_query_obj(self):
+        """
+        Return the underlying flowmachine object.
+
+        Returns
+        -------
+        Query
+        """
+        return RedactedUnmovingAtReferenceLocationCounts(
+            unmoving_at_reference_location_counts=UnmovingAtReferenceLocationCounts(
+                UnmovingAtReferenceLocation(
+                    locations=self.locations._flowmachine_query_obj,
+                    reference_locations=self.reference_locations._flowmachine_query_obj,
+                )
+            )
+        )

flowmachine/flowmachine/core/server/query_schemas/unmoving_counts.py

@@ -2,7 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marshmallow import fields, post_load
+from marshmallow import fields, post_load, Schema
 from marshmallow.validate import OneOf
 
 from flowmachine.features.location.redacted_unmoving_counts import (
@@ -11,13 +11,12 @@
 from flowmachine.features.location.unmoving_counts import UnmovingCounts
 from flowmachine.features.subscriber.unmoving import Unmoving
 from . import BaseExposedQuery
-from .base_query_with_sampling import BaseQueryWithSamplingSchema
 from .unique_locations import UniqueLocationsSchema
 
 __all__ = ["UnmovingCountsSchema", "UnmovingCountsExposed"]
 
 
-class UnmovingCountsSchema(BaseQueryWithSamplingSchema):
+class UnmovingCountsSchema(Schema):
     # query_kind parameter is required here for claims validation
     query_kind = fields.String(validate=OneOf(["unmoving_counts"]))
     locations = fields.Nested(UniqueLocationsSchema)