fix: propagate context for cancellable imports

.gitignore

@@ -21,3 +21,4 @@ testdata/perf/
 loadtest/results/
 loadtest/k6/*-summary.json
 loadtest/profiles/
+*.cov

gtfsdb/bulk_insert_frequency_test.go

@@ -23,7 +23,7 @@ func createFrequencyTestClient(t *testing.T) *Client {
 	client, err := NewClient(config)
 	require.NoError(t, err)
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	_, err = client.Queries.CreateAgency(ctx, CreateAgencyParams{
 		ID:       "test_agency",
@@ -94,7 +94,7 @@ func TestBulkInsertFrequencies(t *testing.T) {
 			client := createFrequencyTestClient(t)
 			defer func() { _ = client.Close() }()
 
-			ctx := context.Background()
+			ctx := t.Context()
 
 			var frequencies []CreateFrequencyParams
 			if tc.name == "Multiple windows for one trip" {
@@ -172,7 +172,7 @@ func TestBulkInsertFrequencies_MultipleTrips(t *testing.T) {
 	client := createFrequencyTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	frequencies := []CreateFrequencyParams{
 		{TripID: "trip_1", StartTime: int64(6 * 3600 * 1e9), EndTime: int64(9 * 3600 * 1e9), HeadwaySecs: 600, ExactTimes: 0},
@@ -207,7 +207,7 @@ func TestBulkInsertFrequencies_DuplicatePrimaryKey(t *testing.T) {
 	client := createFrequencyTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	startTime := int64(6 * 3600 * 1e9)
 	frequencies := []CreateFrequencyParams{
@@ -224,11 +224,34 @@ func TestBulkInsertFrequencies_DuplicatePrimaryKey(t *testing.T) {
 	assert.Equal(t, int64(600), rows[0].HeadwaySecs)
 	assert.Equal(t, int64(0), rows[0].ExactTimes)
 }
+
+func TestGetFrequencyTripIDs(t *testing.T) {
+	client := createFrequencyTestClient(t)
+	defer func() { _ = client.Close() }()
+
+	ctx := t.Context()
+
+	frequencies := []CreateFrequencyParams{
+		{TripID: "trip_1", StartTime: int64(6 * 3600 * 1e9), EndTime: int64(9 * 3600 * 1e9), HeadwaySecs: 600, ExactTimes: 0},
+		{TripID: "trip_1", StartTime: int64(16 * 3600 * 1e9), EndTime: int64(19 * 3600 * 1e9), HeadwaySecs: 600, ExactTimes: 0},
+		{TripID: "trip_3", StartTime: int64(8 * 3600 * 1e9), EndTime: int64(12 * 3600 * 1e9), HeadwaySecs: 1200, ExactTimes: 0},
+	}
+
+	err := client.bulkInsertFrequencies(ctx, frequencies, nil)
+	require.NoError(t, err)
+
+	tripIDs, err := client.Queries.GetFrequencyTripIDs(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, 2, len(tripIDs), "Should return 2 distinct trip IDs")
+	assert.Contains(t, tripIDs, "trip_1")
+	assert.Contains(t, tripIDs, "trip_3")
+}
+
 func TestClearFrequencies(t *testing.T) {
 	client := createFrequencyTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	frequencies := []CreateFrequencyParams{
 		{TripID: "trip_1", StartTime: int64(6 * 3600 * 1e9), EndTime: int64(9 * 3600 * 1e9), HeadwaySecs: 600, ExactTimes: 0},
@@ -258,7 +281,7 @@ func TestGetFrequenciesForTrips(t *testing.T) {
 	client := createFrequencyTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	frequencies := []CreateFrequencyParams{
 		{TripID: "trip_1", StartTime: int64(6 * 3600 * 1e9), EndTime: int64(9 * 3600 * 1e9), HeadwaySecs: 600, ExactTimes: 0},

gtfsdb/bulk_insert_test.go

@@ -21,7 +21,7 @@ func TestBulkInsertStopTimes(t *testing.T) {
 	require.NoError(t, err)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Create prerequisite data (agency, route, calendar/service)
 	_, err = client.Queries.CreateAgency(ctx, CreateAgencyParams{
@@ -143,7 +143,7 @@ func TestBulkInsertShapes(t *testing.T) {
 	require.NoError(t, err)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	testCases := []struct {
 		name  string
@@ -214,7 +214,7 @@ func TestBulkInsertWithNullValues(t *testing.T) {
 	require.NoError(t, err)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Test shapes with NULL distance traveled
 	shapes := []CreateShapeParams{
@@ -265,7 +265,7 @@ func TestBulkInsertPerformance(t *testing.T) {
 	require.NoError(t, err)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Create prerequisite data (agency, route, calendar/service)
 	_, err = client.Queries.CreateAgency(ctx, CreateAgencyParams{

gtfsdb/conditional_import_test.go

@@ -52,7 +52,7 @@ func TestConditionalImport_InitialImport(t *testing.T) {
 	require.NoError(t, err, "Failed to create client")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 	originalData, _ := createTestData(t)
 
 	// Perform initial import
@@ -85,7 +85,7 @@ func TestConditionalImport_SkipUnchangedData(t *testing.T) {
 	require.NoError(t, err, "Failed to create client")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 	originalData, _ := createTestData(t)
 
 	// Perform initial import
@@ -139,7 +139,7 @@ func TestConditionalImport_ReloadChangedData(t *testing.T) {
 	require.NoError(t, err, "Failed to create client")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 	originalData, modifiedData := createTestData(t)
 
 	// Perform initial import
@@ -182,7 +182,7 @@ func TestConditionalImport_DifferentSources(t *testing.T) {
 	require.NoError(t, err, "Failed to create client")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 	originalData, _ := createTestData(t)
 
 	// Perform initial import with source A
@@ -218,7 +218,7 @@ func TestConditionalImport_FileImport(t *testing.T) {
 	require.NoError(t, err, "Failed to create client")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 	testFilePath := getTestFixturePath(t, "raba.zip")
 
 	// Perform initial import from file
@@ -260,7 +260,7 @@ func TestClearAllGTFSData(t *testing.T) {
 	require.NoError(t, err, "Failed to create client")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 	originalData, _ := createTestData(t)
 
 	// Perform initial import

gtfsdb/connection_pool_test.go

@@ -51,7 +51,7 @@ func TestConnectionPoolBehavior(t *testing.T) {
 	db := client.DB
 
 	// Test that we can make sequential queries since :memory: uses 1 connection
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Make sequential queries to test that the single connection works
 	for i := 0; i < 10; i++ {
@@ -87,7 +87,7 @@ func TestConnectionLifetime(t *testing.T) {
 	initialStats := db.Stats()
 
 	// Make a query to create at least one connection
-	ctx := context.Background()
+	ctx := t.Context()
 	row := db.QueryRowContext(ctx, "SELECT 1")
 	var result int
 	err = row.Scan(&result)
@@ -115,7 +115,7 @@ func TestConnectionPoolConfiguration(t *testing.T) {
 	assert.Equal(t, 1, stats.MaxOpenConnections, "MaxOpenConns should be 1 for :memory: databases")
 
 	// Test that we can ping the database
-	ctx := context.Background()
+	ctx := t.Context()
 	err = db.PingContext(ctx)
 	assert.NoError(t, err, "Should be able to ping configured database")
 }

gtfsdb/error_handling_test.go

@@ -1,7 +1,6 @@
 package gtfsdb
 
 import (
-	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -98,7 +97,7 @@ func TestImportFromFile_ErrorHandling(t *testing.T) {
 	require.NoError(t, err, "NewClient should succeed")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Test with non-existent file
 	err = client.ImportFromFile(ctx, "/nonexistent/file.zip")
@@ -117,7 +116,7 @@ func TestDownloadAndStore_ErrorHandling(t *testing.T) {
 	require.NoError(t, err, "NewClient should succeed")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Test with invalid URL
 	err = client.DownloadAndStore(ctx, "invalid-url", "", "")
@@ -160,7 +159,7 @@ func TestDownloadAndStore_AuthenticationHeaders(t *testing.T) {
 	require.NoError(t, err, "NewClient should succeed")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Test with authentication headers
 	_ = client.DownloadAndStore(ctx, server.URL, expectedHeaderName, expectedHeaderValue)
@@ -197,7 +196,7 @@ func TestDownloadAndStore_NoAuthHeaders(t *testing.T) {
 	require.NoError(t, err, "NewClient should succeed")
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Test without authentication headers (empty strings)
 	_ = client.DownloadAndStore(ctx, server.URL, "", "")

gtfsdb/fts_queries_test.go

@@ -19,7 +19,7 @@ func createFTSTestClient(t *testing.T) *Client {
 	client, err := NewClient(config)
 	require.NoError(t, err)
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	_, err = client.Queries.CreateAgency(ctx, CreateAgencyParams{
 		ID:       "agency1",
@@ -36,7 +36,7 @@ func TestSearchRoutesByFullText(t *testing.T) {
 	client := createFTSTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Insert test routes with varied fields to verify correct column scanning
 	routes := []CreateRouteParams{
@@ -192,7 +192,7 @@ func TestSearchRoutesByFullTextEmptyDB(t *testing.T) {
 	client := createFTSTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	results, err := client.Queries.SearchRoutesByFullText(ctx, SearchRoutesByFullTextParams{
 		Query: "anything",
@@ -206,7 +206,7 @@ func TestSearchStopsByName(t *testing.T) {
 	client := createFTSTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Insert test stops with varied fields to verify correct column scanning
 	stops := []CreateStopParams{
@@ -323,7 +323,7 @@ func TestSearchStopsByNameEmptyDB(t *testing.T) {
 	client := createFTSTestClient(t)
 	defer func() { _ = client.Close() }()
 
-	ctx := context.Background()
+	ctx := t.Context()
 
 	results, err := client.Queries.SearchStopsByName(ctx, SearchStopsByNameParams{
 		SearchQuery: "anything",