fix(firestore): Add ID to all documents in Firestore provider (#94)

Noticed a bug in the Firestore provider while taking this on a test run
- many of the methods with the Firestore provider do not add the ID to
the document and therefore fail model validation. Additionally, some
places had incorrect conversion of documents into dict (using
`dict(doc)` instead of `doc.to_dict()`). This PR resolves both of these
issues.

---------

Co-authored-by: Anirudh Murali <[email protected]>

extension_service/datastore/providers/firestore.py

@@ -161,13 +161,17 @@ async def get_airport_by_id(self, id: int) -> Optional[models.Airport]:
         query = self.__client.collection("airports").where(
             filter=FieldFilter("id", "==", id)
         )
-        return models.Airport.model_validate(await query.get().to_dict())
+        airport_doc = await query.get()
+        airport_dict = airport_doc.to_dict() | {"id": airport_doc.id}
+        return models.Airport.model_validate(airport_dict)
 
     async def get_airport_by_iata(self, iata: str) -> Optional[models.Airport]:
         query = self.__client.collection("airports").where(
             filter=FieldFilter("iata", "==", iata)
         )
-        return models.Airport.model_validate(await query.get().to_dict())
+        airport_doc = await query.get()
+        airport_dict = airport_doc.to_dict() | {"id": airport_doc.id}
+        return models.Airport.model_validate(airport_dict)
 
     async def search_airports(
         self,
@@ -187,14 +191,19 @@ async def search_airports(
             query = query.where("name", ">=", name).where("name", "<=", name + "\uf8ff")
 
         docs = query.stream()
-        airports = [models.Airport.model_validate(dict(doc)) async for doc in docs]
+        airports = []
+        async for doc in docs:
+            airport_dict = doc.to_dict() | {"id": doc.id}
+            airports.append(models.Airport.model_validate(airport_dict))
         return airports
 
     async def get_amenity(self, id: int) -> Optional[models.Amenity]:
         query = self.__client.collection("amenities").where(
             filter=FieldFilter("id", "==", id)
         )
-        return models.Amenity.model_validate(await query.get().to_dict())
+        amenity_doc = await query.get()
+        amenity_dict = amenity_doc.to_dict() | {"id": amenity_doc.id}
+        return models.Amenity.model_validate(amenity_dict)
 
     async def amenities_search(
         self, query_embedding: list[float], similarity_threshold: float, top_k: int
@@ -203,9 +212,11 @@ async def amenities_search(
 
     async def get_flight(self, flight_id: int) -> Optional[models.Flight]:
         query = self.__client.collection("flights").where(
-            filter=FieldFilter("id", "==", id)
+            filter=FieldFilter("id", "==", flight_id)
         )
-        return models.Flight.model_validate(await query.get().to_dict())
+        flight_doc = await query.get()
+        flight_dict = flight_doc.to_dict() | {"id": flight_doc.id}
+        return models.Flight.model_validate(flight_dict)
 
     async def search_flights_by_number(
         self,
@@ -218,9 +229,11 @@ async def search_flights_by_number(
             .where(filter=FieldFilter("flight_number", "==", number))
         )
 
-        flights = [
-            models.Flight.model_validate(dict(doc)) async for doc in query.stream()
-        ]
+        docs = query.stream()
+        flights = []
+        async for doc in docs:
+            flight_dict = doc.to_dict() | {"id": doc.id}
+            flights.append(models.Flight.model_validate(flight_dict))
         return flights
 
     async def search_flights_by_airports(
@@ -243,7 +256,10 @@ async def search_flights_by_airports(
             query = query.where("arrival_airport", "==", arrival_airport)
 
         docs = query.stream()
-        flights = [models.Flight.model_validate(dict(doc)) async for doc in docs]
+        flights = []
+        async for doc in docs:
+            flight_dict = doc.to_dict() | {"id": doc.id}
+            flights.append(models.Flight.model_validate(flight_dict))
         return flights
 
     async def close(self):