use std::sync::Arc;
use serde_json::Value as JsonValue;
use sqlx::PgPool;
use crate::broker::RequestContext;
use crate::runtime::core::{
bind_typed_generic_pg_params, dispatch_param_types, dispatch_params, pg_rows_to_json,
set_request_local_settings, validate_pg_mutation_sql, validate_pg_read_sql,
};
use crate::runtime::executor_utils::{build_probe, json_required_str, with_executor_timeout};
use crate::runtime::executors::{
BackendExecutor, BackendHealth, BackendProbe, MutationExecutor, ObjectExecutor, QueryExecutor,
ResourceAdminExecutor, SearchExecutor,
};
pub(crate) struct PostgresExecutor {
pub(crate) pool: PgPool,
pub(crate) context: Option<Arc<RequestContext>>,
}
impl PostgresExecutor {
pub(crate) fn with_pool(pool: PgPool) -> Self {
Self {
pool,
context: None,
}
}
pub(crate) fn with_context(pool: PgPool, context: Arc<RequestContext>) -> Self {
Self {
pool,
context: Some(context),
}
}
}
impl crate::runtime::backend_context::BackendContextEnforcer for PostgresExecutor {
fn backend_label(&self) -> &str {
"postgres"
}
fn enforce(
&self,
ctx: &crate::runtime::backend_context::AppliedContext,
) -> crate::runtime::backend_context::ContextEffect {
crate::runtime::backend_context::enforce_with_mechanism(
ctx,
"SET LOCAL app.current_* in request-scoped transaction (RLS policies)",
)
}
}
impl BackendHealth for PostgresExecutor {
async fn ping(&self) -> Result<(), String> {
sqlx::query("SELECT 1")
.execute(&self.pool)
.await
.map(|_| ())
.map_err(|e| e.to_string())
}
}
impl QueryExecutor for PostgresExecutor {
async fn query(&self, request_json: &str) -> Result<String, tonic::Status> {
with_executor_timeout("PostgreSQL", "query", async {
let spec: JsonValue = serde_json::from_str(request_json).map_err(|e| {
tonic::Status::invalid_argument(format!("invalid request json: {e}"))
})?;
let sql = json_required_str(&spec, "sql")?;
validate_pg_read_sql(sql)?;
let params = dispatch_params(&spec)?;
let param_types = dispatch_param_types(&spec)?;
let rows = if let Some(ctx) = &self.context {
let mut tx = self.pool.begin().await.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL transaction start failed: {err}"))
})?;
set_request_local_settings(&mut tx, ctx).await?;
let rows = bind_typed_generic_pg_params(
sqlx::query(sql),
¶ms,
param_types.as_deref(),
)?
.fetch_all(&mut *tx)
.await
.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL generic query failed: {err}"))
})?;
tx.commit().await.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL transaction commit failed: {err}"))
})?;
rows
} else {
bind_typed_generic_pg_params(sqlx::query(sql), ¶ms, param_types.as_deref())?
.fetch_all(&self.pool)
.await
.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL generic query failed: {err}"))
})?
};
let rows_json = pg_rows_to_json(rows)?;
serde_json::to_string(&rows_json).map_err(|e| tonic::Status::internal(e.to_string()))
})
.await
}
}
impl MutationExecutor for PostgresExecutor {
async fn mutate(&self, request_json: &str) -> Result<String, tonic::Status> {
with_executor_timeout("PostgreSQL", "mutate", async {
let spec: JsonValue = serde_json::from_str(request_json).map_err(|e| {
tonic::Status::invalid_argument(format!("invalid request json: {e}"))
})?;
let sql = json_required_str(&spec, "sql")?;
validate_pg_mutation_sql(sql)?;
let params = dispatch_params(&spec)?;
let param_types = dispatch_param_types(&spec)?;
let return_rows = spec
.get("return_rows")
.and_then(JsonValue::as_bool)
.unwrap_or_else(|| sql.to_ascii_lowercase().contains(" returning "));
if let Some(ctx) = &self.context {
let mut tx = self.pool.begin().await.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL transaction start failed: {err}"))
})?;
set_request_local_settings(&mut tx, ctx).await?;
let result = if return_rows {
let rows = bind_typed_generic_pg_params(
sqlx::query(sql),
¶ms,
param_types.as_deref(),
)?
.fetch_all(&mut *tx)
.await
.map_err(|err| {
tonic::Status::internal(format!(
"PostgreSQL generic mutation failed: {err}"
))
})?;
let rows_json = pg_rows_to_json(rows)?;
serde_json::json!({
"affected_rows": rows_json.len(),
"rows": rows_json
})
.to_string()
} else {
let result = bind_typed_generic_pg_params(
sqlx::query(sql),
¶ms,
param_types.as_deref(),
)?
.execute(&mut *tx)
.await
.map_err(|err| {
tonic::Status::internal(format!(
"PostgreSQL generic mutation failed: {err}"
))
})?;
serde_json::json!({ "affected_rows": result.rows_affected() }).to_string()
};
tx.commit().await.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL transaction commit failed: {err}"))
})?;
return Ok(result);
}
if return_rows {
let rows = bind_typed_generic_pg_params(
sqlx::query(sql),
¶ms,
param_types.as_deref(),
)?
.fetch_all(&self.pool)
.await
.map_err(|err| {
tonic::Status::internal(format!("PostgreSQL generic mutation failed: {err}"))
})?;
let rows_json = pg_rows_to_json(rows)?;
return Ok(serde_json::json!({
"affected_rows": rows_json.len(),
"rows": rows_json
})
.to_string());
}
let result =
bind_typed_generic_pg_params(sqlx::query(sql), ¶ms, param_types.as_deref())?
.execute(&self.pool)
.await
.map_err(|err| {
tonic::Status::internal(format!(
"PostgreSQL generic mutation failed: {err}"
))
})?;
Ok(serde_json::json!({ "affected_rows": result.rows_affected() }).to_string())
})
.await
}
}
impl SearchExecutor for PostgresExecutor {
async fn search(&self, _request_json: &str) -> Result<String, tonic::Status> {
Err(tonic::Status::failed_precondition(
"postgres does not support vector search; use query",
))
}
}
impl ObjectExecutor for PostgresExecutor {
async fn get_object(&self, _request_json: &str) -> Result<Vec<u8>, tonic::Status> {
Err(tonic::Status::failed_precondition(
"postgres is not an object store",
))
}
async fn put_object(
&self,
_request_json: &str,
_bytes: Vec<u8>,
) -> Result<String, tonic::Status> {
Err(tonic::Status::failed_precondition(
"postgres is not an object store",
))
}
}
impl ResourceAdminExecutor for PostgresExecutor {
async fn ensure_resource(
&self,
_resource_name: &str,
_spec_json: &str,
) -> Result<(), tonic::Status> {
Err(tonic::Status::failed_precondition(
"postgres resource lifecycle is managed via catalog migrations",
))
}
async fn drop_resource(&self, _resource_name: &str) -> Result<(), tonic::Status> {
Err(tonic::Status::failed_precondition(
"postgres resource lifecycle is managed via catalog migrations",
))
}
async fn list_resources(&self) -> Result<Vec<String>, tonic::Status> {
Err(tonic::Status::failed_precondition(
"postgres resource lifecycle is managed via catalog migrations",
))
}
}
impl BackendExecutor for PostgresExecutor {
async fn transaction(&self, _request_json: &str) -> Result<String, tonic::Status> {
Err(tonic::Status::failed_precondition(
"use the typed transaction RPC for relational transactions",
))
}
async fn probe(&self) -> Result<BackendProbe, tonic::Status> {
Ok(build_probe(
"postgres",
<Self as BackendHealth>::ping(self).await,
))
}
}