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
use anyhow::Result;
use tokio::time::{self, Duration};
use tokio_postgres::{error::SqlState, Client, Error, NoTls, Row};
use ore::task;
pub async fn client(host: &str, port: u16) -> Result<Client> {
let (mz_client, conn) = tokio_postgres::Config::new()
.user("materialize")
.host(host)
.port(port)
.connect(NoTls)
.await?;
task::spawn(|| "test_util_mz_client", async move {
if let Err(e) = conn.await {
panic!("connection error: {}", e);
}
});
Ok(mz_client)
}
pub async fn try_query(mz_client: &Client, query: &str, delay: Duration) -> Result<Vec<Row>> {
loop {
let timer = std::time::Instant::now();
match mz_client.query(&*query, &[]).await {
Ok(rows) => return Ok(rows),
Err(e) => check_error(e)?,
}
delay_for(timer.elapsed(), delay).await;
}
}
pub async fn try_query_one(mz_client: &Client, query: &str, delay: Duration) -> Result<Row> {
loop {
let timer = std::time::Instant::now();
match mz_client.query_one(&*query, &[]).await {
Ok(rows) => return Ok(rows),
Err(e) => check_error(e)?,
}
delay_for(timer.elapsed(), delay).await;
}
}
fn check_error(e: Error) -> Result<()> {
if e.code() == Some(&SqlState::SQL_STATEMENT_NOT_YET_COMPLETE) {
tracing::info!("Error querying, will try again... {}", e.to_string());
Ok(())
} else {
Err(anyhow::Error::from(e))
}
}
async fn delay_for(elapsed: Duration, delay: Duration) {
if elapsed < delay {
time::sleep(delay - elapsed).await;
} else {
tracing::info!(
"Expected to query for records in {:#?}, took {:#?}",
delay,
elapsed
);
}
}
pub async fn show_sources(mz_client: &Client) -> Result<Vec<String>> {
let mut res = Vec::new();
for row in mz_client.query("SHOW SOURCES", &[]).await? {
res.push(row.get(0))
}
Ok(res)
}
pub async fn drop_source(mz_client: &Client, name: &str) -> Result<()> {
let q = format!("DROP SOURCE IF EXISTS {} CASCADE", name);
tracing::debug!("deleting source=> {}", q);
mz_client.execute(&*q, &[]).await?;
Ok(())
}
pub async fn execute(mz_client: &Client, query: &str) -> Result<u64> {
tracing::debug!("exec=> {}", query);
Ok(mz_client.execute(query, &[]).await?)
}
pub async fn drop_index(mz_client: &Client, name: &str) -> Result<()> {
let q = format!("DROP INDEX {}", name);
tracing::debug!("deleting index=> {}", q);
mz_client.execute(&*q, &[]).await?;
Ok(())
}