Attempt to fix nested model retrival

opensensor · Oct 3, 2024 · 98dea19 · 98dea19
1 parent b24ca81
commit 98dea19
Showing 1 changed file with 24 additions and 9 deletions.
diff --git a/opensensor/collection_apis.py b/opensensor/collection_apis.py
@@ -206,7 +206,7 @@ def create_nested_pipeline(model: Type[BaseModel], prefix="", pipeline=None):
             "timestamp": "$timestamp",
         }
 
-    for field_name, _ in model.__fields__.items():
+    for field_name, field_type in model.__fields__.items():
         if field_name == "timestamp":
             continue
         lookup_field = (
@@ -221,11 +221,25 @@ def create_nested_pipeline(model: Type[BaseModel], prefix="", pipeline=None):
             match_conditions[unit_field_name] = {"$exists": True}
         elif field_name in nested_fields:
             nested_model = nested_fields[field_name]
-            nested_prefix = f"{full_mongo_field_name}."
-            nested_pipeline, nested_match = create_nested_pipeline(nested_model, nested_prefix)
-            pipeline[field_name] = nested_pipeline
-            for k, v in nested_match.items():
-                match_conditions[f"{nested_prefix}{k}"] = v
+            if get_origin(field_type) is List:
+                # Handle array of nested objects (like RelayStatus)
+                nested_pipeline, nested_match = create_nested_pipeline(nested_model, "")
+                pipeline[field_name] = {
+                    "$map": {
+                        "input": f"${full_mongo_field_name}",
+                        "as": "item",
+                        "in": {
+                            k: f"$$item.{v.replace('$', '')}" for k, v in nested_pipeline.items()
+                        },
+                    }
+                }
+                match_conditions[full_mongo_field_name] = {"$exists": True}
+            else:
+                nested_prefix = f"{full_mongo_field_name}."
+                nested_pipeline, nested_match = create_nested_pipeline(nested_model, nested_prefix)
+                pipeline[field_name] = nested_pipeline
+                for k, v in nested_match.items():
+                    match_conditions[f"{nested_prefix}{k}"] = v
         else:
             pipeline[field_name] = f"${full_mongo_field_name}"
             match_conditions[full_mongo_field_name] = {"$exists": True}
@@ -257,10 +271,11 @@ def create_model_instance(model: Type[BaseModel], data: dict, target_unit: Optio
             instance_data[field_name] = data.get(unit_field)
         elif field_name in nested_fields:
             nested_model = nested_fields[field_name]
-            if isinstance(field.type_, List):
+            if get_origin(field.type_) is List:
+                # Handle array of nested objects (like RelayStatus)
+                nested_data = data.get(field_name, [])
                 instance_data[field_name] = [
-                    create_model_instance(nested_model, item, target_unit)
-                    for item in data.get(field_name, [])
+                    create_model_instance(nested_model, item, target_unit) for item in nested_data
                 ]
             else:
                 nested_data = data.get(field_name, {})