Make TransactionBuilder's constructor public

[mihai-dinculescu]

TransactionBuilder is already generic over AsyncConnection<Backend = Pg, TransactionManager = AnsiTransactionManager>, and I would like to take advantage of this to use it with a connnection other than AsyncPgConnection.

In my testing, it all should work just fine if

diesel_async/src/pg/transaction_builder.rs

Line 30 in 1c36b65

pub(crate) fn new(connection: &'a mut C) -> Self {

is made public.

Are there any reasons why the constructor shouldn't be made public?
If not, I'm happy to submit a PR to make it so.

[weiznich]

I converted this to a discussion as this is more a feature request than a bug report.
Could you explain your use case a bit more in detail to help us understanding how common that use case might be?

[mihai-dinculescu]

My use case is that I've got my own struct implementing AsyncConnection and SimpleAsyncConnection that wraps up an AsyncPgConnection so that I can hook tracing up.

That's all working fine for the most common uses, except the ones that need more fine-grained control over the transaction options, like the isolation level.

This is where I need a TransactionBuilder<MyAsyncPgConnectionWrapper>, but because build_transaction returns TransactionBuilder<AsyncPgConnection>, I cannot use it with my wrapper.

TransactionBuilder::new is able to create a TransactionBuilder<MyAsyncPgConnectionWrapper>, if only it would be public, hence the ask :)

[weiznich]

Thanks for providing these details. Would you mind explaining why you plan to write such a wrapper in the first place instead of just using the existing instrumentation infrastructure: https://docs.rs/diesel-async/latest/diesel_async/trait.AsyncConnection.html#tymethod.set_instrumentation. The later should allow you to easily implement that via the provided hooks without needing to have a custom connection implementation at all.

[mihai-dinculescu]

The wrapper is built to handle TLS configuration (if necessary) and call the instrumentation hooks with a custom span provider, ensuring that connections retrieved from the pool are pre-configured for the runtime environment and instrumented according to my organization's conventions.

Most of this functionality is already facilitated by diesel_async through its clever use of generics—for instance, allowing AsyncDieselConnectionManager to seamlessly integrate with our custom wrapper. Given this, it would be reasonable for TransactionBuilder to adopt a similar pattern for consistency, especially since it is already sufficiently generic to support a custom wrapper.

[weiznich]

Thanks for providing these details. As already written above: The instrumentation part of that is certainly provided by existing hook mechanisms. If you still feel that this doesn't sufficiently cover your use-case I really would like to see an example for that.
So you can certainly implement a new connection type for that, but you shouldn't expect that this is then well supported. As for the TLS configuration: I can understand that this requires a bit of setup code, but I'm not sure if that on it's own warrants a custom connection implementation.

Generally we try to keep such details like this constructor as private function as we do not want guarantee stability for that.

[mihai-dinculescu]

Here’s a summary of the wrapper we’ve developed. The primary focus was on achieving consistent instrumentation, which is particularly valuable for my team given its usage across several dozen repositories and services.

pub type Pool = deadpool::Pool<MyAsyncPgConnection>;
pub type PoolBuilder = deadpool::PoolBuilder<MyAsyncPgConnection>;

pub struct MyAsyncPgConnection {
    conn: AsyncPgConnection,
    span_provider: PgQuerySpan,
}

impl MyAsyncPgConnection {
    /// Build a transaction, specifying additional details such as isolation level
    pub fn build_transaction(&mut self) -> TransactionBuilder<Self> {
        // This is what I would like to be able to do
        TransactionBuilder::new(self)
    }
}

impl AsyncConnection for MyAsyncPgConnection {
    async fn establish(database_url: &str) -> ConnectionResult<Self> {
        async {
            // create the connection with or without TLS depending on env

            // let mut conn = ...

            let namespace = diesel::select(current_database())
                .get_result(&mut conn)
                .await
                .map_err(ConnectionError::CouldntSetupConfiguration)?;

            Ok(PgConnection {
                conn,
                span_provider: PgQuerySpan::new(namespace),
            })
        }
        .inspect_err(|e| {
            // record error with the necessary semantic fields
            // https://github.com/open-telemetry/semantic-conventions/blob/main/docs/exceptions/exceptions-spans.md
        })
        .instrument(PgEstablishSpan.span())
        .await
    }

    fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query>
    where
        T: AsQuery + 'query,
        T::Query: QueryFragment<Self::Backend> + QueryId + 'query,
    {
        self.conn.load(source).instrument(self.span_provider.span())
    }

    // ...
}

pub struct PgQuerySpan {
    namespace: String,
}

impl PgQuerySpan {
    pub fn new(namespace: String) -> Self {
        Self { namespace }
    }
}

pub trait SpanProvider: Sized {
    fn span(&self) -> Span;
}

impl<F> SpanProvider for F
where
    F: Fn() -> Span,
{
    fn span(&self) -> Span {
        (self)()
    }
}

impl SpanProvider for PgQuerySpan {
    fn span(&self) -> Span {
        info_span!(
            "query",
            // set up a handful of standard fields
        )
    }
}

pub struct PgEstablishSpan;

impl SpanProvider for PgEstablishSpan {
    fn span(&self) -> Span {
        info_span!(
            "establish_connection",
            // set up a handful of standard fields
        )
    }
}

/// Extension trait for deadpool [`Pool`]s that handles the required setup for a pool of
/// [`MyAsyncPgConnection`]s.
pub trait PoolExt {
    /// Returns a [`PoolBuilder`] after first handling the details of setting up the
    /// `ConnectionManager` and installing connection instrumentation.
    fn my_builder(db_url: &str) -> PoolBuilder;
}

impl PoolExt for Pool {
    fn my_builder(db_url: &str) -> PoolBuilder {
        connection::set_default_instrumentation(|| Some(Box::new(ConnectionEvents)))
            .expect("the RwLock backing the global instrumentation has not been poisoned");

        let mut config = ManagerConfig::default();
        config.recycling_method = RecyclingMethod::Fast;

        let manager =
            AsyncDieselConnectionManager::<MyAsyncPgConnection>::new_with_config(db_url, config);

        Pool::builder(manager)
    }
}

pub struct ConnectionEvents;

impl Instrumentation for ConnectionEvents {
    fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
        match event {
            // We're not inside a span at this point
            InstrumentationEvent::StartQuery { .. } => {}
            // At this point in time we are inside the span, we have both the query and might have
            // an error, this is the interesting event.
            InstrumentationEvent::FinishQuery { query, error, .. } => {
                let span = Span::current();

                // Check to make sure we're actually in a DB span, just in case.
                if span.has_field(DB_QUERY_TEXT) {
                    // ...

                    span.record(SPAN_NAME, op_name(query));
                    span.record(DB_QUERY_TEXT, query);

                    if let Some(err) = error {
                        // record error
                    }
                }
            }
            _ => {}
        }
    }
}