microcosm.blue
Server tools to backfill, tail, mirror, and verify PLC logs
keep track of pg connection tasks
bad-example.com
+41
-59
src/plc_pg.rs
+41
-59
src/plc_pg.rs
···
4
use std::path::PathBuf;
5
use std::pin::pin;
6
use std::time::Instant;
7
-
use tokio::sync::{mpsc, oneshot};
8
use tokio_postgres::{
9
Client, Error as PgError, NoTls,
10
binary_copy::BinaryCopyInWriter,
11
connect,
12
types::{Json, Type},
13
};
14
15
-
fn get_tls(cert: PathBuf) -> MakeTlsConnector {
16
-
let cert = std::fs::read(cert).expect("to read cert file");
17
-
let cert = Certificate::from_pem(&cert).expect("to build cert");
18
-
let connector = TlsConnector::builder()
19
-
.add_root_certificate(cert)
20
-
.build()
21
-
.expect("to build tls connector");
22
-
MakeTlsConnector::new(connector)
23
}
24
25
/// a little tokio-postgres helper
26
///
27
/// it's clone for easiness. it doesn't share any resources underneath after
28
-
/// cloning at all so it's not meant for
29
#[derive(Clone)]
30
pub struct Db {
31
pg_uri: String,
···
38
// it's what we expect: check for db migrations.
39
log::trace!("checking migrations...");
40
41
-
let connector = cert.map(get_tls);
42
43
-
let (client, connection_task) = if let Some(ref connector) = connector {
44
-
let (client, connection) = connect(pg_uri, connector.clone()).await?;
45
-
let task = tokio::task::spawn(async move {
46
-
connection
47
-
.await
48
-
.inspect_err(|e| log::error!("connection ended with error: {e}"))
49
-
.expect("pg validation connection not to blow up");
50
-
});
51
-
(client, task)
52
} else {
53
-
let (client, connection) = connect(pg_uri, NoTls).await?;
54
-
let task = tokio::task::spawn(async move {
55
-
connection
56
-
.await
57
-
.inspect_err(|e| log::error!("connection ended with error: {e}"))
58
-
.expect("pg validation connection not to blow up");
59
-
});
60
-
(client, task)
61
};
62
63
let migrations: Vec<String> = client
···
77
);
78
drop(client);
79
// make sure the connection worker thing doesn't linger
80
-
connection_task.await?;
81
log::info!("db connection succeeded and plc migrations appear as expected");
82
83
Ok(Self {
···
86
})
87
}
88
89
-
pub async fn connect(&self) -> Result<Client, PgError> {
90
log::trace!("connecting postgres...");
91
-
let client = if let Some(ref connector) = self.cert {
92
-
let (client, connection) = connect(&self.pg_uri, connector.clone()).await?;
93
-
94
-
// send the connection away to do the actual communication work
95
-
// apparently the connection will complete when the client drops
96
-
tokio::task::spawn(async move {
97
-
connection
98
-
.await
99
-
.inspect_err(|e| log::error!("connection ended with error: {e}"))
100
-
.expect("pg connection not to blow up");
101
-
});
102
-
client
103
} else {
104
-
let (client, connection) = connect(&self.pg_uri, NoTls).await?;
105
-
106
-
// send the connection away to do the actual communication work
107
-
// apparently the connection will complete when the client drops
108
-
tokio::task::spawn(async move {
109
-
connection
110
-
.await
111
-
.inspect_err(|e| log::error!("connection ended with error: {e}"))
112
-
.expect("pg connection not to blow up");
113
-
});
114
-
client
115
-
};
116
-
117
-
Ok(client)
118
}
119
120
pub async fn get_latest(&self) -> Result<Option<Dt>, PgError> {
121
-
let client = self.connect().await?;
122
let dt: Option<Dt> = client
123
.query_opt(
124
r#"SELECT "createdAt"
···
129
)
130
.await?
131
.map(|row| row.get(0));
132
Ok(dt)
133
}
134
}
···
137
db: Db,
138
mut pages: mpsc::Receiver<ExportPage>,
139
) -> anyhow::Result<&'static str> {
140
-
let mut client = db.connect().await?;
141
142
let ops_stmt = client
143
.prepare(
···
174
}
175
tx.commit().await?;
176
}
177
178
log::info!(
179
"no more pages. inserted {ops_inserted} ops and {dids_inserted} dids in {:?}",
···
201
mut pages: mpsc::Receiver<ExportPage>,
202
notify_last_at: Option<oneshot::Sender<Option<Dt>>>,
203
) -> anyhow::Result<&'static str> {
204
-
let mut client = db.connect().await?;
205
206
let t0 = Instant::now();
207
let tx = client.transaction().await?;
···
316
log::trace!("set tables LOGGED: {:?}", t_step.elapsed());
317
318
tx.commit().await?;
319
log::info!("total backfill time: {:?}", t0.elapsed());
320
321
Ok("backfill_to_pg")
···
4
use std::path::PathBuf;
5
use std::pin::pin;
6
use std::time::Instant;
7
+
use tokio::{
8
+
task::{spawn, JoinHandle},
9
+
sync::{mpsc, oneshot},
10
+
};
11
use tokio_postgres::{
12
Client, Error as PgError, NoTls,
13
binary_copy::BinaryCopyInWriter,
14
connect,
15
+
Socket,
16
types::{Json, Type},
17
+
tls::MakeTlsConnect
18
};
19
20
+
fn get_tls(cert: PathBuf) -> anyhow::Result<MakeTlsConnector> {
21
+
let cert = std::fs::read(cert)?;
22
+
let cert = Certificate::from_pem(&cert)?;
23
+
let connector = TlsConnector::builder().add_root_certificate(cert).build()?;
24
+
Ok(MakeTlsConnector::new(connector))
25
+
}
26
+
27
+
async fn get_client_and_task<T>(
28
+
uri: &str,
29
+
connector: T,
30
+
) -> Result<(Client, JoinHandle<Result<(), PgError>>), PgError>
31
+
where
32
+
T: MakeTlsConnect<Socket>,
33
+
<T as MakeTlsConnect<Socket>>::Stream: Send + 'static,
34
+
{
35
+
let (client, connection) = connect(uri, connector).await?;
36
+
Ok((client, spawn(connection)))
37
}
38
39
/// a little tokio-postgres helper
40
///
41
/// it's clone for easiness. it doesn't share any resources underneath after
42
+
/// cloning *at all* so it's not meant for eg. handling public web requests
43
#[derive(Clone)]
44
pub struct Db {
45
pg_uri: String,
···
52
// it's what we expect: check for db migrations.
53
log::trace!("checking migrations...");
54
55
+
let connector = cert.map(get_tls).transpose()?;
56
57
+
let (client, conn_task) = if let Some(ref connector) = connector {
58
+
get_client_and_task(pg_uri, connector.clone()).await?
59
} else {
60
+
get_client_and_task(pg_uri, NoTls).await?
61
};
62
63
let migrations: Vec<String> = client
···
77
);
78
drop(client);
79
// make sure the connection worker thing doesn't linger
80
+
conn_task.await??;
81
log::info!("db connection succeeded and plc migrations appear as expected");
82
83
Ok(Self {
···
86
})
87
}
88
89
+
#[must_use]
90
+
pub async fn connect(&self) -> Result<(Client, JoinHandle<Result<(), PgError>>), PgError> {
91
log::trace!("connecting postgres...");
92
+
if let Some(ref connector) = self.cert {
93
+
get_client_and_task(&self.pg_uri, connector.clone()).await
94
} else {
95
+
get_client_and_task(&self.pg_uri, NoTls).await
96
+
}
97
}
98
99
pub async fn get_latest(&self) -> Result<Option<Dt>, PgError> {
100
+
let (client, task) = self.connect().await?;
101
let dt: Option<Dt> = client
102
.query_opt(
103
r#"SELECT "createdAt"
···
108
)
109
.await?
110
.map(|row| row.get(0));
111
+
drop(task);
112
Ok(dt)
113
}
114
}
···
117
db: Db,
118
mut pages: mpsc::Receiver<ExportPage>,
119
) -> anyhow::Result<&'static str> {
120
+
let (mut client, task) = db.connect().await?;
121
122
let ops_stmt = client
123
.prepare(
···
154
}
155
tx.commit().await?;
156
}
157
+
drop(task);
158
159
log::info!(
160
"no more pages. inserted {ops_inserted} ops and {dids_inserted} dids in {:?}",
···
182
mut pages: mpsc::Receiver<ExportPage>,
183
notify_last_at: Option<oneshot::Sender<Option<Dt>>>,
184
) -> anyhow::Result<&'static str> {
185
+
let (mut client, task) = db.connect().await?;
186
187
let t0 = Instant::now();
188
let tx = client.transaction().await?;
···
297
log::trace!("set tables LOGGED: {:?}", t_step.elapsed());
298
299
tx.commit().await?;
300
+
drop(task);
301
log::info!("total backfill time: {:?}", t0.elapsed());
302
303
Ok("backfill_to_pg")