dogdata-sqlx 0.0.4

SQLx instrumentation for tracing to Datadog APM via OpenTelemetry
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use sqlx::{Database, Executor};

use crate::execute::{InstrumentedExecute, InstrumentedFetch, InstrumentedPool};

pub struct InstrumentedQueryBuilder<Q> {
    pub query: Q,
    pub sql: &'static str,
}

impl<'q, DB, A> InstrumentedQueryBuilder<sqlx::query::Query<'q, DB, A>>
where
    DB: Database,
    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
{
    pub async fn execute<P>(self, pool: &P) -> Result<DB::QueryResult, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.execute_instrumented(pool, self.sql).await
    }

    pub async fn fetch_one<P>(self, pool: &P) -> Result<DB::Row, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_one_instrumented(pool, self.sql).await
    }

    pub async fn fetch_optional<P>(self, pool: &P) -> Result<Option<DB::Row>, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_optional_instrumented(pool, self.sql).await
    }

    pub async fn fetch_all<P>(self, pool: &P) -> Result<Vec<DB::Row>, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_all_instrumented(pool, self.sql).await
    }
}

// Implementation for sqlx::query::Map<'q, DB, F, A> (from query_as! macro)
impl<'q, DB, F, A, O> InstrumentedQueryBuilder<sqlx::query::Map<'q, DB, F, A>>
where
    DB: Database,
    F: FnMut(DB::Row) -> Result<O, sqlx::Error> + Send,
    O: Send + Unpin,
    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
{
    pub async fn fetch_one<P>(self, pool: &P) -> Result<O, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_one_instrumented(pool, self.sql).await
    }

    pub async fn fetch_optional<P>(self, pool: &P) -> Result<Option<O>, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_optional_instrumented(pool, self.sql).await
    }

    pub async fn fetch_all<P>(self, pool: &P) -> Result<Vec<O>, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_all_instrumented(pool, self.sql).await
    }
}

// Implementation for sqlx::query::QueryScalar<'q, DB, O, A> (from query_scalar! macro)
impl<'q, DB, O, A> InstrumentedQueryBuilder<sqlx::query::QueryScalar<'q, DB, O, A>>
where
    DB: Database,
    O: Send + Unpin + for<'r> sqlx::Decode<'r, DB> + sqlx::Type<DB>,
    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
    usize: sqlx::ColumnIndex<DB::Row>,
{
    pub async fn fetch_one<P>(self, pool: &P) -> Result<O, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_one_instrumented(pool, self.sql).await
    }

    pub async fn fetch_optional<P>(self, pool: &P) -> Result<Option<O>, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_optional_instrumented(pool, self.sql).await
    }

    pub async fn fetch_all<P>(self, pool: &P) -> Result<Vec<O>, sqlx::Error>
    where
        P: InstrumentedPool<Database = DB> + Send + Sync,
        for<'c> &'c P: Executor<'c, Database = DB>,
    {
        self.query.fetch_all_instrumented(pool, self.sql).await
    }
}