Fix/n plus one routes for stops

internal/restapi/trips_for_location_handler.go

@@ -504,29 +504,47 @@ func (rb *referenceBuilder) collectTripIDs(trips []models.TripsForLocationListEn
 
 func (rb *referenceBuilder) buildStopList(stops []gtfsdb.Stop) {
 	rb.stopList = make([]models.Stop, 0, len(stops))
+	if len(stops) == 0 {
+		return
+	}
+
+	stopIDs := make([]string, 0, len(stops))
 	for _, stop := range stops {
-		if rb.ctx.Err() != nil {
-			return
-		}
+		stopIDs = append(stopIDs, stop.ID)
+	}
 
-		routeIds, err := rb.api.GtfsManager.GtfsDB.Queries.GetRouteIDsForStop(rb.ctx, stop.ID)
-		if err != nil {
+	routesForStops, err := rb.api.GtfsManager.GtfsDB.Queries.GetRouteIDsForStops(rb.ctx, stopIDs)
+	if err != nil {
+		logging.LogError(rb.api.Logger, "failed to batch fetch routes for stops", err)
+		return
+	}
+
+	// Build in-memory map: stopID → []combinedRouteID (e.g. "40_100479").
+	// Also register the raw route ID (e.g. "100479") in presentRoutes so that
+	// collectAgenciesAndRoutes can fetch full route details via GetRoutesByIDs,
+	// which queries WHERE routes.id IN (?), using raw IDs — not combined ones.
+	stopRouteMap := make(map[string][]string)
+	for _, r := range routesForStops {
+		combinedID, ok := r.RouteID.(string)
+		if !ok {
 			continue
 		}
-
-		routeIdsString := rb.processRouteIds(routeIds)
-		rb.stopList = append(rb.stopList, rb.createStop(stop, routeIdsString))
+		stopRouteMap[r.StopID] = append(stopRouteMap[r.StopID], combinedID)
+		if rawID, err := utils.ExtractCodeID(combinedID); err == nil {
+			rb.presentRoutes[rawID] = models.Route{}
+		}
 	}
-}
 
-func (rb *referenceBuilder) processRouteIds(routeIds []interface{}) []string {
-	routeIdsString := make([]string, len(routeIds))
-	for i, id := range routeIds {
-		routeId := id.(string)
-		rb.presentRoutes[routeId] = models.Route{}
-		routeIdsString[i] = routeId
+	for _, stop := range stops {
+		if rb.ctx.Err() != nil {
+			return
+		}
+		combinedRouteIDs := stopRouteMap[stop.ID]
+		if combinedRouteIDs == nil {
+			combinedRouteIDs = []string{}
+		}
+		rb.stopList = append(rb.stopList, rb.createStop(stop, combinedRouteIDs))
 	}
-	return routeIdsString
 }
 
 func (rb *referenceBuilder) createStop(stop gtfsdb.Stop, routeIds []string) models.Stop {

internal/restapi/trips_for_location_handler_test.go

@@ -2,11 +2,12 @@ package restapi
 
 import (
 	"fmt"
-	"github.com/stretchr/testify/require"
 	"net/http"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/stretchr/testify/assert"
 )
 
@@ -106,6 +107,85 @@ func TestTripsForLocationHandler_DifferentAreas(t *testing.T) {
 	}
 }
 
+func TestTripsForLocationHandler_ReferencesContainStopsAndRoutes(t *testing.T) {
+	api, cleanup := createTestApiWithRealTimeData(t)
+	defer cleanup()
+
+	time.Sleep(500 * time.Millisecond)
+
+	tests := []struct {
+		name    string
+		lat     float64
+		lon     float64
+		latSpan float64
+		lonSpan float64
+	}{
+		{
+			name:    "Transit Center Area",
+			lat:     40.5865,
+			lon:     -122.3917,
+			latSpan: 1.0,
+			lonSpan: 1.0,
+		},
+		{
+			name:    "Wide Area Coverage",
+			lat:     40.5865,
+			lon:     -122.3917,
+			latSpan: 2.0,
+			lonSpan: 3.0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			url := fmt.Sprintf("/api/where/trips-for-location.json?key=TEST&lat=%f&lon=%f&latSpan=%f&lonSpan=%f&includeSchedule=true",
+				tt.lat, tt.lon, tt.latSpan, tt.lonSpan)
+
+			resp, model := serveApiAndRetrieveEndpoint(t, api, url)
+			assert.Equal(t, http.StatusOK, resp.StatusCode)
+
+			data, ok := model.Data.(map[string]interface{})
+			require.True(t, ok, "response data should be a map")
+
+			refs, ok := data["references"].(map[string]interface{})
+			require.True(t, ok, "references should be present in response data")
+
+			stops, ok := refs["stops"].([]interface{})
+			require.True(t, ok, "references.stops should be a []interface{}, not null")
+
+			routes, ok := refs["routes"].([]interface{})
+			require.True(t, ok, "references.routes should be a []interface{}, not null")
+
+			if len(stops) > 0 {
+				routeIDSet := make(map[string]bool, len(routes))
+				for _, r := range routes {
+					route, ok := r.(map[string]interface{})
+					require.True(t, ok, "each route in references.routes should be a map")
+					if id, ok := route["id"].(string); ok {
+						routeIDSet[id] = true
+					}
+				}
+
+				for _, s := range stops {
+					stop, ok := s.(map[string]interface{})
+					require.True(t, ok, "each stop in references.stops should be a map")
+
+					routeIds, ok := stop["routeIds"].([]interface{})
+					assert.True(t, ok, "stop.routeIds should be a []interface{}, not null (stop: %v)", stop["id"])
+
+					for _, rid := range routeIds {
+						routeID, ok := rid.(string)
+						require.True(t, ok)
+						assert.True(t, routeIDSet[routeID],
+							"route %q referenced by stop %q must appear in references.routes",
+							routeID, stop["id"])
+					}
+				}
+			}
+		})
+	}
+}
+
 func TestTripsForLocationHandler_ScheduleInclusion(t *testing.T) {
 	api, cleanup := createTestApiWithRealTimeData(t)
 	defer cleanup()