use crate::server::type_map::{encode_value, ent_to_pg_type};
use async_trait::async_trait;
use entdb::catalog::Catalog;
use entdb::error::EntDbError;
use entdb::query::binder::Binder;
use entdb::query::executor::{build_executor, ExecutionContext, TxExecutionContext};
use entdb::query::history::OptimizerHistoryRecord;
use entdb::query::optimizer::Optimizer;
use entdb::query::planner::Planner;
use entdb::query::polyglot::{transpile_with_meta, PolyglotOptions};
use entdb::storage::table::Table;
use entdb::tx::TransactionHandle;
use futures::{stream, Sink};
use parking_lot::Mutex;
use pgwire::api::portal::{Format, Portal};
use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
use pgwire::api::results::{
DataRowEncoder, DescribePortalResponse, DescribeStatementResponse, FieldFormat, FieldInfo,
QueryResponse, Response, Tag,
};
use pgwire::api::stmt::{NoopQueryParser, StoredStatement};
use pgwire::api::{ClientInfo, Type};
use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
use pgwire::messages::PgWireBackendMessage;
use sqlparser::ast::{visit_expressions_mut, Expr, Statement, Value as SqlValue};
use sqlparser::dialect::PostgreSqlDialect;
use sqlparser::parser::Parser;
use std::fmt::Debug;
use std::ops::ControlFlow;
use std::sync::Arc;
use tracing::{debug, info_span};
use super::metrics::ServerMetrics;
use super::{optimizer_history_schema_hash, Database};
pub struct EntHandler {
db: Arc<Database>,
current_txn: Mutex<Option<TransactionHandle>>,
query_parser: Arc<NoopQueryParser>,
max_statement_bytes: usize,
query_timeout_ms: u64,
metrics: Arc<ServerMetrics>,
polyglot: PolyglotOptions,
await_durable: bool,
}
impl EntHandler {
pub fn new(
db: Arc<Database>,
max_statement_bytes: usize,
query_timeout_ms: u64,
metrics: Arc<ServerMetrics>,
await_durable: bool,
) -> Self {
Self {
db,
current_txn: Mutex::new(None),
query_parser: Arc::new(NoopQueryParser::new()),
max_statement_bytes,
query_timeout_ms,
metrics,
await_durable,
polyglot: PolyglotOptions {
enabled: std::env::var("ENTDB_POLYGLOT")
.map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
.unwrap_or(false),
},
}
}
fn execute_sql(
&self,
sql: &str,
max_rows: Option<usize>,
) -> PgWireResult<Vec<Response<'static>>> {
if sql.len() > self.max_statement_bytes {
return Err(user_error(
"54000",
format!(
"statement exceeds configured max_statement_bytes ({} > {})",
sql.len(),
self.max_statement_bytes
),
));
}
let _span = info_span!(
"pgwire.execute_sql",
sql_len = sql.len(),
max_rows = max_rows.unwrap_or(0)
)
.entered();
let statements = parse_sql_with_polyglot(sql, self.polyglot)?;
let mut responses = Vec::new();
let started = std::time::Instant::now();
for stmt in &statements {
if started.elapsed().as_millis() as u64 > self.query_timeout_ms {
return Err(user_error("57014", "query timeout exceeded"));
}
debug!(statement = %tag_for_statement_name(stmt), "processing statement");
match stmt {
Statement::StartTransaction { .. } => {
let mut guard = self.current_txn.lock();
if guard.is_some() {
return Err(user_error(
"25001",
"transaction already active for this session",
));
}
*guard = Some(self.db.txn_manager.begin());
responses.push(Response::TransactionStart(Tag::new("BEGIN")));
}
Statement::Commit { .. } => {
let tx = {
let mut guard = self.current_txn.lock();
guard
.take()
.ok_or_else(|| user_error("25000", "no active transaction to COMMIT"))?
};
self.db
.txn_manager
.commit_with_options(tx.txn_id, Some(self.await_durable))
.map_err(map_entdb_error)?;
responses.push(Response::TransactionEnd(Tag::new("COMMIT")));
}
Statement::Rollback { .. } => {
let tx = {
let mut guard = self.current_txn.lock();
guard.take().ok_or_else(|| {
user_error("25000", "no active transaction to ROLLBACK")
})?
};
self.db.txn_manager.abort(tx.txn_id);
responses.push(Response::TransactionEnd(Tag::new("ROLLBACK")));
}
_ => {
let active_tx = *self.current_txn.lock();
let exec_resp = if let Some(tx) = active_tx {
self.execute_statement_in_txn(&tx, stmt, max_rows, None)?
} else {
let tx = self.db.txn_manager.begin();
match self.execute_statement_in_txn(&tx, stmt, max_rows, None) {
Ok(resp) => {
if let Err(e) = self
.db
.txn_manager
.commit_with_options(tx.txn_id, Some(self.await_durable))
{
self.db.txn_manager.abort(tx.txn_id);
return Err(map_entdb_error(e));
}
resp
}
Err(e) => {
self.db.txn_manager.abort(tx.txn_id);
return Err(e);
}
}
};
responses.push(exec_resp);
}
}
}
Ok(responses)
}
fn execute_statement_in_txn(
&self,
tx: &TransactionHandle,
stmt: &Statement,
max_rows: Option<usize>,
row_format: Option<&Format>,
) -> PgWireResult<Response<'static>> {
let stmt_started = std::time::Instant::now();
let _span = info_span!(
"pgwire.execute_statement",
statement = %tag_for_statement_name(stmt),
txn_id = tx.txn_id
)
.entered();
let binder = Binder::new(Arc::clone(&self.db.catalog));
let planner = Planner;
let bound = match binder.bind(stmt) {
Ok(b) => b,
Err(e) => {
self.record_optimizer_history_error(stmt, "bind_error", stmt_started.elapsed());
return Err(map_entdb_error(e));
}
};
let fingerprint = Optimizer::fingerprint_bound_statement(&bound);
let history = self.db.optimizer_history.read_for_fingerprint(&fingerprint);
self.ensure_statement_timeout(stmt_started)?;
let plan = match planner.plan(bound) {
Ok(p) => p,
Err(e) => {
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint,
plan_signature: "error".to_string(),
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: "{\"root\":0}".to_string(),
latency_ms: stmt_started.elapsed().as_millis() as u64,
memory_peak_bytes: 0,
success: false,
error_class: Some("plan_error".to_string()),
confidence: 0.0,
});
return Err(map_entdb_error(e));
}
};
self.ensure_statement_timeout(stmt_started)?;
let optimized_outcome = Optimizer::optimize_with_trace_and_history(
plan,
&fingerprint,
self.db.optimizer_config,
&history,
);
let chosen_plan_signature = optimized_outcome
.trace
.chosen_plan_signature
.clone()
.unwrap_or_else(|| "baseline".to_string());
let optimized = optimized_outcome.plan;
self.ensure_statement_timeout(stmt_started)?;
let ctx = ExecutionContext {
catalog: Arc::clone(&self.db.catalog),
tx: Some(TxExecutionContext {
txn_id: tx.txn_id,
snapshot_ts: tx.snapshot_ts,
txn_manager: Arc::clone(&self.db.txn_manager),
}),
};
let mut exec = match build_executor(&optimized, &ctx) {
Ok(exec) => exec,
Err(e) => {
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint,
plan_signature: chosen_plan_signature,
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: "{\"root\":0}".to_string(),
latency_ms: stmt_started.elapsed().as_millis() as u64,
memory_peak_bytes: 0,
success: false,
error_class: Some("executor_build_error".to_string()),
confidence: 0.0,
});
return Err(map_entdb_error(e));
}
};
self.ensure_statement_timeout(stmt_started)?;
if let Err(e) = exec.open() {
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint,
plan_signature: chosen_plan_signature,
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: "{\"root\":0}".to_string(),
latency_ms: stmt_started.elapsed().as_millis() as u64,
memory_peak_bytes: 0,
success: false,
error_class: Some("executor_open_error".to_string()),
confidence: 0.0,
});
return Err(map_entdb_error(e));
}
self.ensure_statement_timeout(stmt_started)?;
let schema = exec.schema().clone();
let mut rows = Vec::new();
while let Some(row) = match exec.next() {
Ok(r) => r,
Err(e) => {
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint: fingerprint.clone(),
plan_signature: chosen_plan_signature.clone(),
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: format!("{{\"root\":{}}}", rows.len()),
latency_ms: stmt_started.elapsed().as_millis() as u64,
memory_peak_bytes: 0,
success: false,
error_class: Some("executor_next_error".to_string()),
confidence: 0.0,
});
return Err(map_entdb_error(e));
}
} {
self.ensure_statement_timeout(stmt_started)?;
rows.push(row);
if let Some(limit) = max_rows {
if limit > 0 && rows.len() >= limit {
break;
}
}
}
if let Err(e) = exec.close() {
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint,
plan_signature: chosen_plan_signature,
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: format!("{{\"root\":{}}}", rows.len()),
latency_ms: stmt_started.elapsed().as_millis() as u64,
memory_peak_bytes: 0,
success: false,
error_class: Some("executor_close_error".to_string()),
confidence: 0.0,
});
return Err(map_entdb_error(e));
}
let observed_rows = if schema.columns.is_empty() {
exec.affected_rows().unwrap_or(0)
} else {
rows.len() as u64
};
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint,
plan_signature: chosen_plan_signature,
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: format!("{{\"root\":{observed_rows}}}"),
latency_ms: stmt_started.elapsed().as_millis() as u64,
memory_peak_bytes: 0,
success: true,
error_class: None,
confidence: 1.0,
});
if schema.columns.is_empty() {
let affected = exec.affected_rows().unwrap_or(0) as usize;
let tag = execution_tag_for_statement(stmt, affected);
return Ok(Response::Execution(tag));
}
let fields = schema
.columns
.iter()
.enumerate()
.map(|(idx, c)| {
FieldInfo::new(
c.name.clone(),
None,
None,
ent_to_pg_type(&c.data_type),
row_format
.map(|f| f.format_for(idx))
.unwrap_or(FieldFormat::Text),
)
})
.collect::<Vec<_>>();
let schema = Arc::new(fields);
let mut encoded = Vec::with_capacity(rows.len());
for row in &rows {
let mut encoder = DataRowEncoder::new(Arc::clone(&schema));
for (idx, value) in row.iter().enumerate() {
let field = &schema[idx];
encode_value(&mut encoder, value, field.datatype(), field.format())?;
}
encoded.push(encoder.finish()?);
}
let data_row_stream = stream::iter(encoded.into_iter().map(Ok));
let mut query = QueryResponse::new(schema, data_row_stream);
query.set_command_tag(tag_for_statement_name(stmt));
Ok(Response::Query(query))
}
fn describe_statement_from_stmt(
&self,
stmt: &Statement,
format: &Format,
) -> PgWireResult<Vec<FieldInfo>> {
if !statement_returns_rows(stmt) {
return Ok(Vec::new());
}
let tx = self.db.txn_manager.begin();
let binder = Binder::new(Arc::clone(&self.db.catalog));
let planner = Planner;
let bound = binder.bind(stmt).map_err(map_entdb_error)?;
let fingerprint = Optimizer::fingerprint_bound_statement(&bound);
let history = self.db.optimizer_history.read_for_fingerprint(&fingerprint);
let plan = planner.plan(bound).map_err(map_entdb_error)?;
let optimized = Optimizer::optimize_with_trace_and_history(
plan,
&fingerprint,
self.db.optimizer_config,
&history,
)
.plan;
let ctx = ExecutionContext {
catalog: Arc::clone(&self.db.catalog),
tx: Some(TxExecutionContext {
txn_id: tx.txn_id,
snapshot_ts: tx.snapshot_ts,
txn_manager: Arc::clone(&self.db.txn_manager),
}),
};
let exec = build_executor(&optimized, &ctx).map_err(map_entdb_error)?;
let fields = exec
.schema()
.columns
.iter()
.enumerate()
.map(|(idx, col)| {
FieldInfo::new(
col.name.clone(),
None,
None,
ent_to_pg_type(&col.data_type),
format.format_for(idx),
)
})
.collect::<Vec<_>>();
self.db.txn_manager.abort(tx.txn_id);
Ok(fields)
}
fn describe_statement_inner(&self, sql: &str, format: &Format) -> PgWireResult<Vec<FieldInfo>> {
let stmt = parse_single_statement_with_polyglot(
sql,
"describe supports a single statement",
self.polyglot,
)?;
self.describe_statement_from_stmt(&stmt, format)
}
}
impl EntHandler {
fn record_optimizer_history(&self, entry: OptimizerHistoryRecord) {
self.db.optimizer_history.try_record(entry);
}
fn record_optimizer_history_error(
&self,
stmt: &Statement,
class: &str,
elapsed: std::time::Duration,
) {
self.record_optimizer_history(OptimizerHistoryRecord {
fingerprint: format!("stmt:{}", tag_for_statement_name(stmt)),
plan_signature: "error".to_string(),
schema_hash: optimizer_history_schema_hash().to_string(),
captured_at_ms: now_epoch_millis(),
rowcount_observed_json: "{\"root\":0}".to_string(),
latency_ms: elapsed.as_millis() as u64,
memory_peak_bytes: 0,
success: false,
error_class: Some(class.to_string()),
confidence: 0.0,
});
}
}
fn now_epoch_millis() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
#[async_trait]
impl SimpleQueryHandler for EntHandler {
async fn do_query<'a, C>(
&self,
_client: &mut C,
query: &'a str,
) -> PgWireResult<Vec<Response<'a>>>
where
C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
C::Error: Debug,
PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
{
let started = std::time::Instant::now();
let result = self.execute_sql(query, None);
self.record_query_metric(started, result.as_ref().err());
result
}
}
#[async_trait]
impl ExtendedQueryHandler for EntHandler {
type Statement = String;
type QueryParser = NoopQueryParser;
fn query_parser(&self) -> Arc<Self::QueryParser> {
Arc::clone(&self.query_parser)
}
async fn do_query<'a, C>(
&self,
_client: &mut C,
portal: &'a Portal<Self::Statement>,
max_rows: usize,
) -> PgWireResult<Response<'a>>
where
C: ClientInfo + Unpin + Send + Sync,
{
let started = std::time::Instant::now();
let result = (|| -> PgWireResult<Response<'a>> {
let stmt = statement_from_portal(portal, self.polyglot)?;
if matches!(
stmt,
Statement::StartTransaction { .. }
| Statement::Commit { .. }
| Statement::Rollback { .. }
) {
let mut responses = self.execute_sql(&stmt.to_string(), Some(max_rows))?;
if responses.is_empty() {
return Ok(Response::EmptyQuery);
}
return Ok(responses.remove(0));
}
let active_tx = *self.current_txn.lock();
if let Some(tx) = active_tx {
return self.execute_statement_in_txn(
&tx,
&stmt,
Some(max_rows),
Some(&portal.result_column_format),
);
}
let tx = self.db.txn_manager.begin();
match self.execute_statement_in_txn(
&tx,
&stmt,
Some(max_rows),
Some(&portal.result_column_format),
) {
Ok(resp) => {
if let Err(e) = self
.db
.txn_manager
.commit_with_options(tx.txn_id, Some(self.await_durable))
{
self.db.txn_manager.abort(tx.txn_id);
Err(map_entdb_error(e))
} else {
Ok(resp)
}
}
Err(e) => {
self.db.txn_manager.abort(tx.txn_id);
Err(e)
}
}
})();
self.record_query_metric(started, result.as_ref().err());
result
}
async fn do_describe_statement<C>(
&self,
_client: &mut C,
stmt: &StoredStatement<Self::Statement>,
) -> PgWireResult<DescribeStatementResponse>
where
C: ClientInfo + Unpin + Send + Sync,
{
let mut parameter_types = stmt.parameter_types.clone();
let inferred_count = max_placeholder_index(&stmt.statement);
if parameter_types.len() < inferred_count {
parameter_types.resize(inferred_count, Type::INT4);
}
for ty in &mut parameter_types {
if *ty == Type::UNKNOWN {
*ty = Type::INT4;
}
}
let fields = self
.describe_statement_inner(&stmt.statement, &Format::UnifiedBinary)
.or_else(|_| {
let normalized = normalize_placeholders_for_describe(&stmt.statement);
self.describe_statement_inner(&normalized, &Format::UnifiedBinary)
})
.unwrap_or_default();
Ok(DescribeStatementResponse::new(parameter_types, fields))
}
async fn do_describe_portal<C>(
&self,
_client: &mut C,
portal: &Portal<Self::Statement>,
) -> PgWireResult<DescribePortalResponse>
where
C: ClientInfo + Unpin + Send + Sync,
{
let stmt = statement_from_portal(portal, self.polyglot)?;
let fields = self.describe_statement_from_stmt(&stmt, &portal.result_column_format)?;
Ok(DescribePortalResponse::new(fields))
}
}
#[cfg(test)]
fn parse_sql(sql: &str) -> PgWireResult<Vec<Statement>> {
let dialect = PostgreSqlDialect {};
Parser::parse_sql(&dialect, sql).map_err(|e| user_error("42601", format!("parse error: {e}")))
}
fn parse_sql_with_polyglot(sql: &str, opts: PolyglotOptions) -> PgWireResult<Vec<Statement>> {
let transpiled = transpile_with_meta(sql, opts).map_err(map_entdb_error)?;
let dialect = PostgreSqlDialect {};
Parser::parse_sql(&dialect, &transpiled.transpiled_sql).map_err(|e| {
if transpiled.changed {
user_error(
"42601",
format!(
"parse error: {e}; original_sql={:?}; transpiled_sql={:?}",
transpiled.original_sql, transpiled.transpiled_sql
),
)
} else {
user_error("42601", format!("parse error: {e}"))
}
})
}
impl EntHandler {
fn ensure_statement_timeout(&self, started: std::time::Instant) -> PgWireResult<()> {
if started.elapsed().as_millis() as u64 > self.query_timeout_ms {
return Err(user_error("57014", "query timeout exceeded"));
}
Ok(())
}
fn record_query_metric(&self, started: std::time::Instant, err: Option<&PgWireError>) {
let elapsed = started.elapsed().as_nanos() as u64;
let sqlstate = err.and_then(sqlstate_from_error);
self.metrics.on_query_finished(elapsed, sqlstate);
}
}
fn sqlstate_from_error(err: &PgWireError) -> Option<&str> {
match err {
PgWireError::UserError(info) => Some(info.code.as_str()),
_ => Some("XX000"),
}
}
#[cfg(test)]
fn parse_single_statement(sql: &str, message: &'static str) -> PgWireResult<Statement> {
let mut statements = parse_sql(sql)?;
if statements.len() != 1 {
return Err(user_error("0A000", message));
}
Ok(statements.remove(0))
}
fn parse_single_statement_with_polyglot(
sql: &str,
message: &'static str,
opts: PolyglotOptions,
) -> PgWireResult<Statement> {
let mut statements = parse_sql_with_polyglot(sql, opts)?;
if statements.len() != 1 {
return Err(user_error("0A000", message));
}
Ok(statements.remove(0))
}
fn statement_from_portal(
portal: &Portal<String>,
opts: PolyglotOptions,
) -> PgWireResult<Statement> {
let mut stmt = parse_single_statement_with_polyglot(
&portal.statement.statement,
"extended protocol supports a single statement per execute",
opts,
)?;
if portal.parameter_len() == 0 {
return Ok(stmt);
}
let values = collect_portal_values(portal)?;
bind_statement_placeholders(&mut stmt, &values)?;
Ok(stmt)
}
fn collect_portal_values(portal: &Portal<String>) -> PgWireResult<Vec<SqlValue>> {
let mut values = Vec::with_capacity(portal.parameter_len());
for i in 0..portal.parameter_len() {
let pg_type = portal
.statement
.parameter_types
.get(i)
.unwrap_or(&Type::UNKNOWN);
values.push(parameter_sql_value(portal, i, pg_type)?);
}
Ok(values)
}
fn bind_statement_placeholders(stmt: &mut Statement, values: &[SqlValue]) -> PgWireResult<()> {
let mut question_mark_idx = 0usize;
let outcome = visit_expressions_mut(stmt, |expr| {
if let Expr::Value(v) = expr {
if let SqlValue::Placeholder(placeholder) = &v.value {
let idx = match placeholder_index(placeholder, values.len(), &mut question_mark_idx)
{
Ok(idx) => idx,
Err(err) => return ControlFlow::Break(err),
};
v.value = values[idx - 1].clone();
}
}
ControlFlow::Continue(())
});
match outcome {
ControlFlow::Continue(()) => Ok(()),
ControlFlow::Break(err) => Err(err),
}
}
fn placeholder_index(
placeholder: &str,
total: usize,
question_mark_idx: &mut usize,
) -> PgWireResult<usize> {
let idx = if placeholder == "?" {
*question_mark_idx += 1;
*question_mark_idx
} else if let Some(raw) = placeholder.strip_prefix('$') {
raw.parse::<usize>()
.map_err(|_| user_error("22023", format!("invalid placeholder '{placeholder}'")))?
} else {
return Err(user_error(
"22023",
format!("unsupported placeholder syntax '{placeholder}'"),
));
};
if idx == 0 || idx > total {
return Err(user_error(
"22023",
format!("placeholder '{placeholder}' is out of range"),
));
}
Ok(idx)
}
fn parameter_sql_value(
portal: &Portal<String>,
idx: usize,
pg_type: &Type,
) -> PgWireResult<SqlValue> {
if *pg_type == Type::UNKNOWN {
return parameter_sql_value_unknown(portal, idx);
}
if *pg_type == Type::BOOL {
return Ok(match portal.parameter::<bool>(idx, pg_type)? {
Some(v) => SqlValue::Boolean(v),
None => SqlValue::Null,
});
}
if *pg_type == Type::INT2 {
return Ok(match portal.parameter::<i16>(idx, pg_type)? {
Some(v) => SqlValue::Number(v.to_string(), false),
None => SqlValue::Null,
});
}
if *pg_type == Type::INT4 {
return Ok(match portal.parameter::<i32>(idx, pg_type)? {
Some(v) => SqlValue::Number(v.to_string(), false),
None => SqlValue::Null,
});
}
if *pg_type == Type::INT8 {
return Ok(match portal.parameter::<i64>(idx, pg_type)? {
Some(v) => SqlValue::Number(v.to_string(), false),
None => SqlValue::Null,
});
}
if *pg_type == Type::FLOAT4 {
return Ok(match portal.parameter::<f32>(idx, pg_type)? {
Some(v) => SqlValue::Number((v as f64).to_string(), false),
None => SqlValue::Null,
});
}
if *pg_type == Type::FLOAT8 {
return Ok(match portal.parameter::<f64>(idx, pg_type)? {
Some(v) => SqlValue::Number(v.to_string(), false),
None => SqlValue::Null,
});
}
Ok(match portal.parameter::<String>(idx, pg_type)? {
Some(v) => SqlValue::SingleQuotedString(v),
None => SqlValue::Null,
})
}
fn parameter_sql_value_unknown(portal: &Portal<String>, idx: usize) -> PgWireResult<SqlValue> {
let Some(raw) = portal.parameters.get(idx) else {
return Ok(SqlValue::Null);
};
let Some(bytes) = raw else {
return Ok(SqlValue::Null);
};
if portal.parameter_format.is_binary(idx) {
match bytes.len() {
1 => {
let b = bytes[0];
if b == 0 || b == 1 {
return Ok(SqlValue::Boolean(b == 1));
}
}
4 => {
let v = i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
return Ok(SqlValue::Number(v.to_string(), false));
}
8 => {
let v = i64::from_be_bytes([
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
]);
return Ok(SqlValue::Number(v.to_string(), false));
}
_ => {}
}
}
match std::str::from_utf8(bytes) {
Ok(s) => {
if let Ok(i) = s.parse::<i64>() {
return Ok(SqlValue::Number(i.to_string(), false));
}
if let Ok(f) = s.parse::<f64>() {
return Ok(SqlValue::Number(f.to_string(), false));
}
if s.eq_ignore_ascii_case("true") || s.eq_ignore_ascii_case("false") {
return Ok(SqlValue::Boolean(s.eq_ignore_ascii_case("true")));
}
Ok(SqlValue::SingleQuotedString(s.to_string()))
}
Err(_) => Ok(SqlValue::SingleQuotedString(
String::from_utf8_lossy(bytes).into_owned(),
)),
}
}
fn tag_for_statement(stmt: &Statement) -> Tag {
Tag::new(tag_for_statement_name(stmt))
}
fn statement_returns_rows(stmt: &Statement) -> bool {
match stmt {
Statement::Query(_) => true,
Statement::Insert(insert) => insert.returning.is_some(),
Statement::Update { returning, .. } => returning.is_some(),
Statement::Delete(delete) => delete.returning.is_some(),
_ => false,
}
}
fn execution_tag_for_statement(stmt: &Statement, affected: usize) -> Tag {
match stmt {
Statement::Insert(_) => Tag::new("INSERT 0").with_rows(affected),
_ => tag_for_statement(stmt).with_rows(affected),
}
}
fn tag_for_statement_name(stmt: &Statement) -> &'static str {
match stmt {
Statement::Query(_) => "SELECT",
Statement::Insert(_) => "INSERT",
Statement::Update { .. } => "UPDATE",
Statement::Delete { .. } => "DELETE",
Statement::Truncate { .. } => "TRUNCATE",
Statement::CreateTable { .. } => "CREATE TABLE",
Statement::CreateIndex(_) => "CREATE INDEX",
Statement::AlterTable { .. } => "ALTER TABLE",
Statement::Drop {
object_type: sqlparser::ast::ObjectType::Table,
..
} => "DROP TABLE",
Statement::Drop {
object_type: sqlparser::ast::ObjectType::Index,
..
} => "DROP INDEX",
Statement::StartTransaction { .. } => "BEGIN",
Statement::Commit { .. } => "COMMIT",
Statement::Rollback { .. } => "ROLLBACK",
_ => "OK",
}
}
fn map_entdb_error(err: EntDbError) -> PgWireError {
let (code, message) = match err {
EntDbError::Query(m) => (query_sqlstate(&m), m),
EntDbError::BufferPoolFull => ("53200", "buffer pool full".to_string()),
EntDbError::PagePinned(pid) => ("55006", format!("page {pid} is pinned")),
EntDbError::PageNotFound(pid) => ("XX000", format!("page {pid} not found")),
EntDbError::Io(e) => ("58000", format!("io error: {e}")),
EntDbError::Wal(m) => ("XX000", format!("wal error: {m}")),
EntDbError::Corruption(m) => ("XX001", format!("corruption: {m}")),
EntDbError::InvalidPage(m) => ("XX000", format!("invalid page: {m}")),
EntDbError::PageAlreadyPresent(pid) => ("XX000", format!("page {pid} already present")),
};
user_error(code, message)
}
fn query_sqlstate(message: &str) -> &'static str {
let lower = message.to_ascii_lowercase();
if lower.contains("parse error") {
return "42601";
}
if lower.contains("already active") {
return "25001";
}
if lower.contains("no active transaction") {
return "25000";
}
if lower.contains("write-write conflict") {
return "40001";
}
if lower.contains("does not exist") {
if lower.contains("table") {
return "42P01";
}
if lower.contains("column") {
return "42703";
}
}
if lower.contains("cannot cast value")
|| lower.contains("invalid")
|| lower.contains("unsupported")
|| lower.contains("must be numeric")
{
return "22000";
}
"XX000"
}
fn user_error(code: impl Into<String>, message: impl Into<String>) -> PgWireError {
PgWireError::UserError(Box::new(ErrorInfo::new(
"ERROR".to_string(),
code.into(),
message.into(),
)))
}
fn max_placeholder_index(sql: &str) -> usize {
let bytes = sql.as_bytes();
let mut max_idx = 0usize;
let mut i = 0usize;
while i < bytes.len() {
if bytes[i] == b'$' {
let start = i + 1;
let mut j = start;
while j < bytes.len() && bytes[j].is_ascii_digit() {
j += 1;
}
if j > start {
if let Ok(n) = sql[start..j].parse::<usize>() {
max_idx = max_idx.max(n);
}
i = j;
continue;
}
}
i += 1;
}
max_idx
}
fn normalize_placeholders_for_describe(sql: &str) -> String {
let bytes = sql.as_bytes();
let mut out = String::with_capacity(sql.len());
let mut i = 0usize;
while i < bytes.len() {
if bytes[i] == b'$' {
let start = i + 1;
let mut j = start;
while j < bytes.len() && bytes[j].is_ascii_digit() {
j += 1;
}
if j > start {
out.push('0');
i = j;
continue;
}
}
out.push(bytes[i] as char);
i += 1;
}
out
}
pub fn scan_max_txn_id_from_storage(catalog: &Catalog) -> Result<u64, EntDbError> {
let mut max_txn = 0_u64;
for table in catalog.list_tables() {
let t = Table::open(table.table_id, table.first_page_id, catalog.buffer_pool());
for (_, tuple) in t.scan() {
let decoded = entdb::query::executor::decode_stored_row(&tuple.data)?;
if let entdb::query::executor::DecodedRow::Versioned(v) = decoded {
max_txn = max_txn.max(v.created_txn);
if let Some(d) = v.deleted_txn {
max_txn = max_txn.max(d);
}
}
}
}
Ok(max_txn)
}
#[cfg(test)]
mod tests {
use super::{
bind_statement_placeholders, execution_tag_for_statement, max_placeholder_index,
normalize_placeholders_for_describe, parse_single_statement, parse_sql, EntHandler,
};
use crate::server::Database;
use entdb::catalog::Catalog;
use entdb::query::history::OptimizerHistoryRecorder;
use entdb::query::optimizer::OptimizerConfig;
use entdb::storage::buffer_pool::BufferPool;
use entdb::storage::disk_manager::DiskManager;
use entdb::tx::TransactionManager;
use entdb::wal::log_manager::LogManager;
use pgwire::api::results::Tag;
use pgwire::error::PgWireError;
use proptest::prelude::*;
use sqlparser::ast::Value as SqlValue;
use std::sync::Arc;
use tempfile::tempdir;
fn test_handler() -> EntHandler {
let dir = tempdir().expect("tempdir");
let db_path = dir.path().join("handler.db");
let wal_path = dir.path().join("handler.wal");
let dm = Arc::new(DiskManager::new(&db_path).expect("disk"));
let lm = Arc::new(LogManager::new(&wal_path, 4096).expect("wal"));
let bp = Arc::new(BufferPool::with_log_manager(
32,
Arc::clone(&dm),
Arc::clone(&lm),
));
let catalog = Arc::new(Catalog::load(Arc::clone(&bp)).expect("catalog"));
let txn = Arc::new(TransactionManager::new());
let history_path = dir.path().join("handler.optimizer_history.json");
let optimizer_history = Arc::new(
OptimizerHistoryRecorder::new(
&history_path,
super::optimizer_history_schema_hash(),
16,
128,
)
.expect("optimizer history"),
);
let db = Arc::new(Database {
disk_manager: dm,
log_manager: lm,
buffer_pool: bp,
catalog,
txn_manager: txn,
optimizer_history,
optimizer_config: OptimizerConfig::default(),
});
EntHandler::new(
db,
1024 * 1024,
30_000,
Arc::new(super::ServerMetrics::default()),
false,
)
}
#[test]
fn bind_statement_placeholders_replaces_parameter_nodes() {
let mut stmt = parse_single_statement(
"SELECT id FROM t WHERE v > $1 ORDER BY id LIMIT $2",
"single",
)
.expect("parse");
bind_statement_placeholders(
&mut stmt,
&[
SqlValue::Number("7".to_string(), false),
SqlValue::Number("3".to_string(), false),
],
)
.expect("bind placeholders");
let rendered = stmt.to_string();
assert!(rendered.contains("v > 7"), "rendered SQL: {rendered}");
assert!(rendered.contains("LIMIT 3"), "rendered SQL: {rendered}");
}
#[test]
fn bind_statement_placeholders_rejects_out_of_range_index() {
let mut stmt = parse_single_statement("SELECT $3", "single").expect("parse");
let err = match bind_statement_placeholders(
&mut stmt,
&[SqlValue::Number("1".to_string(), false)],
) {
Ok(_) => panic!("expected out-of-range placeholder failure"),
Err(e) => e,
};
match err {
PgWireError::UserError(info) => {
assert_eq!(info.code, "22023");
}
other => panic!("unexpected error: {other:?}"),
}
}
#[test]
fn parse_single_statement_rejects_multi_statement_sql() {
let err = match parse_single_statement("SELECT 1; SELECT 2", "single") {
Ok(_) => panic!("expected single-statement rejection"),
Err(e) => e,
};
match err {
PgWireError::UserError(info) => {
assert_eq!(info.code, "0A000");
}
other => panic!("unexpected error: {other:?}"),
}
}
#[test]
fn max_placeholder_index_finds_highest_parameter() {
assert_eq!(max_placeholder_index("SELECT $1, $2, $10"), 10);
assert_eq!(max_placeholder_index("SELECT 1"), 0);
}
#[test]
fn normalize_placeholders_for_describe_rewrites_all_numbered_params() {
let sql = "SELECT id FROM t WHERE v > $1 ORDER BY id LIMIT $2";
let normalized = normalize_placeholders_for_describe(sql);
assert_eq!(
normalized,
"SELECT id FROM t WHERE v > 0 ORDER BY id LIMIT 0"
);
}
#[test]
fn execution_tag_for_insert_matches_postgres_format() {
let stmt = parse_single_statement("INSERT INTO t VALUES (1)", "single").expect("parse");
let tag = execution_tag_for_statement(&stmt, 2);
assert_eq!(tag, Tag::new("INSERT 0").with_rows(2));
}
#[test]
fn execution_tag_for_update_and_delete_include_row_count() {
let update = parse_single_statement("UPDATE t SET v = 1", "single").expect("parse");
let delete = parse_single_statement("DELETE FROM t", "single").expect("parse");
assert_eq!(
execution_tag_for_statement(&update, 3),
Tag::new("UPDATE").with_rows(3)
);
assert_eq!(
execution_tag_for_statement(&delete, 4),
Tag::new("DELETE").with_rows(4)
);
}
#[test]
fn handler_begin_commit_round_trip() {
let handler = test_handler();
let out = handler
.execute_sql("BEGIN; COMMIT;", None)
.expect("execute");
assert_eq!(out.len(), 2);
}
#[test]
fn handler_enforces_max_statement_size() {
let handler = test_handler();
let huge = "X".repeat(2 * 1024 * 1024);
let err = match handler.execute_sql(&huge, None) {
Ok(_) => panic!("max statement size should be enforced"),
Err(e) => e,
};
assert!(err.to_string().contains("max_statement_bytes"));
}
proptest! {
#[test]
fn parse_sql_never_panics_on_random_input(input in ".*") {
let _ = parse_sql(&input);
}
}
}