feat: add test for database export

retrieval_service/alloydb.tests.cloudbuild.yaml

@@ -76,6 +76,30 @@ steps:
         export DB_PASS=$$PGPASSWORD
         python -m pytest
 
+  - id: Run database export
+    name: python:3.11
+    dir: retrieval_service
+    entrypoint: /bin/bash
+    args:
+      - "-c"
+      - |
+        # Run script
+        python run_database_export.py
+        cd ../data
+        diff --strip-trailing-cr -Z airport_dataset.csv airport_dataset.csv.new || (echo "airport dataset export fail" && exit 1)
+        diff --strip-trailing-cr -Z amenity_dataset.csv amenity_dataset.csv.new || (echo "amenity dataset export fail" && exit 1)
+        diff --strip-trailing-cr -Z flights_dataset.csv flights_dataset.csv.new || (echo "flight dataset export fail" && exit 1)
+
+  - id: Clean exported files
+    name: python:3.11
+    dir: data
+    entrypoint: /bin/bash
+    args:
+      - "-c"
+      - |
+        rm airport_dataset.csv.new amenity_dataset.csv.new flights_dataset.csv.new
+
+
   - id: Clean database
     name: postgres
     entrypoint: /bin/bash
@@ -99,6 +123,7 @@ availableSecrets:
     - versionName: projects/$PROJECT_ID/secrets/alloy_db_user/versions/latest
       env: PGUSER
 options:
+  substitutionOption: 'ALLOW_LOOSE'
   dynamic_substitutions: true
   pool:
     name: projects/$PROJECT_ID/locations/us-central1/workerPools/alloy-private-pool # Necessary for VPC network connection

retrieval_service/datastore/providers/postgres.py

@@ -170,13 +170,13 @@ async def export_data(
         self,
     ) -> tuple[list[models.Airport], list[models.Amenity], list[models.Flight]]:
         airport_task = asyncio.create_task(
-            self.__pool.fetch("""SELECT * FROM airports""")
+            self.__pool.fetch("""SELECT * FROM airports ORDER BY id ASC""")
         )
         amenity_task = asyncio.create_task(
-            self.__pool.fetch("""SELECT * FROM amenities""")
+            self.__pool.fetch("""SELECT * FROM amenities ORDER BY id ASC""")
         )
         flight_task = asyncio.create_task(
-            self.__pool.fetch("""SELECT * FROM flights""")
+            self.__pool.fetch("""SELECT * FROM flights ORDER BY id ASC""")
         )
 
         airports = [models.Airport.model_validate(dict(a)) for a in await airport_task]