break out long lines

GoogleCloudPlatform · Jan 8, 2024 · c578717 · c578717
1 parent 58e5f2c
commit c578717
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 17 deletions.
diff --git a/retrieval_service/datastore/providers/cloudsql_postgres.py b/retrieval_service/datastore/providers/cloudsql_postgres.py
@@ -154,7 +154,14 @@ async def initialize_data(
             # Insert all the data
             await conn.execute(
                 text(
-                    """INSERT INTO amenities VALUES (:id, :name, :description, :location, :terminal, :category, :hour, :sunday_start_hour, :sunday_end_hour, :monday_start_hour, :monday_end_hour, :tuesday_start_hour, :tuesday_end_hour, :wednesday_start_hour, :wednesday_end_hour, :thursday_start_hour, :thursday_end_hour, :friday_start_hour, :friday_end_hour, :saturday_start_hour, :saturday_end_hour, :content, :embedding)"""
+                    """
+                    INSERT INTO amenities VALUES (:id, :name, :description, :location,
+                      :terminal, :category, :hour, :sunday_start_hour, :sunday_end_hour,
+                      :monday_start_hour, :monday_end_hour, :tuesday_start_hour,
+                      :tuesday_end_hour, :wednesday_start_hour, :wednesday_end_hour,
+                      :thursday_start_hour, :thursday_end_hour, :friday_start_hour,
+                      :friday_end_hour, :saturday_start_hour, :saturday_end_hour, :content, :embedding)
+                    """
                 ),
                 [
                     {
@@ -209,7 +216,11 @@ async def initialize_data(
             # Insert all the data
             await conn.execute(
                 text(
-                    """INSERT INTO flights VALUES (:id, :airline, :flight_number, :departure_airport, :arrival_airport, :departure_time, :arrival_time, :departure_gate, :arrival_gate)"""
+                    """
+                    INSERT INTO flights VALUES (:id, :airline, :flight_number,
+                      :departure_airport, :arrival_airport, :departure_time,
+                      :arrival_time, :departure_gate, :arrival_gate)
+                    """
                 ),
                 [
                     {
@@ -305,8 +316,12 @@ async def get_amenity(self, id: int) -> Optional[models.Amenity]:
         async with self.__pool.connect() as conn:
             s = text(
                 """
-                SELECT id, name, description, location, terminal, category, hour, sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour, tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour, thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour, saturday_start_hour, saturday_end_hour
-                  FROM amenities WHERE id=:id
+                SELECT id, name, description, location, terminal, category, hour,
+                  sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour,
+                  tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
+                  thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour,
+                  saturday_start_hour, saturday_end_hour
+                FROM amenities WHERE id=:id
                 """
             )
             params = {"id": id}
@@ -324,9 +339,17 @@ async def amenities_search(
         async with self.__pool.connect() as conn:
             s = text(
                 """
-                SELECT id, name, description, location, terminal, category, hour, sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour, tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour, thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour, saturday_start_hour, saturday_end_hour
+                SELECT id, name, description, location, terminal, category, hour,
+                  sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour,
+                  tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
+                  thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour,
+                  saturday_start_hour, saturday_end_hour
                   FROM (
-                      SELECT id, name, description, location, terminal, category, hour, sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour, tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour, thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour, saturday_start_hour, saturday_end_hour, 1 - (embedding <=> :query_embedding) AS similarity
+                      SELECT id, name, description, location, terminal, category, hour,
+                        sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour,
+                        tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
+                        thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour,
+                        saturday_start_hour, saturday_end_hour, 1 - (embedding <=> :query_embedding) AS similarity
                       FROM amenities
                       WHERE 1 - (embedding <=> :query_embedding) > :similarity_threshold
                       ORDER BY similarity DESC

diff --git a/retrieval_service/datastore/providers/postgres.py b/retrieval_service/datastore/providers/postgres.py
@@ -126,7 +126,14 @@ async def initialize_data(
             )
             # Insert all the data
             await conn.executemany(
-                """INSERT INTO amenities VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23)""",
+                """
+                INSERT INTO amenities VALUES (
+                  $1, $2, $3, $4, $5,
+                  $6, $7, $8, $9, $10,
+                  $11, $12, $13, $14, $15,
+                  $16, $17, $18, $19, $20,
+                  $21, $22, $23)
+                """,
                 [
                     (
                         a.id,
@@ -265,8 +272,13 @@ async def search_airports(
     async def get_amenity(self, id: int) -> Optional[models.Amenity]:
         result = await self.__pool.fetchrow(
             """
-                SELECT id, name, description, location, terminal, category, hour, sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour, tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour, thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour, saturday_start_hour, saturday_end_hour
-                FROM amenities WHERE id=$1
+            SELECT id, name, description, location, terminal,
+              category, hour, sunday_start_hour, sunday_end_hour,
+              monday_start_hour, monday_end_hour, tuesday_start_hour,
+              tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
+              thursday_start_hour, thursday_end_hour, friday_start_hour,
+              friday_end_hour, saturday_start_hour, saturday_end_hour
+            FROM amenities WHERE id=$1
             """,
             id,
         )
@@ -282,14 +294,25 @@ async def amenities_search(
     ) -> list[models.Amenity]:
         results = await self.__pool.fetch(
             """
-                SELECT id, name, description, location, terminal, category, hour, sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour, tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour, thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour, saturday_start_hour, saturday_end_hour
-                FROM (
-                    SELECT id, name, description, location, terminal, category, hour, sunday_start_hour, sunday_end_hour, monday_start_hour, monday_end_hour, tuesday_start_hour, tuesday_end_hour, wednesday_start_hour, wednesday_end_hour, thursday_start_hour, thursday_end_hour, friday_start_hour, friday_end_hour, saturday_start_hour, saturday_end_hour, 1 - (embedding <=> $1) AS similarity
-                    FROM amenities
-                    WHERE 1 - (embedding <=> $1) > $2
-                    ORDER BY similarity DESC
-                    LIMIT $3
-                ) AS sorted_amenities
+            SELECT id, name, description, location, terminal,
+              category, hour, sunday_start_hour, sunday_end_hour,
+              monday_start_hour, monday_end_hour, tuesday_start_hour,
+              tuesday_end_hour, wednesday_start_hour, wednesday_end_hour,
+              thursday_start_hour, thursday_end_hour, friday_start_hour,
+              friday_end_hour, saturday_start_hour, saturday_end_hour
+            FROM (
+                SELECT id, name, description, location, terminal, category,
+                  hour, sunday_start_hour, sunday_end_hour, monday_start_hour,
+                  monday_end_hour, tuesday_start_hour, tuesday_end_hour,
+                  wednesday_start_hour, wednesday_end_hour, thursday_start_hour,
+                  thursday_end_hour, friday_start_hour, friday_end_hour,
+                  saturday_start_hour, saturday_end_hour,
+                  1 - (embedding <=> $1) AS similarity
+                FROM amenities
+                WHERE 1 - (embedding <=> $1) > $2
+                ORDER BY similarity DESC
+                LIMIT $3
+            ) AS sorted_amenities
             """,
             query_embedding,
             similarity_threshold,