WIP Convert get_shapes to use SQLAlchemy ORM, query builder

Author: raebot

_shared_utils/requirements.txt

@@ -13,6 +13,8 @@ pytest-recording (>=0.13.4,<0.14.0)
 pytest-unordered (>=0.7.0,<0.8.0)
 quarto==0.1.0
 quarto-cli==1.6.40
+sqlalchemy==1.4.46
+sqlalchemy-bigquery==1.11.0
 vegafusion==2.0.2
 vl-convert-python>=1.6.0
 movingpandas==0.22.4

_shared_utils/shared_utils/gtfs_utils_v2.py

@@ -17,6 +17,7 @@
 from shared_utils import DBSession, schedule_rt_utils
 from shared_utils.models.dim_gtfs_dataset import DimGtfsDataset
 from shared_utils.models.fct_daily_schedule_feeds import FctDailyScheduleFeeds
+from shared_utils.models.fct_daily_scheduled_shapes import FctDailyScheduledShapes
 from shared_utils.models.fct_scheduled_trips import FctScheduledTrips
 from siuba import *
 from sqlalchemy import and_, create_engine, func, or_, select
@@ -407,7 +408,7 @@ def get_shapes(
     get_df: bool = True,
     crs: str = geography_utils.WGS84,
     custom_filtering: dict = None,
-) -> gpd.GeoDataFrame:
+) -> Union[gpd.GeoDataFrame | sqlalchemy.sql.selectable.Select]:
     """
     Query fct_daily_scheduled_shapes.
 
@@ -416,21 +417,15 @@ def get_shapes(
     """
     check_operator_feeds(operator_feeds)
 
-    # If pt_array is not kept in the final, we still need it
-    # to turn this into a gdf
-    if "pt_array" not in shape_cols:
-        shape_cols_with_geom = shape_cols + ["pt_array"]
-    elif shape_cols:
-        shape_cols_with_geom = shape_cols[:]
+    search_conditions = [
+        FctDailyScheduledShapes.service_date == selected_date,
+        FctDailyScheduledShapes.feed_key.in_(operator_feeds),
+    ]
 
-    tables = _get_tables()
+    for k, v in (custom_filtering or {}).items():
+        search_conditions.append(getattr(FctDailyScheduledShapes, k).in_(v))
 
-    shapes = (
-        getattr(tables, "test_shared_utils").fct_daily_scheduled_shapes()
-        >> filter_date(selected_date, date_col="service_date")
-        >> filter_operator(operator_feeds, include_name=False)
-        >> filter_custom_col(custom_filtering)
-    )
+    shapes = select(FctDailyScheduledShapes).where(and_(*search_conditions))
 
     if get_df:
         shapes = shapes >> collect()
@@ -442,7 +437,12 @@ def get_shapes(
         return shapes_gdf
 
     else:
-        return shapes >> subset_cols(shape_cols_with_geom)
+        columns = {FctDailyScheduledShapes.pt_array}
+
+        for column in shape_cols:
+            columns.add(getattr(FctDailyScheduledShapes, column))
+
+        return shapes.with_only_columns(*list(columns))
 
 
 def get_stops(

_shared_utils/shared_utils/models/fct_daily_scheduled_shapes.py

@@ -1 +1,27 @@
-# mart_gtfs.fct_daily_scheduled_shapes
+from shared_utils.models.base import get_table_name
+from sqlalchemy import Boolean, Column, Date, DateTime, Integer, String
+from sqlalchemy.orm import declarative_base, declared_attr
+from sqlalchemy_bigquery import GEOGRAPHY
+
+Base = declarative_base()
+
+
+class FctDailyScheduledShapes(Base):
+    dataset = "mart_gtfs"
+    table = "fct_daily_scheduled_shapes"
+
+    @declared_attr
+    def __tablename__(cls):
+        return get_table_name(cls.dataset, cls.table)
+
+    key = Column(String, primary_key=True)
+    feed_key = Column(String)
+    service_date = Column(Date)
+    shape_id = Column(String)
+    shape_array_key = Column(String)
+    feed_timezone = Column(String)
+    n_trips = Column(Integer)
+    shape_first_departure_datetime_pacific = Column(DateTime)
+    shape_last_arrival_datetime_pacific = Column(DateTime)
+    contains_warning_duplicate_trip_primary_key = Column(Boolean)
+    pt_array = Column(GEOGRAPHY)

_shared_utils/tests/shared_utils/test_gtfs_utils_v2.py

@@ -731,3 +731,11 @@ def test_get_shapes_custom_filtering(self):
     def test_get_shapes_no_operator_feeds(self):
         with pytest.raises(ValueError, match="Supply list of feed keys or operator names!"):
             get_shapes(selected_date="2025-09-19")
+
+    def test_get_shapes_get_df_false(self):
+        result = get_shapes(
+            selected_date="2025-09-01", operator_feeds=["3ea60aa240ddc543da5415ccc759fd6d"], get_df=False
+        )
+
+        assert isinstance(result, sqlalchemy.sql.selectable.Select)
+        # assert result select includes pt_array