petl 0.1.0

Framework of piped ETL
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use anyhow::Result;
use async_trait::async_trait;
// use deadpool_postgres::tokio_postgres::Row;
// use deadpool_postgres::tokio_postgres::types::{FromSql, ToSql};
use tokio_postgres::Row;
use tokio_postgres::types::{FromSql, ToSql};
use tracing::{debug, error, trace};

/// An extension for shortcutting some Postgres-scraping protocols.
#[async_trait]
pub trait PgExtractExt {
    /// If a single entity exists, return it.
    async fn select<'a, T>(&self, fetch_stmt: &'a str, params: &[&(dyn ToSql + Sync)]) -> Result<T>
    where
        T: for<'b> FromSql<'b>;

    /// If a single entity exists, return it; if not, insert it.
    async fn select_else_insert<'a, T>(
        &self,
        fetch_stmt: &'a str,
        insert_stmt: &'a str,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<T>
    where
        T: for<'b> FromSql<'b> + 'static;

    /// Return some collection from the database.
    async fn select_collection<C, T, F, 'a>(
        &self,
        fetch_stmt: &'a str,
        params: &[&(dyn ToSql + Sync)],
        f: F,
    ) -> Result<C>
    where
        C: FromIterator<T>,
        T: Send,
        F: FnMut(&Row) -> T + Send;
}

#[async_trait]
impl PgExtractExt for tokio_postgres::Client {
    async fn select<'a, T>(&self, fetch_stmt: &'a str, params: &[&(dyn ToSql + Sync)]) -> Result<T>
    where
        T: for<'b> FromSql<'b>,
    {
        trace!(fetch_stmt = %fetch_stmt, "Fetching data for query");
        let data: Row = match self.query_one(fetch_stmt, params).await {
            Ok(response) => response,
            Err(e) => {
                error!(stmt = %fetch_stmt, "{e}");
                return Err(anyhow::anyhow!(e));
            }
        };

        Ok(data.get(0))
    }

    async fn select_else_insert<'a, T>(
        &self,
        fetch_stmt: &'a str,
        insert_stmt: &'a str,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<T>
    where
        T: for<'b> FromSql<'b>,
    {
        // Attempt to find the Source PK in the existing table.
        let data: Row = match self.query_one(fetch_stmt, params).await {
            Ok(response) => response,

            // If no PK is found, insert a new one, and reattempt to find it.
            Err(e) => {
                debug!("Did not find data for query: {fetch_stmt} - inserting data instead: {e}");
                self.query_one(insert_stmt, params).await?;

                match self.query_one(fetch_stmt, params).await {
                    Ok(second_response) => second_response,
                    Err(e) => {
                        error!(fetch_stmt = %fetch_stmt, insert_stmt = %insert_stmt, "Failed to insert and retrieve new data");
                        return Err(anyhow::anyhow!(e));
                    }
                }
            }
        };
        Ok(data.get(0))
    }

    async fn select_collection<'a, C, T, F>(
        &self,
        fetch_stmt: &'a str,
        params: &[&(dyn ToSql + Sync)],
        f: F,
    ) -> Result<C>
    where
        C: FromIterator<T>,
        T: Send,
        F: FnMut(&Row) -> T + Send,
    {
        // Return the collection from the pg database.
        let data: Vec<Row> = match self.query(fetch_stmt, params).await {
            Ok(response) => response,
            Err(e) => {
                error!(fetch_stmt = %fetch_stmt, "Failed to fetch collection");
                return Err(anyhow::anyhow!(e));
            }
        };

        // Transform the array of [`tokio_postgres::Row`] with some closure.
        let output: C = data.iter().map(f).collect();
        Ok(output)
    }
}

#[async_trait]
impl PgExtractExt for deadpool_postgres::Pool {
    async fn select<'a, T>(&self, fetch_stmt: &'a str, params: &[&(dyn ToSql + Sync)]) -> Result<T>
    where
        T: for<'b> FromSql<'b>,
    {
        let pg_client = self.get().await?;
        Ok(pg_client.select(fetch_stmt, params).await?)
    }

    async fn select_else_insert<'a, T>(
        &self,
        fetch_stmt: &'a str,
        insert_stmt: &'a str,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<T>
    where
        T: for<'b> FromSql<'b> + 'static,
    {
        let pg_client = self.get().await?;
        let data = pg_client
            .select_else_insert(fetch_stmt, insert_stmt, params)
            .await?;
        Ok(data)
    }

    async fn select_collection<'a, C, T, F>(
        &self,
        fetch_stmt: &'a str,
        params: &[&(dyn ToSql + Sync)],
        f: F,
    ) -> Result<C>
    where
        C: FromIterator<T>,
        T: Send,
        F: FnMut(&Row) -> T + Send,
    {
        let pg_client = self.get().await?;
        let data = pg_client.select_collection(fetch_stmt, params, f).await?;
        Ok(data)
    }
}