AfterRelease hook does not provide context access

[gcfabri]

Hello.

I'm using pgxpool and specifically two hooks (BeforeAcquire and AfterRelease) to set a session variable and control tenant isolation.

In the BeforeAcquire hook the variable is set and then in AfterRelease it reverts the setting to an empty value.

My problem is from an observability point of view. To get the spans linked to the API request I am propagating the context and then it is shown in DataDog in the same trace. However, the AfterRelease hook does not provide access to the context and this function call is consequently not shown in the same trace as the others.

// BeforeAcquire is called before a connection is acquired from the pool. It must return true to allow the
// acquisition or false to indicate that the connection should be destroyed and a different connection should be
// acquired.
BeforeAcquire func(context.Context, *pgx.Conn) bool

// AfterRelease is called after a connection is released, but before it is returned to the pool. It must return true to
// return the connection to the pool or false to destroy the connection.
AfterRelease func(*pgx.Conn) bool

Is there any alternative/suggestion to achieve this?

[jackc]

AfterRelease doesn't take a context because it is called asynchronously after Release which is itself cannot block. i.e. It's not clear what cancelling a Release would even mean. I hadn't considered needing to pass arbitrary extra data.

I'm not sure of a good way to do this... but I can think of a way...

You could have a global map[*pgx.Conn]context.Context. You would need to guard access with a mutex, but you could set the value in BeforeAcquire and retrieve and delete it in AfterRelease.

[Jtcruthers]

Since v5, the main db methods require context to actually send a statement to pg. My use case is in AfterRelease call conn.Exec("RESET ALL") so that I can clear my session tenant_id with RLS. Without having the context passed in, I can't call conn.Exec in the callback.

Is there something else I should be doing? I don't need any of the data on the context, I just need a context to pass to the conn.Exec function. I can just pass a random context.Background() in via a closure, but that feels wrong.

[jackc]

As mentioned above, cleanup actions such as RESET ALL shouldn't be cancelled by original caller's context. What you probably want is to create a context with a 1-5 second timeout and use that.