Function testConnContextCanceledCancelsRunningQueryOnServer

pgconn/pgconn_test.go:2890–2946

(t *testing.T, connString, dbType string)

}

func testConnContextCanceledCancelsRunningQueryOnServer(t *testing.T, connString, dbType string) {
	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
	defer cancel()

	pgConn, err := pgconn.Connect(ctx, connString)
	require.NoError(t, err)
	defer closeConn(t, pgConn)

	ctx, cancel = context.WithTimeout(ctx, 100*time.Millisecond)
	defer cancel()

	// Getting the actual PostgreSQL server process ID (PID) from a query executed through pgbouncer is not straightforward
	// because pgbouncer abstracts the underlying database connections, and it doesn't expose the PID of the PostgreSQL
	// server process to clients. However, we can check if the query is running by checking the generated query ID.
	queryID := fmt.Sprintf("%s testConnContextCanceled %d", dbType, time.Now().UnixNano())

	multiResult := pgConn.Exec(ctx, fmt.Sprintf(`
	-- %v
	select 'Hello, world', pg_sleep(30)
	`, queryID))

	for multiResult.NextResult() {
	}
	err = multiResult.Close()
	assert.True(t, pgconn.Timeout(err))
	assert.True(t, pgConn.IsClosed())
	select {
	case <-pgConn.CleanupDone():
	case <-time.After(5 * time.Second):
		t.Fatal("Connection cleanup exceeded maximum time")
	}

	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	otherConn, err := pgconn.Connect(ctx, connString)
	require.NoError(t, err)
	defer closeConn(t, otherConn)

	ctx, cancel = context.WithTimeout(ctx, time.Second*5)
	defer cancel()

	for {
		result := otherConn.ExecParams(ctx,
			`select 1 from pg_stat_activity where query like $1`,
			[][]byte{[]byte("%" + queryID + "%")},
			nil,
			nil,
			nil,
		).Read()
		require.NoError(t, result.Err)

		if len(result.Rows) == 0 {
			break
		}
	}
}