mirror of
https://github.com/fluxerapp/fluxer.git
synced 2026-09-02 21:04:06 +03:00
fix(svc): allow disabling named Postgres prepared statements (#2256)
This commit is contained in:
@@ -114,9 +114,10 @@ FLUXER_DISCOVERY_ENABLED=true
|
||||
#FLUXER_ERLANG_SCHEDULERS_MIN=2
|
||||
#FLUXER_ERLANG_SCHEDULERS_MAX=16
|
||||
|
||||
# The api names its fixed Postgres statement shapes so the server can reuse their
|
||||
# plans. Named prepared statements require a session that outlives the
|
||||
# transaction, so set this to false if you put a transaction-pooling connection
|
||||
# pooler such as PgBouncer in front of Postgres. The bundled compose talks to
|
||||
# Postgres directly, where naming is a win and the default is correct.
|
||||
# The api and the Rust services name their fixed Postgres statement shapes so the
|
||||
# server can reuse their plans. Named prepared statements require a session that
|
||||
# outlives the transaction, so set this to false if you put a transaction-pooling
|
||||
# connection pooler such as PgBouncer in front of Postgres. One setting governs
|
||||
# every service. The bundled compose talks to Postgres directly, where naming is
|
||||
# a win and the default is correct.
|
||||
#FLUXER_POSTGRES_PREPARED_STATEMENTS=true
|
||||
|
||||
@@ -56,6 +56,7 @@ pub struct AppProxyConfig {
|
||||
pub postgres_ssl_ca: Option<String>,
|
||||
pub postgres_max_connections: usize,
|
||||
pub postgres_kv_table: String,
|
||||
pub postgres_prepared_statements: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
@@ -250,6 +251,10 @@ impl AppProxyConfig {
|
||||
postgres_ssl_ca: cfg::non_empty_env("FLUXER_POSTGRES_SSL_CA"),
|
||||
postgres_max_connections,
|
||||
postgres_kv_table: cfg::read_env("FLUXER_POSTGRES_KV_TABLE", "fluxer_kv"),
|
||||
postgres_prepared_statements: cfg::read_bool_env(
|
||||
&["FLUXER_POSTGRES_PREPARED_STATEMENTS"],
|
||||
true,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -103,7 +103,7 @@ impl InviteMetaResolver {
|
||||
pub async fn connect(config: &AppProxyConfig) -> anyhow::Result<Self> {
|
||||
match config.database_backend {
|
||||
DatabaseBackend::Postgres => {
|
||||
let pool = fluxer_svc::postgres::connect(&fluxer_svc::postgres::PostgresConfig {
|
||||
let postgres_config = fluxer_svc::postgres::PostgresConfig {
|
||||
url: config.postgres_url.clone(),
|
||||
host: config.postgres_host.clone(),
|
||||
port: config.postgres_port,
|
||||
@@ -114,9 +114,10 @@ impl InviteMetaResolver {
|
||||
ssl_ca: config.postgres_ssl_ca.clone(),
|
||||
max_connections: config.postgres_max_connections,
|
||||
kv_table: config.postgres_kv_table.clone(),
|
||||
})
|
||||
.await?;
|
||||
let kv = postgres::KvClient::new(pool, &config.postgres_kv_table)?;
|
||||
prepared_statements: config.postgres_prepared_statements,
|
||||
};
|
||||
let pool = fluxer_svc::postgres::connect(&postgres_config).await?;
|
||||
let kv = postgres::KvClient::new(pool, &postgres_config)?;
|
||||
Self::new_postgres(kv, config)
|
||||
}
|
||||
DatabaseBackend::Cassandra => {
|
||||
|
||||
@@ -31,7 +31,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
let postgres_config =
|
||||
fluxer_svc::postgres::PostgresConfig::from_service_config(&config);
|
||||
let pool = fluxer_svc::postgres::connect(&postgres_config).await?;
|
||||
let kv = fluxer_svc::postgres::KvClient::new(pool, &postgres_config.kv_table)?;
|
||||
let kv = fluxer_svc::postgres::KvClient::new(pool, &postgres_config)?;
|
||||
MessagesShard::new_postgres(kv, transport.clone())?
|
||||
}
|
||||
DatabaseBackend::Cassandra => {
|
||||
|
||||
@@ -35,6 +35,7 @@ pub struct ServiceConfig {
|
||||
pub postgres_ssl_ca: Option<String>,
|
||||
pub postgres_max_connections: usize,
|
||||
pub postgres_kv_table: String,
|
||||
pub postgres_prepared_statements: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
@@ -144,6 +145,11 @@ impl ServiceConfig {
|
||||
.transpose()?
|
||||
.unwrap_or(20)
|
||||
.max(1);
|
||||
let postgres_prepared_statements =
|
||||
optional_from(&get, "FLUXER_POSTGRES_PREPARED_STATEMENTS")
|
||||
.map(|v| parse_bool(&v))
|
||||
.transpose()?
|
||||
.unwrap_or(true);
|
||||
|
||||
let max_concurrent_requests = optional_from(&get, "FLUXER_SVC_MAX_CONCURRENT_REQUESTS")
|
||||
.map(|v| v.parse::<usize>())
|
||||
@@ -186,6 +192,7 @@ impl ServiceConfig {
|
||||
postgres_max_connections,
|
||||
postgres_kv_table: optional_from(&get, "FLUXER_POSTGRES_KV_TABLE")
|
||||
.unwrap_or_else(|| "fluxer_kv".to_owned()),
|
||||
postgres_prepared_statements,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -344,6 +351,7 @@ mod tests {
|
||||
assert_eq!(None, cfg.postgres_ssl_ca);
|
||||
assert_eq!(20, cfg.postgres_max_connections);
|
||||
assert_eq!("fluxer_kv", cfg.postgres_kv_table);
|
||||
assert!(cfg.postgres_prepared_statements);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -360,6 +368,7 @@ mod tests {
|
||||
("FLUXER_POSTGRES_SSL_CA", "ca-pem"),
|
||||
("FLUXER_POSTGRES_MAX_CONNECTIONS", "7"),
|
||||
("FLUXER_POSTGRES_KV_TABLE", "fluxer_kv_dev"),
|
||||
("FLUXER_POSTGRES_PREPARED_STATEMENTS", "false"),
|
||||
]);
|
||||
|
||||
assert_eq!(DatabaseBackend::Postgres, cfg.database_backend);
|
||||
@@ -376,6 +385,16 @@ mod tests {
|
||||
assert_eq!(Some("ca-pem".to_owned()), cfg.postgres_ssl_ca);
|
||||
assert_eq!(7, cfg.postgres_max_connections);
|
||||
assert_eq!("fluxer_kv_dev", cfg.postgres_kv_table);
|
||||
assert!(!cfg.postgres_prepared_statements);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_a_non_boolean_prepared_statements_value() {
|
||||
let result = ServiceConfig::from_env_reader(|name| {
|
||||
(name == "FLUXER_POSTGRES_PREPARED_STATEMENTS").then(|| "maybe".to_owned())
|
||||
});
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+194
-48
@@ -10,7 +10,11 @@ use rustls::{
|
||||
};
|
||||
use serde_json::{Map, Number, Value};
|
||||
use std::str::FromStr;
|
||||
use tokio_postgres::{Config as PgConfig, Row, Statement, config::SslMode, types::ToSql};
|
||||
use tokio_postgres::{
|
||||
Config as PgConfig, Row,
|
||||
config::SslMode,
|
||||
types::{ToSql, Type},
|
||||
};
|
||||
use tokio_postgres_rustls::MakeRustlsConnect;
|
||||
|
||||
const POSTGRES_KV_SCHEMA_LOCK_NAMESPACE: i32 = 0x4658_4b56;
|
||||
@@ -29,6 +33,7 @@ pub struct PostgresConfig {
|
||||
pub ssl_ca: Option<String>,
|
||||
pub max_connections: usize,
|
||||
pub kv_table: String,
|
||||
pub prepared_statements: bool,
|
||||
}
|
||||
|
||||
impl PostgresConfig {
|
||||
@@ -44,6 +49,7 @@ impl PostgresConfig {
|
||||
ssl_ca: config.postgres_ssl_ca.clone(),
|
||||
max_connections: config.postgres_max_connections,
|
||||
kv_table: config.postgres_kv_table.clone(),
|
||||
prepared_statements: config.postgres_prepared_statements,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -92,6 +98,7 @@ pub async fn connect(config: &PostgresConfig) -> anyhow::Result<Pool> {
|
||||
database = config.database,
|
||||
max_connections = config.max_connections,
|
||||
kv_table = config.kv_table,
|
||||
prepared_statements = config.prepared_statements,
|
||||
"connected to Postgres"
|
||||
);
|
||||
Ok(pool)
|
||||
@@ -250,6 +257,7 @@ fn is_safe_identifier(identifier: &str) -> bool {
|
||||
#[derive(Clone)]
|
||||
pub struct KvClient {
|
||||
pool: Pool,
|
||||
prepared_statements: bool,
|
||||
get_row_sql: String,
|
||||
get_rows_sql: String,
|
||||
get_partition_rows_sql: String,
|
||||
@@ -258,10 +266,11 @@ pub struct KvClient {
|
||||
}
|
||||
|
||||
impl KvClient {
|
||||
pub fn new(pool: Pool, kv_table: &str) -> anyhow::Result<Self> {
|
||||
let table = quote_identifier(kv_table)?;
|
||||
pub fn new(pool: Pool, config: &PostgresConfig) -> anyhow::Result<Self> {
|
||||
let table = quote_identifier(&config.kv_table)?;
|
||||
Ok(Self {
|
||||
pool,
|
||||
prepared_statements: config.prepared_statements,
|
||||
get_row_sql: format!(
|
||||
"SELECT row_data FROM {table} WHERE table_name = $1 AND row_key = $2 AND (expires_at IS NULL OR expires_at > now()) LIMIT 1"
|
||||
),
|
||||
@@ -278,11 +287,75 @@ impl KvClient {
|
||||
})
|
||||
}
|
||||
|
||||
async fn query_rows(
|
||||
&self,
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
params: &[(&(dyn ToSql + Sync), Type)],
|
||||
) -> anyhow::Result<Vec<Row>> {
|
||||
if !self.prepared_statements {
|
||||
return Ok(client.query_typed(sql, params).await?);
|
||||
}
|
||||
let statement = client.prepare_cached(sql).await?;
|
||||
Ok(client.query(&statement, &statement_params(params)).await?)
|
||||
}
|
||||
|
||||
async fn query_row(
|
||||
&self,
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
params: &[(&(dyn ToSql + Sync), Type)],
|
||||
) -> anyhow::Result<Option<Row>> {
|
||||
if !self.prepared_statements {
|
||||
return Ok(client.query_typed_opt(sql, params).await?);
|
||||
}
|
||||
let statement = client.prepare_cached(sql).await?;
|
||||
Ok(client
|
||||
.query_opt(&statement, &statement_params(params))
|
||||
.await?)
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
params: &[(&(dyn ToSql + Sync), Type)],
|
||||
) -> anyhow::Result<u64> {
|
||||
if !self.prepared_statements {
|
||||
return Ok(client.execute_typed(sql, params).await?);
|
||||
}
|
||||
let statement = client.prepare_cached(sql).await?;
|
||||
Ok(client
|
||||
.execute(&statement, &statement_params(params))
|
||||
.await?)
|
||||
}
|
||||
|
||||
async fn query_dynamic_rows(
|
||||
&self,
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
field_name: &str,
|
||||
params: &[(&(dyn ToSql + Sync), Type)],
|
||||
) -> anyhow::Result<Vec<Row>> {
|
||||
if !self.prepared_statements {
|
||||
return Ok(client.query_typed(sql, params).await?);
|
||||
}
|
||||
let statement = if is_cached_json_field(field_name) {
|
||||
client.prepare_cached(sql).await?
|
||||
} else {
|
||||
client.prepare(sql).await?
|
||||
};
|
||||
Ok(client.query(&statement, &statement_params(params)).await?)
|
||||
}
|
||||
|
||||
pub async fn get_row(&self, table_name: &str, row_key: &str) -> anyhow::Result<Option<Value>> {
|
||||
let client = self.pool.get().await?;
|
||||
let statement = client.prepare_cached(&self.get_row_sql).await?;
|
||||
let row = client
|
||||
.query_opt(&statement, &[&table_name, &row_key])
|
||||
let row = self
|
||||
.query_row(
|
||||
&client,
|
||||
&self.get_row_sql,
|
||||
&[(&table_name, Type::TEXT), (&row_key, Type::TEXT)],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.map(|row| row.get::<_, Value>("row_data")))
|
||||
}
|
||||
@@ -296,8 +369,13 @@ impl KvClient {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let client = self.pool.get().await?;
|
||||
let statement = client.prepare_cached(&self.get_rows_sql).await?;
|
||||
let rows = client.query(&statement, &[&table_name, &row_keys]).await?;
|
||||
let rows = self
|
||||
.query_rows(
|
||||
&client,
|
||||
&self.get_rows_sql,
|
||||
&[(&table_name, Type::TEXT), (&row_keys, Type::TEXT_ARRAY)],
|
||||
)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(row_key_and_data).collect())
|
||||
}
|
||||
|
||||
@@ -307,9 +385,12 @@ impl KvClient {
|
||||
partition_key: &str,
|
||||
) -> anyhow::Result<Vec<(String, Value)>> {
|
||||
let client = self.pool.get().await?;
|
||||
let statement = client.prepare_cached(&self.get_partition_rows_sql).await?;
|
||||
let rows = client
|
||||
.query(&statement, &[&table_name, &partition_key])
|
||||
let rows = self
|
||||
.query_rows(
|
||||
&client,
|
||||
&self.get_partition_rows_sql,
|
||||
&[(&table_name, Type::TEXT), (&partition_key, Type::TEXT)],
|
||||
)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(row_key_and_data).collect())
|
||||
}
|
||||
@@ -321,11 +402,16 @@ impl KvClient {
|
||||
) -> anyhow::Result<Vec<(String, Value)>> {
|
||||
let client = self.pool.get().await?;
|
||||
let upper = format!("{row_key_prefix}\u{10ffff}");
|
||||
let statement = client
|
||||
.prepare_cached(&self.get_row_key_prefix_rows_sql)
|
||||
.await?;
|
||||
let rows = client
|
||||
.query(&statement, &[&table_name, &row_key_prefix, &upper])
|
||||
let rows = self
|
||||
.query_rows(
|
||||
&client,
|
||||
&self.get_row_key_prefix_rows_sql,
|
||||
&[
|
||||
(&table_name, Type::TEXT),
|
||||
(&row_key_prefix, Type::TEXT),
|
||||
(&upper, Type::TEXT),
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(row_key_and_data).collect())
|
||||
}
|
||||
@@ -351,26 +437,49 @@ impl KvClient {
|
||||
let sql = format!(
|
||||
"{base} AND {field_expr} < $3 ORDER BY {field_expr} {direction} LIMIT $4"
|
||||
);
|
||||
let statement = prepare_dynamic(&client, &sql, field_name).await?;
|
||||
client
|
||||
.query(&statement, &[&table_name, &partition_key, &value, &limit])
|
||||
.await?
|
||||
self.query_dynamic_rows(
|
||||
&client,
|
||||
&sql,
|
||||
field_name,
|
||||
&[
|
||||
(&table_name, Type::TEXT),
|
||||
(&partition_key, Type::TEXT),
|
||||
(&value, Type::INT8),
|
||||
(&limit, Type::INT8),
|
||||
],
|
||||
)
|
||||
.await?
|
||||
}
|
||||
Some(BigIntBound::GreaterThan(value)) => {
|
||||
let sql = format!(
|
||||
"{base} AND {field_expr} > $3 ORDER BY {field_expr} {direction} LIMIT $4"
|
||||
);
|
||||
let statement = prepare_dynamic(&client, &sql, field_name).await?;
|
||||
client
|
||||
.query(&statement, &[&table_name, &partition_key, &value, &limit])
|
||||
.await?
|
||||
self.query_dynamic_rows(
|
||||
&client,
|
||||
&sql,
|
||||
field_name,
|
||||
&[
|
||||
(&table_name, Type::TEXT),
|
||||
(&partition_key, Type::TEXT),
|
||||
(&value, Type::INT8),
|
||||
(&limit, Type::INT8),
|
||||
],
|
||||
)
|
||||
.await?
|
||||
}
|
||||
None => {
|
||||
let sql = format!("{base} ORDER BY {field_expr} {direction} LIMIT $3");
|
||||
let statement = prepare_dynamic(&client, &sql, field_name).await?;
|
||||
client
|
||||
.query(&statement, &[&table_name, &partition_key, &limit])
|
||||
.await?
|
||||
self.query_dynamic_rows(
|
||||
&client,
|
||||
&sql,
|
||||
field_name,
|
||||
&[
|
||||
(&table_name, Type::TEXT),
|
||||
(&partition_key, Type::TEXT),
|
||||
(&limit, Type::INT8),
|
||||
],
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
Ok(rows.into_iter().map(row_key_and_data).collect())
|
||||
@@ -390,46 +499,55 @@ impl KvClient {
|
||||
let field_expr = json_field_expr(field_name)?;
|
||||
let base = &self.get_partition_rows_sql;
|
||||
let sql = format!("{base} AND {field_expr} = ANY($3::bigint[])");
|
||||
let statement = prepare_dynamic(&client, &sql, field_name).await?;
|
||||
let rows = client
|
||||
.query(&statement, &[&table_name, &partition_key, &values])
|
||||
let rows = self
|
||||
.query_dynamic_rows(
|
||||
&client,
|
||||
&sql,
|
||||
field_name,
|
||||
&[
|
||||
(&table_name, Type::TEXT),
|
||||
(&partition_key, Type::TEXT),
|
||||
(&values, Type::INT8_ARRAY),
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(row_key_and_data).collect())
|
||||
}
|
||||
|
||||
pub async fn delete_row(&self, table_name: &str, row_key: &str) -> anyhow::Result<()> {
|
||||
let client = self.pool.get().await?;
|
||||
let statement = client.prepare_cached(&self.delete_row_sql).await?;
|
||||
client.execute(&statement, &[&table_name, &row_key]).await?;
|
||||
self.execute(
|
||||
&client,
|
||||
&self.delete_row_sql,
|
||||
&[(&table_name, Type::TEXT), (&row_key, Type::TEXT)],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn query(
|
||||
&self,
|
||||
sql: &str,
|
||||
params: &[&(dyn ToSql + Sync)],
|
||||
params: &[(&(dyn ToSql + Sync), Type)],
|
||||
) -> anyhow::Result<Vec<Row>> {
|
||||
let client = self.pool.get().await?;
|
||||
Ok(client.query(sql, params).await?)
|
||||
if !self.prepared_statements {
|
||||
return Ok(client.query_typed(sql, params).await?);
|
||||
}
|
||||
Ok(client.query(sql, &statement_params(params)).await?)
|
||||
}
|
||||
}
|
||||
|
||||
fn statement_params<'a>(
|
||||
params: &'a [(&'a (dyn ToSql + Sync), Type)],
|
||||
) -> Vec<&'a (dyn ToSql + Sync)> {
|
||||
params.iter().map(|(value, _)| *value).collect()
|
||||
}
|
||||
|
||||
fn is_cached_json_field(field_name: &str) -> bool {
|
||||
CACHED_JSON_FIELDS.contains(&field_name)
|
||||
}
|
||||
|
||||
async fn prepare_dynamic(
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
field_name: &str,
|
||||
) -> anyhow::Result<Statement> {
|
||||
if is_cached_json_field(field_name) {
|
||||
Ok(client.prepare_cached(sql).await?)
|
||||
} else {
|
||||
Ok(client.prepare(sql).await?)
|
||||
}
|
||||
}
|
||||
|
||||
fn row_key_and_data(row: Row) -> (String, Value) {
|
||||
(
|
||||
row.get::<_, String>("row_key"),
|
||||
@@ -611,11 +729,27 @@ mod tests {
|
||||
assert_eq!(decoded["metadata"], json!([["kind", 9]]));
|
||||
}
|
||||
|
||||
fn test_postgres_config(kv_table: &str) -> PostgresConfig {
|
||||
PostgresConfig {
|
||||
url: None,
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 5432,
|
||||
database: "fluxer".to_owned(),
|
||||
username: "fluxer".to_owned(),
|
||||
password: None,
|
||||
ssl: false,
|
||||
ssl_ca: None,
|
||||
max_connections: 1,
|
||||
kv_table: kv_table.to_owned(),
|
||||
prepared_statements: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_kv_client(kv_table: &str) -> KvClient {
|
||||
let pg = PgConfig::from_str("postgres://fluxer@127.0.0.1:5432/fluxer").unwrap();
|
||||
let manager = Manager::new(pg, build_disabled_tls_connector());
|
||||
let pool = Pool::builder(manager).max_size(1).build().unwrap();
|
||||
KvClient::new(pool, kv_table).unwrap()
|
||||
KvClient::new(pool, &test_postgres_config(kv_table)).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -644,6 +778,18 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn carries_the_prepared_statement_switch_onto_the_client() {
|
||||
let mut config = test_postgres_config("fluxer_kv");
|
||||
config.prepared_statements = false;
|
||||
let pg = PgConfig::from_str("postgres://fluxer@127.0.0.1:5432/fluxer").unwrap();
|
||||
let manager = Manager::new(pg, build_disabled_tls_connector());
|
||||
let pool = Pool::builder(manager).max_size(1).build().unwrap();
|
||||
|
||||
assert!(!KvClient::new(pool, &config).unwrap().prepared_statements);
|
||||
assert!(test_kv_client("fluxer_kv").prepared_statements);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn caches_only_closed_set_json_fields() {
|
||||
assert!(is_cached_json_field("message_id"));
|
||||
|
||||
@@ -886,6 +886,7 @@ mod tests {
|
||||
postgres_ssl_ca: None,
|
||||
postgres_max_connections: 1,
|
||||
postgres_kv_table: "fluxer_kv".to_owned(),
|
||||
postgres_prepared_statements: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -435,6 +435,7 @@ mod tests {
|
||||
postgres_ssl_ca: None,
|
||||
postgres_max_connections: 1,
|
||||
postgres_kv_table: "fluxer_kv".to_owned(),
|
||||
postgres_prepared_statements: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,336 @@
|
||||
// SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
use deadpool_postgres::Pool;
|
||||
use fluxer_svc::postgres::{BigIntBound, KvClient, PostgresConfig, connect};
|
||||
use serde_json::{Value, json};
|
||||
use std::net::TcpListener;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::time::Duration;
|
||||
use tokio_postgres::types::Type;
|
||||
|
||||
const KV_TABLE: &str = "svc_stmt_names";
|
||||
const KV_TABLE_POOLED: &str = "svc_pooled_stmt_names";
|
||||
const TABLE_NAME: &str = "messages";
|
||||
|
||||
fn docker_available() -> bool {
|
||||
Command::new("docker")
|
||||
.arg("version")
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.is_ok_and(|status| status.success())
|
||||
}
|
||||
|
||||
fn docker(args: &[&str]) -> anyhow::Result<()> {
|
||||
let output = Command::new("docker").args(args).output()?;
|
||||
if output.status.success() {
|
||||
return Ok(());
|
||||
}
|
||||
anyhow::bail!(
|
||||
"docker {} failed: {}",
|
||||
args.join(" "),
|
||||
String::from_utf8_lossy(&output.stderr).trim().to_owned()
|
||||
)
|
||||
}
|
||||
|
||||
fn executed_statements(container: &str, kv_table: &str) -> anyhow::Result<Vec<String>> {
|
||||
let output = Command::new("docker").args(["logs", container]).output()?;
|
||||
let needle = format!("\"{kv_table}\"");
|
||||
Ok(String::from_utf8_lossy(&output.stderr)
|
||||
.lines()
|
||||
.chain(String::from_utf8_lossy(&output.stdout).lines())
|
||||
.filter(|line| line.contains("execute ") && line.contains(&needle))
|
||||
.map(str::to_owned)
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn free_port() -> anyhow::Result<u16> {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")?;
|
||||
Ok(listener.local_addr()?.port())
|
||||
}
|
||||
|
||||
fn postgres_config(port: u16, kv_table: &str, prepared_statements: bool) -> PostgresConfig {
|
||||
PostgresConfig {
|
||||
url: None,
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
database: "fluxer".to_owned(),
|
||||
username: "fluxer".to_owned(),
|
||||
password: Some("fluxer".to_owned()),
|
||||
ssl: false,
|
||||
ssl_ca: None,
|
||||
max_connections: 1,
|
||||
kv_table: kv_table.to_owned(),
|
||||
prepared_statements,
|
||||
}
|
||||
}
|
||||
|
||||
fn tagged_bigint(value: i64) -> Value {
|
||||
json!({"__fluxer_type": "bigint", "value": value.to_string()})
|
||||
}
|
||||
|
||||
fn seed_rows() -> Vec<(&'static str, &'static str, Value)> {
|
||||
vec![
|
||||
(
|
||||
"p0",
|
||||
"r1",
|
||||
json!({"message_id": tagged_bigint(10), "user_id": tagged_bigint(1010), "payload": "one"}),
|
||||
),
|
||||
(
|
||||
"p0",
|
||||
"r2",
|
||||
json!({"message_id": tagged_bigint(20), "user_id": tagged_bigint(1020), "payload": "two"}),
|
||||
),
|
||||
(
|
||||
"p0",
|
||||
"r3",
|
||||
json!({"message_id": tagged_bigint(30), "user_id": tagged_bigint(1030), "payload": "three"}),
|
||||
),
|
||||
(
|
||||
"p1",
|
||||
"z9",
|
||||
json!({"message_id": tagged_bigint(40), "user_id": tagged_bigint(1040), "payload": "four"}),
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
async fn seed(kv: &KvClient, kv_table: &str) -> anyhow::Result<()> {
|
||||
let sql = format!(
|
||||
"INSERT INTO \"{kv_table}\" (table_name, partition_key, row_key, row_data) VALUES ($1, $2, $3, $4)"
|
||||
);
|
||||
for (partition_key, row_key, row_data) in seed_rows() {
|
||||
kv.query(
|
||||
&sql,
|
||||
&[
|
||||
(&TABLE_NAME, Type::TEXT),
|
||||
(&partition_key, Type::TEXT),
|
||||
(&row_key, Type::TEXT),
|
||||
(&row_data, Type::JSONB),
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn render(rows: Vec<(String, Value)>) -> String {
|
||||
rows.into_iter()
|
||||
.map(|(row_key, row_data)| format!("{row_key}={row_data}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("|")
|
||||
}
|
||||
|
||||
fn render_sorted(mut rows: Vec<(String, Value)>) -> String {
|
||||
rows.sort_by(|left, right| left.0.cmp(&right.0));
|
||||
render(rows)
|
||||
}
|
||||
|
||||
async fn exercise_kv_shapes(kv: &KvClient) -> anyhow::Result<Vec<String>> {
|
||||
let mut observations = Vec::new();
|
||||
observations.push(format!("{:?}", kv.get_row(TABLE_NAME, "r1").await?));
|
||||
observations.push(format!("{:?}", kv.get_row(TABLE_NAME, "missing").await?));
|
||||
observations.push(render_sorted(
|
||||
kv.get_rows(TABLE_NAME, &["r1".to_owned(), "r3".to_owned()])
|
||||
.await?,
|
||||
));
|
||||
observations.push(render_sorted(kv.get_rows(TABLE_NAME, &[]).await?));
|
||||
observations.push(render_sorted(
|
||||
kv.get_partition_rows(TABLE_NAME, "p0").await?,
|
||||
));
|
||||
observations.push(render_sorted(
|
||||
kv.get_row_key_prefix_rows(TABLE_NAME, "r").await?,
|
||||
));
|
||||
observations.push(render(
|
||||
kv.get_partition_rows_by_bigint_field(
|
||||
TABLE_NAME,
|
||||
"p0",
|
||||
"message_id",
|
||||
Some(BigIntBound::LessThan(30)),
|
||||
true,
|
||||
10,
|
||||
)
|
||||
.await?,
|
||||
));
|
||||
observations.push(render(
|
||||
kv.get_partition_rows_by_bigint_field(
|
||||
TABLE_NAME,
|
||||
"p0",
|
||||
"message_id",
|
||||
Some(BigIntBound::GreaterThan(10)),
|
||||
false,
|
||||
10,
|
||||
)
|
||||
.await?,
|
||||
));
|
||||
observations.push(render(
|
||||
kv.get_partition_rows_by_bigint_field(TABLE_NAME, "p0", "message_id", None, true, 2)
|
||||
.await?,
|
||||
));
|
||||
observations.push(render(
|
||||
kv.get_partition_rows_by_bigint_field(TABLE_NAME, "p0", "user_id", None, false, 10)
|
||||
.await?,
|
||||
));
|
||||
observations.push(render(
|
||||
kv.get_partition_rows_by_bigint_field(TABLE_NAME, "p0", "message_id", None, true, 0)
|
||||
.await?,
|
||||
));
|
||||
observations.push(render_sorted(
|
||||
kv.get_partition_rows_by_bigint_field_values(TABLE_NAME, "p0", "message_id", &[10, 30])
|
||||
.await?,
|
||||
));
|
||||
observations.push(render_sorted(
|
||||
kv.get_partition_rows_by_bigint_field_values(TABLE_NAME, "p0", "message_id", &[])
|
||||
.await?,
|
||||
));
|
||||
kv.delete_row(TABLE_NAME, "r2").await?;
|
||||
observations.push(render_sorted(
|
||||
kv.get_partition_rows(TABLE_NAME, "p0").await?,
|
||||
));
|
||||
Ok(observations)
|
||||
}
|
||||
|
||||
struct KvRun {
|
||||
observations: Vec<String>,
|
||||
prepared: Vec<String>,
|
||||
after_session_reset: String,
|
||||
}
|
||||
|
||||
async fn prepared_statement_texts(pool: &Pool) -> anyhow::Result<Vec<String>> {
|
||||
let client = pool.get().await?;
|
||||
let rows = client
|
||||
.query_typed(
|
||||
"SELECT statement FROM pg_prepared_statements ORDER BY statement",
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
Ok(rows
|
||||
.iter()
|
||||
.map(|row| row.get::<_, String>("statement"))
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn discard_session_state(pool: &Pool) -> anyhow::Result<()> {
|
||||
let client = pool.get().await?;
|
||||
client.simple_query("DISCARD ALL").await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_kv_shapes(
|
||||
port: u16,
|
||||
kv_table: &str,
|
||||
prepared_statements: bool,
|
||||
) -> anyhow::Result<KvRun> {
|
||||
let config = postgres_config(port, kv_table, prepared_statements);
|
||||
let pool = connect(&config).await?;
|
||||
let kv = KvClient::new(pool.clone(), &config)?;
|
||||
seed(&kv, kv_table).await?;
|
||||
let observations = exercise_kv_shapes(&kv).await?;
|
||||
let prepared = prepared_statement_texts(&pool).await?;
|
||||
discard_session_state(&pool).await?;
|
||||
let after_session_reset = match kv.get_row(TABLE_NAME, "r1").await {
|
||||
Ok(row) => format!("ok:{row:?}"),
|
||||
Err(err) => format!("err:{err:#}"),
|
||||
};
|
||||
pool.close();
|
||||
Ok(KvRun {
|
||||
observations,
|
||||
prepared,
|
||||
after_session_reset,
|
||||
})
|
||||
}
|
||||
|
||||
async fn wait_for_postgres(container: &str, port: u16) -> anyhow::Result<()> {
|
||||
for _ in 0..180 {
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
if docker(&[
|
||||
"exec",
|
||||
container,
|
||||
"pg_isready",
|
||||
"-U",
|
||||
"fluxer",
|
||||
"-d",
|
||||
"fluxer",
|
||||
])
|
||||
.is_err()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let config = postgres_config(port, KV_TABLE, true);
|
||||
if let Ok(pool) = connect(&config).await {
|
||||
pool.close();
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
anyhow::bail!("postgres never came up")
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn drops_server_side_prepared_statements_when_disabled() -> anyhow::Result<()> {
|
||||
if !docker_available() {
|
||||
eprintln!("skipping: docker is not available");
|
||||
return Ok(());
|
||||
}
|
||||
let port = free_port()?;
|
||||
let container = format!("fluxer-svcstmt-{}-{port}", std::process::id());
|
||||
docker(&[
|
||||
"run",
|
||||
"-d",
|
||||
"--name",
|
||||
&container,
|
||||
"-e",
|
||||
"POSTGRES_USER=fluxer",
|
||||
"-e",
|
||||
"POSTGRES_PASSWORD=fluxer",
|
||||
"-e",
|
||||
"POSTGRES_DB=fluxer",
|
||||
"-p",
|
||||
&format!("127.0.0.1:{port}:5432"),
|
||||
"postgres:16-alpine",
|
||||
"-c",
|
||||
"fsync=off",
|
||||
"-c",
|
||||
"log_statement=all",
|
||||
])?;
|
||||
let result = async {
|
||||
wait_for_postgres(&container, port).await?;
|
||||
let named = run_kv_shapes(port, KV_TABLE, true).await?;
|
||||
let unnamed = run_kv_shapes(port, KV_TABLE_POOLED, false).await?;
|
||||
let named_executions = executed_statements(&container, KV_TABLE)?;
|
||||
let unnamed_executions = executed_statements(&container, KV_TABLE_POOLED)?;
|
||||
anyhow::Ok((named, unnamed, named_executions, unnamed_executions))
|
||||
}
|
||||
.await;
|
||||
let _ = docker(&["rm", "-f", &container]);
|
||||
let (named, unnamed, named_executions, unnamed_executions) = result?;
|
||||
|
||||
assert_eq!(9, named.prepared.len(), "{:#?}", named.prepared);
|
||||
assert!(named.prepared.iter().all(|sql| sql.contains(KV_TABLE)));
|
||||
assert!(
|
||||
named.after_session_reset.contains("prepared statement"),
|
||||
"{}",
|
||||
named.after_session_reset
|
||||
);
|
||||
assert!(
|
||||
named_executions
|
||||
.iter()
|
||||
.any(|line| !line.contains("execute <unnamed>:")),
|
||||
"{named_executions:#?}"
|
||||
);
|
||||
|
||||
assert!(unnamed.prepared.is_empty(), "{:#?}", unnamed.prepared);
|
||||
assert_eq!(
|
||||
format!("ok:{}", unnamed.observations[0]),
|
||||
unnamed.after_session_reset
|
||||
);
|
||||
assert!(!unnamed_executions.is_empty());
|
||||
assert!(
|
||||
unnamed_executions
|
||||
.iter()
|
||||
.all(|line| line.contains("execute <unnamed>:")),
|
||||
"{unnamed_executions:#?}"
|
||||
);
|
||||
|
||||
assert_eq!(named.observations, unnamed.observations);
|
||||
Ok(())
|
||||
}
|
||||
@@ -35,7 +35,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
let postgres_config =
|
||||
fluxer_svc::postgres::PostgresConfig::from_service_config(&config);
|
||||
let pool = fluxer_svc::postgres::connect(&postgres_config).await?;
|
||||
let kv = fluxer_svc::postgres::KvClient::new(pool, &postgres_config.kv_table)?;
|
||||
let kv = fluxer_svc::postgres::KvClient::new(pool, &postgres_config)?;
|
||||
UsersShard::new_postgres(
|
||||
kv,
|
||||
transport.clone(),
|
||||
|
||||
Reference in New Issue
Block a user