diff --git a/gtfsdb/block_trip_entry_index_test.go b/gtfsdb/block_trip_entry_index_test.go new file mode 100644 index 000000000..c3e5a7e0b --- /dev/null +++ b/gtfsdb/block_trip_entry_index_test.go @@ -0,0 +1,42 @@ +package gtfsdb + +import ( + "database/sql" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// GetBlocksForBlockTripIndexIDs joins block_trip_entry to itself on block_id. Without +// an index there, SQLite scans every block entry for every route trip, which made +// trips-for-route take seconds on NJ Transit's 33k-trip feed. +func TestBlockTripEntryBlockIDJoinUsesIndex(t *testing.T) { + client := newTestClientWithRABA(t) + ctx := t.Context() + + var routeID, serviceID string + var indexID int64 + require.NoError(t, client.DB.QueryRowContext(ctx, ` + SELECT trips.route_id, bte.service_id, bte.block_trip_index_id + FROM block_trip_entry bte JOIN trips ON trips.id = bte.trip_id + WHERE bte.block_id IS NOT NULL LIMIT 1`).Scan(&routeID, &serviceID, &indexID)) + + recorder := &queryRecorder{DBTX: client.DB} + _, err := New(recorder).GetBlocksForBlockTripIndexIDs(ctx, GetBlocksForBlockTripIndexIDsParams{ + FromTime: sql.NullInt64{Int64: 0, Valid: true}, + ToTime: sql.NullInt64{Int64: int64(24 * time.Hour), Valid: true}, + RouteID: routeID, + IndexIds: []int64{indexID}, + ActiveServiceIds: []string{serviceID}, + RouteServiceIds: []string{serviceID}, + }) + require.NoError(t, err) + + // The planner may start the self-join from either side, so accept the index + // on either alias. + plan := strings.Join(explainRecordedQuery(t, client, recorder), "\n") + assert.Regexp(t, `SEARCH \w+ USING (COVERING )?INDEX idx_block_trip_entry_block_id \(block_id=\?`, plan) +} diff --git a/gtfsdb/query_plan_test.go b/gtfsdb/query_plan_test.go new file mode 100644 index 000000000..9e25181ef --- /dev/null +++ b/gtfsdb/query_plan_test.go @@ -0,0 +1,41 @@ +package gtfsdb + +import ( + "context" + "database/sql" + "testing" + + "github.com/stretchr/testify/require" +) + +// queryRecorder captures the SQL and arguments a generated query method sends, +// so a test can explain the real query rather than a copy of it. +type queryRecorder struct { + DBTX + query string + args []any +} + +func (db *queryRecorder) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { + db.query, db.args = query, args + return db.DBTX.QueryContext(ctx, query, args...) +} + +// explainRecordedQuery returns the detail column of each EXPLAIN QUERY PLAN row +// for the query the recorder last captured. +func explainRecordedQuery(t *testing.T, client *Client, recorder *queryRecorder) []string { + t.Helper() + rows, err := client.DB.QueryContext(t.Context(), "EXPLAIN QUERY PLAN "+recorder.query, recorder.args...) + require.NoError(t, err) + defer func() { _ = rows.Close() }() + + var details []string + for rows.Next() { + var id, parent, unused int + var detail string + require.NoError(t, rows.Scan(&id, &parent, &unused, &detail)) + details = append(details, detail) + } + require.NoError(t, rows.Err()) + return details +} diff --git a/gtfsdb/schema.sql b/gtfsdb/schema.sql index d55ca2bed..39ea32156 100644 --- a/gtfsdb/schema.sql +++ b/gtfsdb/schema.sql @@ -454,6 +454,11 @@ CREATE INDEX IF NOT EXISTS idx_block_trip_index_service_ids ON block_trip_index -- migrate CREATE INDEX IF NOT EXISTS idx_block_trip_entry_index_id ON block_trip_entry (block_trip_index_id); +-- migrate +-- GetBlocksForBlockTripIndexIDs (a route's active blocks, for trips-for-route) +-- and GetBlockTripIndexIDsForBlocks both match block_trip_entry rows by block_id. +CREATE INDEX IF NOT EXISTS idx_block_trip_entry_block_id ON block_trip_entry (block_id); + -- migrate CREATE INDEX IF NOT EXISTS idx_block_trip_entry_trip_id ON block_trip_entry (trip_id); diff --git a/gtfsdb/stops_rtree_test.go b/gtfsdb/stops_rtree_test.go index a82fdc64a..6bd190e89 100644 --- a/gtfsdb/stops_rtree_test.go +++ b/gtfsdb/stops_rtree_test.go @@ -2,7 +2,6 @@ package gtfsdb import ( "context" - "database/sql" "strings" "testing" @@ -47,44 +46,26 @@ func TestGetActiveRoutesWithinBounds_CenteredOnStop(t *testing.T) { assert.ElementsMatch(t, []string{"15", "154", "159", "160", "44X"}, stop4000RouteIDs) } -type routeQueryRecorder struct { - DBTX - query string - args []any -} - -func (db *routeQueryRecorder) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { - db.query, db.args = query, args - return db.DBTX.QueryContext(ctx, query, args...) -} - func TestGetActiveRoutesWithinBounds_FilteredStopTimesLookup(t *testing.T) { client := newTestClientWithRABA(t) ctx := t.Context() var routeID string require.NoError(t, client.DB.QueryRowContext(ctx, "SELECT id FROM routes LIMIT 1").Scan(&routeID)) - recorder := &routeQueryRecorder{DBTX: client.DB} + recorder := &queryRecorder{DBTX: client.DB} _, err := New(recorder).GetActiveRoutesWithinBounds(ctx, GetActiveRoutesWithinBoundsParams{ Lat: 40.7347, Lon: -74.1644, MinLat: -90, MaxLat: 90, MinLon: -180, MaxLon: 180, RouteIDs: []string{routeID}, MaxCount: 100, }) require.NoError(t, err) - plan, err := client.DB.QueryContext(ctx, "EXPLAIN QUERY PLAN "+recorder.query, recorder.args...) - require.NoError(t, err) - defer func() { _ = plan.Close() }() var stopTimeLookups []string - for plan.Next() { - var id, parent, unused int - var detail string - require.NoError(t, plan.Scan(&id, &parent, &unused, &detail)) + for _, detail := range explainRecordedQuery(t, client, recorder) { if strings.Contains(detail, "stop_times") { stopTimeLookups = append(stopTimeLookups, detail) } } - require.NoError(t, plan.Err()) require.NotEmpty(t, stopTimeLookups) for _, lookup := range stopTimeLookups { assert.Contains(t, lookup, "(trip_id=?)", "filtered searches must look up stop times by matching trip, not nearby stop")