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.