postgres.rs

#![cfg(feature = "postgres")]

use crate::dialect::{AnsiSqlDialect, Dialect, PostgresDialect};
use crate::docker::DockerParams;
use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext};
use crate::memory::Config;
use crate::util::sql::bench_to_postgres_param;
use crate::value::BenchValue;
use crate::valueprovider::{ColumnType, Columns};
use crate::{
	Benchmark, Index, KeyType, Projection, Scan, VectorDistance, VectorIndexStrategy,
	VectorQuerySpec,
};
use anyhow::{Result, anyhow, bail};
use chrono::{NaiveDateTime, TimeZone, Utc};
use rust_decimal::Decimal;
use std::hint::black_box;
use tokio_postgres::types::{Json, ToSql};
use tokio_postgres::{Client, NoTls, Row};

pub const DEFAULT: &str = "host=127.0.0.1 user=postgres password=postgres";

/// pgvector operator-class for a given [`VectorDistance`].
fn pgvector_ops_for(d: VectorDistance) -> &'static str {
	match d {
		VectorDistance::Cosine => "vector_cosine_ops",
		VectorDistance::Euclidean => "vector_l2_ops",
		VectorDistance::InnerProduct => "vector_ip_ops",
		VectorDistance::Manhattan => "vector_l1_ops",
	}
}

/// pgvector distance operator used in `ORDER BY` for a given [`VectorDistance`].
fn pgvector_op_for(d: VectorDistance) -> &'static str {
	match d {
		VectorDistance::Cosine => "<=>",
		VectorDistance::Euclidean => "<->",
		VectorDistance::InnerProduct => "<#>",
		VectorDistance::Manhattan => "<+>",
	}
}

/// Calculate Postgres specific memory allocation
fn calculate_postgres_memory() -> (u64, u64, u64, u64, u64, u64) {
	// Load the system memory
	let memory = Config::new();
	// Use ~50% of recommended cache allocation. Equal fraction to MySQL
	// `innodb_buffer_pool_size` and MariaDB equivalent so all three SQL
	// adapters get the same in-process cache budget under `--optimised`.
	let shared_buffers_gb = (memory.cache_gb / 2).max(1);
	// `effective_cache_size` is a planner hint, not an allocation; it tells
	// Postgres how much memory is available across shared_buffers + OS page
	// cache. Keep it at the full cache budget.
	let effective_cache_gb = memory.cache_gb;
	// Scale work_mem with shared_buffers
	let work_mem_mb = (shared_buffers_gb * 64).max(32);
	// Use 25% of shared_buffers, max 8GB
	let maintenance_work_mem_gb = (shared_buffers_gb / 4).clamp(1, 8);
	// Scale WAL with shared_buffers
	let max_wal_gb = (shared_buffers_gb).clamp(2, 16);
	let min_wal_gb = (max_wal_gb / 4).max(1);
	// Return configuration
	(
		shared_buffers_gb,
		effective_cache_gb,
		work_mem_mb,
		maintenance_work_mem_gb,
		max_wal_gb,
		min_wal_gb,
	)
}

pub(crate) fn docker(options: &Benchmark) -> DockerParams {
	// Calculate memory allocation
	let (
		shared_buffers_gb,
		effective_cache_gb,
		work_mem_mb,
		maintenance_work_mem_gb,
		max_wal_gb,
		min_wal_gb,
	) = calculate_postgres_memory();
	// `synchronous_commit=on` makes every commit force a `write(2)` of the WAL
	// to the OS page cache before returning; `fsync` then carries the `--sync`
	// flag: `on` fsyncs that WAL to disk per commit (full durability), `off`
	// skips the fsync but keeps the per-commit page-cache write. The no-`--sync`
	// profile is therefore process-crash safe / power-crash unsafe, matching
	// SurrealDB (RocksDB `sync=never`) and ArcadeDB (`txWALFlush=1`) — rather
	// than `synchronous_commit=off`, which defers the `write(2)` itself and so
	// skips work the other adapters still do. See issue #256.
	let fsync_setting = if options.sync {
		"on"
	} else {
		"off"
	};
	// Return Docker parameters
	DockerParams {
		// `pgvector/pgvector` is the upstream pgvector image — same upstream
		// Postgres binaries as `postgres:*` with the `vector` extension
		// pre-installed. Required by the vector-search benchmark; harmless
		// for non-vector workloads (the extension is only created when the
		// schema declares a `FloatVector` column, see `startup`).
		image: "pgvector/pgvector:pg17",
		pre_args:
			"--ulimit nofile=65536:65536 -p 127.0.0.1:5432:5432 -e POSTGRES_PASSWORD=postgres"
				.to_string(),
		post_args: match options.optimised {
			// Optimised configuration
			true => format!(
				"postgres -N 1024 \
				-c shared_buffers={shared_buffers_gb}GB \
				-c effective_cache_size={effective_cache_gb}GB \
				-c work_mem={work_mem_mb}MB \
				-c maintenance_work_mem={maintenance_work_mem_gb}GB \
				-c wal_buffers=16MB \
				-c checkpoint_timeout=15min \
				-c checkpoint_completion_target=0.9 \
				-c random_page_cost=1.1 \
				-c effective_io_concurrency=200 \
				-c min_wal_size={min_wal_gb}GB \
				-c max_wal_size={max_wal_gb}GB \
				-c fsync={fsync_setting} \
				-c synchronous_commit=on"
			),
			// Default configuration
			false => format!(
				"postgres -N 1024 \
				-c fsync={fsync_setting} \
				-c synchronous_commit=on"
			),
		},
	}
}

pub(crate) struct PostgresClientProvider(KeyType, Columns, String);

impl BenchmarkEngine<PostgresClient> for PostgresClientProvider {
	/// Initiates a new datastore benchmarking engine
	async fn setup(kt: KeyType, columns: Columns, options: &Benchmark) -> Result<Self> {
		// Get the custom endpoint if specified
		let url = options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned();
		// Create the client provider
		Ok(Self(kt, columns, url))
	}
	/// Creates a new client for this benchmarking engine
	async fn create_client(&self) -> Result<PostgresClient> {
		// Connect to the database with TLS disabled
		let (client, connection) = tokio_postgres::connect(&self.2, NoTls).await?;
		// Log any errors when the connection is closed
		tokio::spawn(async move {
			if let Err(e) = connection.await {
				eprintln!("connection error: {e}");
			}
		});
		// Create the client
		Ok(PostgresClient {
			client,
			kt: self.0,
			columns: self.1.clone(),
		})
	}
}

pub(crate) struct PostgresClient {
	client: Client,
	kt: KeyType,
	columns: Columns,
}

impl BenchmarkClient for PostgresClient {
	// The return type when reading a row
	type ReadRow = BenchValue;

	async fn startup(&self) -> Result<()> {
		// Ensure pgvector is installed when the schema declares a vector column.
		// Safe to issue unconditionally; CREATE EXTENSION IF NOT EXISTS is a no-op
		// when the extension is already present.
		if self.columns.0.iter().any(|(_, t)| matches!(t, ColumnType::FloatVector(_))) {
			self.client.batch_execute("CREATE EXTENSION IF NOT EXISTS vector;").await?;
		}
		let id_type = match self.kt {
			KeyType::Integer => "SERIAL",
			KeyType::String26 => "VARCHAR(26)",
			KeyType::String90 => "VARCHAR(90)",
			KeyType::String250 => "VARCHAR(250)",
			KeyType::String506 => "VARCHAR(506)",
			KeyType::Uuid => {
				todo!()
			}
		};
		let fields = self
			.columns
			.0
			.iter()
			.map(|(n, t)| {
				let n = AnsiSqlDialect::escape_field(n.clone());
				match t {
					ColumnType::String => format!("{n} TEXT NOT NULL"),
					ColumnType::Integer => format!("{n} INTEGER NOT NULL"),
					ColumnType::Object => format!("{n} JSONB NOT NULL"),
					ColumnType::Array => format!("{n} JSONB NOT NULL"),
					ColumnType::Float => format!("{n} REAL NOT NULL"),
					ColumnType::DateTime => format!("{n} TIMESTAMP NOT NULL"),
					ColumnType::Uuid => format!("{n} UUID NOT NULL"),
					ColumnType::Decimal => format!("{n} NUMERIC(38, 10) NOT NULL"),
					ColumnType::Bool => format!("{n} BOOL NOT NULL"),
					ColumnType::Bytes => format!("{n} BYTEA NOT NULL"),
					ColumnType::FloatVector(dim) => format!("{n} vector({dim}) NOT NULL"),
				}
			})
			.collect::<Vec<String>>()
			.join(", ");
		let stm = format!(
			"DROP TABLE IF EXISTS record; CREATE TABLE record ( id {id_type} PRIMARY KEY, {fields});"
		);
		self.client.batch_execute(&stm).await?;
		Ok(())
	}

	async fn create_u32(&self, key: u32, val: BenchValue) -> Result<()> {
		self.create(key as i32, val).await
	}

	async fn create_string(&self, key: String, val: BenchValue) -> Result<()> {
		self.create(key, val).await
	}

	async fn read_u32(&self, key: u32) -> Result<BenchValue> {
		self.read(key as i32).await
	}

	async fn read_string(&self, key: String) -> Result<BenchValue> {
		self.read(key).await
	}

	async fn update_u32(&self, key: u32, val: BenchValue) -> Result<()> {
		self.update(key as i32, val).await
	}

	async fn update_string(&self, key: String, val: BenchValue) -> Result<()> {
		self.update(key, val).await
	}

	async fn delete_u32(&self, key: u32) -> Result<()> {
		self.delete(key as i32).await
	}

	async fn delete_string(&self, key: String) -> Result<()> {
		self.delete(key).await
	}

	async fn build_index(&self, spec: &Index, name: &str) -> Result<()> {
		// COUNT-style indexes have no Postgres equivalent; the indexed scan
		// leg runs the same query as the baseline so the row still populates.
		if spec.index_type.as_deref() == Some("count") {
			return Ok(());
		}
		// Get the unique flag
		let unique = if spec.unique.unwrap_or(false) {
			"UNIQUE"
		} else {
			""
		}
		.to_string();
		// Get the fields
		let fields = PostgresDialect::btree_index_key_list(&self.columns, spec);
		// Surreal-style `tags.*` specs target a JSONB array/object column that the
		// scan filters with the `@>` containment operator — a B-tree cannot serve
		// it, so these columns need a GIN index instead.
		let gin_columns = PostgresDialect::gin_containment_columns(&self.columns, spec);
		// Check if an index type is specified
		let stmt = match &spec.index_type {
			Some(kind) if kind == "fulltext" => {
				// Create a GIN index for full-text search
				let tsvector_expr = if spec.fields.len() == 1 {
					format!("to_tsvector('english', {})", spec.fields[0])
				} else {
					format!("to_tsvector('english', {})", spec.fields.join(" || ' ' || "))
				};
				format!("CREATE INDEX {name} ON record USING GIN ({tsvector_expr})")
			}
			Some(kind) => {
				format!("CREATE {unique} INDEX {name} ON record USING {kind} ({fields})")
			}
			// Auto-select GIN for JSONB array-containment fields. GIN returns an
			// unordered bitmap, so it indexes only the JSONB column(s) with
			// `jsonb_path_ops`; any companion ordering column (e.g. created_at) is
			// dropped, since it cannot coexist usefully in a GIN index. This makes
			// the `@>` filter index-assisted (the prior B-tree was unusable), though
			// the trailing `ORDER BY ... LIMIT` still forces a Top-N sort.
			None if !gin_columns.is_empty() => {
				let cols = gin_columns
					.iter()
					.map(|c| format!("{c} jsonb_path_ops"))
					.collect::<Vec<_>>()
					.join(", ");
				format!("CREATE INDEX {name} ON record USING GIN ({cols})")
			}
			None => {
				format!("CREATE {unique} INDEX {name} ON record ({fields})")
			}
		};
		// Create the index
		self.client.execute(&stmt, &[]).await?;
		// All ok
		Ok(())
	}

	async fn drop_index(&self, name: &str) -> Result<()> {
		let stmt = format!("DROP INDEX IF EXISTS {name}");
		self.client.execute(&stmt, &[]).await?;
		Ok(())
	}

	async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result<usize> {
		self.scan(scan).await
	}

	async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result<usize> {
		self.scan(scan).await
	}

	async fn build_vector_index(
		&self,
		spec: &Index,
		vq: &VectorQuerySpec,
		_dim: usize,
		name: &str,
	) -> Result<()> {
		let fields = spec.fields.join(", ");
		let ops = pgvector_ops_for(vq.distance);
		let stmt = match vq.index_strategy {
			VectorIndexStrategy::Bruteforce => bail!(crate::benchmark::NOT_SUPPORTED_ERROR),
			VectorIndexStrategy::Hnsw {
				m,
				ef_construction,
				..
			} => format!(
				"CREATE INDEX {name} ON record USING hnsw ({fields} {ops}) WITH (m = {m}, ef_construction = {ef_construction})"
			),
			VectorIndexStrategy::DiskAnn {
				..
			} => bail!(crate::benchmark::NOT_SUPPORTED_ERROR),
		};
		self.client.execute(&stmt, &[]).await?;
		// Set ef_search for HNSW (per-session GUC).
		if let VectorIndexStrategy::Hnsw {
			ef_search,
			..
		} = vq.index_strategy
		{
			let s = format!("SET hnsw.ef_search = {ef_search}");
			self.client.execute(&s, &[]).await?;
		}
		Ok(())
	}

	async fn scan_vector_u32(
		&self,
		scan: &Scan,
		query: &[f32],
		_ctx: ScanContext,
	) -> Result<usize> {
		self.knn_scan(scan, query).await
	}

	async fn scan_vector_string(
		&self,
		scan: &Scan,
		query: &[f32],
		_ctx: ScanContext,
	) -> Result<usize> {
		self.knn_scan(scan, query).await
	}

	async fn batch_create_u32(
		&self,
		key_vals: impl Iterator<Item = (u32, BenchValue)> + Send,
	) -> Result<()> {
		self.batch_create(key_vals.map(|(k, v)| (k as i32, v)).collect()).await
	}

	async fn batch_create_string(
		&self,
		key_vals: impl Iterator<Item = (String, BenchValue)> + Send,
	) -> Result<()> {
		self.batch_create(key_vals.collect()).await
	}

	async fn batch_read_u32(&self, keys: impl Iterator<Item = u32> + Send) -> Result<()> {
		self.batch_read(keys.map(|k| k as i32).collect()).await
	}

	async fn batch_read_string(&self, keys: impl Iterator<Item = String> + Send) -> Result<()> {
		self.batch_read(keys.collect()).await
	}

	async fn batch_update_u32(
		&self,
		key_vals: impl Iterator<Item = (u32, BenchValue)> + Send,
	) -> Result<()> {
		self.batch_update(key_vals.map(|(k, v)| (k as i32, v)).collect()).await
	}

	async fn batch_update_string(
		&self,
		key_vals: impl Iterator<Item = (String, BenchValue)> + Send,
	) -> Result<()> {
		self.batch_update(key_vals.collect()).await
	}

	async fn batch_delete_u32(&self, keys: impl Iterator<Item = u32> + Send) -> Result<()> {
		self.batch_delete(keys.map(|k| k as i32).collect()).await
	}

	async fn batch_delete_string(&self, keys: impl Iterator<Item = String> + Send) -> Result<()> {
		self.batch_delete(keys.collect()).await
	}
}

impl PostgresClient {
	fn consume(&self, row: Row, columns: bool) -> Result<BenchValue> {
		let mut val: Vec<(String, BenchValue)> = Vec::new();
		match self.kt {
			KeyType::Integer => {
				let v: i32 = row.try_get("id")?;
				val.push(("id".into(), BenchValue::Int(v as i64)));
			}
			KeyType::String26 | KeyType::String90 | KeyType::String250 | KeyType::String506 => {
				let v: String = row.try_get("id")?;
				val.push(("id".into(), BenchValue::String(v)));
			}
			KeyType::Uuid => {
				let v: uuid::Uuid = row.try_get("id")?;
				val.push(("id".into(), BenchValue::Uuid(v)));
			}
		}
		if columns {
			for (n, t) in self.columns.0.iter() {
				let bv = match t {
					ColumnType::Bool => {
						let v: bool = row.try_get(n.as_str())?;
						BenchValue::Bool(v)
					}
					ColumnType::Float => {
						let v: f32 = row.try_get(n.as_str())?;
						BenchValue::Float(v as f64)
					}
					ColumnType::Integer => {
						let v: i32 = row.try_get(n.as_str())?;
						BenchValue::Int(v as i64)
					}
					ColumnType::String => {
						let v: String = row.try_get(n.as_str())?;
						BenchValue::String(v)
					}
					ColumnType::DateTime => {
						let v: NaiveDateTime = row.try_get(n.as_str())?;
						BenchValue::DateTime(Utc.from_utc_datetime(&v))
					}
					ColumnType::Uuid => {
						let v: uuid::Uuid = row.try_get(n.as_str())?;
						BenchValue::Uuid(v)
					}
					ColumnType::Decimal => {
						let v: Decimal = row.try_get(n.as_str())?;
						BenchValue::Decimal(v)
					}
					ColumnType::Bytes => {
						let v: Vec<u8> = row.try_get(n.as_str())?;
						BenchValue::Bytes(v)
					}
					ColumnType::Object | ColumnType::Array => {
						let v: Json<serde_json::Value> = row.try_get(n.as_str())?;
						BenchValue::from(&v.0)
					}
					ColumnType::FloatVector(_) => {
						let v: pgvector::Vector = row.try_get(n.as_str())?;
						BenchValue::FloatVector(v.to_vec())
					}
				};
				val.push((n.clone(), bv));
			}
		}
		Ok(BenchValue::Object(val))
	}

	async fn create<T>(&self, key: T, val: BenchValue) -> Result<()>
	where
		T: ToSql + Sync + Send,
	{
		let obj = val.into_object()?;
		let (columns, placeholders) = AnsiSqlDialect::create_clause(&self.columns);
		let stm = format!("INSERT INTO record (id, {columns}) VALUES ($1, {placeholders})");
		let mut owned: Vec<Box<dyn ToSql + Sync + Send>> = vec![Box::new(key)];
		for (column, column_type) in &self.columns.0 {
			let v = obj
				.iter()
				.find(|(k, _)| k == column)
				.map(|(_, v)| v)
				.ok_or_else(|| anyhow!("Missing value for column {column}"))?;
			owned.push(bench_to_postgres_param(column, column_type, v)?);
		}
		let params: Vec<&(dyn ToSql + Sync)> =
			owned.iter().map(|b| b.as_ref() as &(dyn ToSql + Sync)).collect();
		let res = self.client.execute(&stm, &params).await?;
		assert_eq!(res, 1);
		Ok(())
	}

	async fn read<T>(&self, key: T) -> Result<BenchValue>
	where
		T: ToSql + Sync,
	{
		let stm = "SELECT * FROM record WHERE id=$1";
		let res = self.client.query(stm, &[&key]).await?;
		assert_eq!(res.len(), 1);
		Ok(black_box(self.consume(res.into_iter().next().unwrap(), true)?))
	}

	async fn update<T>(&self, key: T, val: BenchValue) -> Result<()>
	where
		T: ToSql + Sync + Send,
	{
		let obj = val.into_object()?;
		let set = AnsiSqlDialect::update_clause(&self.columns);
		let stm = format!("UPDATE record SET {set} WHERE id = $1");
		let mut owned: Vec<Box<dyn ToSql + Sync + Send>> = vec![Box::new(key)];
		for (column, column_type) in &self.columns.0 {
			let v = obj
				.iter()
				.find(|(k, _)| k == column)
				.map(|(_, v)| v)
				.ok_or_else(|| anyhow!("Missing value for column {column}"))?;
			owned.push(bench_to_postgres_param(column, column_type, v)?);
		}
		let params: Vec<&(dyn ToSql + Sync)> =
			owned.iter().map(|b| b.as_ref() as &(dyn ToSql + Sync)).collect();
		let res = self.client.execute(&stm, &params).await?;
		assert_eq!(res, 1);
		Ok(())
	}

	async fn delete<T>(&self, key: T) -> Result<()>
	where
		T: ToSql + Sync,
	{
		let stm = "DELETE FROM record WHERE id=$1";
		let res = self.client.execute(stm, &[&key]).await?;
		assert_eq!(res, 1);
		Ok(())
	}

	async fn scan(&self, scan: &Scan) -> Result<usize> {
		// Extract parameters
		let s = scan.start.map(|s| format!("OFFSET {}", s)).unwrap_or_default();
		let l = scan.limit.map(|s| format!("LIMIT {}", s)).unwrap_or_default();
		let c = PostgresDialect::filter_clause(scan)?;
		let o = AnsiSqlDialect::order_by_clause(scan)?;
		let p = scan.projection()?;
		// Perform the relevant projection scan type
		match p {
			Projection::Id => {
				let stm = format!("SELECT id FROM record {c} {o} {l} {s}");
				let res = self.client.query(&stm, &[]).await?;
				// We use a for loop to iterate over the results, while
				// calling black_box internally. This is necessary as
				// an iterator with `filter_map` or `map` is optimised
				// out by the compiler when calling `count` at the end.
				let mut count = 0;
				for v in res {
					black_box(self.consume(v, false).unwrap());
					count += 1;
				}
				Ok(count)
			}
			Projection::Full => {
				let stm = format!("SELECT * FROM record {c} {o} {l} {s}");
				let res = self.client.query(&stm, &[]).await?;
				// We use a for loop to iterate over the results, while
				// calling black_box internally. This is necessary as
				// an iterator with `filter_map` or `map` is optimised
				// out by the compiler when calling `count` at the end.
				let mut count = 0;
				for v in res {
					black_box(self.consume(v, true).unwrap());
					count += 1;
				}
				Ok(count)
			}
			Projection::Count => {
				let stm = format!("SELECT COUNT(*) FROM (SELECT id FROM record {c} {l} {s})");
				let res = self.client.query(&stm, &[]).await?;
				let count: i64 = res.first().unwrap().get(0);
				Ok(count as usize)
			}
		}
	}

	async fn knn_scan(&self, scan: &Scan, query: &[f32]) -> Result<usize> {
		let vq = scan
			.vector_query
			.as_ref()
			.ok_or_else(|| anyhow!("knn_scan: scan `{}` missing vector_query", scan.name))?;
		let op = pgvector_op_for(vq.distance);
		let field = AnsiSqlDialect::escape_field(vq.field.clone());
		let k = vq.top_k;
		let stm = format!("SELECT id FROM record ORDER BY {field} {op} $1 LIMIT {k}");
		let q = pgvector::Vector::from(query.to_vec());
		let res = self.client.query(&stm, &[&q]).await?;
		let mut count = 0;
		for v in res {
			black_box(self.consume(v, false).unwrap());
			count += 1;
		}
		Ok(count)
	}

	async fn batch_create<T>(&self, key_vals: Vec<(T, BenchValue)>) -> Result<()>
	where
		T: ToSql + Sync,
	{
		// Fetch the columns to insert
		let columns = AnsiSqlDialect::insert_columns(&self.columns);
		// Store the records to insert
		let mut inserts = Vec::with_capacity(key_vals.len());
		// Store the query parameters
		let mut params: Vec<&(dyn ToSql + Sync)> = Vec::new();
		// Store the column values
		let mut values: Vec<Box<dyn ToSql + Sync + Send>> = Vec::new();
		// Store the row index
		let mut index = 1;
		// Iterate over the key-value pairs
		for (_, val) in &key_vals {
			// Add the id placeholder
			let mut row = vec![format!("${index}")];
			index += 1;
			// Process the columns
			if let BenchValue::Object(obj) = val {
				for (column, column_type) in &self.columns.0 {
					// Add the column placeholder
					row.push(format!("${index}"));
					index += 1;
					// Add the column value with proper type conversion
					if let Some(value) = obj.iter().find(|(k, _)| k == column).map(|(_, v)| v) {
						let value = bench_to_postgres_param(column, column_type, value)?;
						values.push(value);
					} else {
						return Err(anyhow::anyhow!("Missing value for column {column}"));
					}
				}
			}
			// Add the row to the inserts
			inserts.push(format!("({})", row.join(", ")));
		}
		// Store the param index
		let mut index = 0;
		// Iterate over the key-value pairs
		for (key, val) in &key_vals {
			params.push(key);
			if let BenchValue::Object(_) = val {
				for _ in &self.columns.0 {
					params.push(values[index].as_ref());
					index += 1;
				}
			}
		}
		// Build and execute the INSERT statement
		let stm = format!("INSERT INTO record (id, {columns}) VALUES {}", inserts.join(", "));
		let res = self.client.execute(&stm, &params).await?;
		assert_eq!(res, key_vals.len() as u64);
		Ok(())
	}

	async fn batch_read<T>(&self, keys: Vec<T>) -> Result<()>
	where
		T: ToSql + Sync,
	{
		// Store the record ids
		let params: Vec<&(dyn ToSql + Sync)> =
			keys.iter().map(|k| k as &(dyn ToSql + Sync)).collect();
		// Build the IN clause
		let ids = (1..=keys.len()).map(|i| format!("${i}")).collect::<Vec<String>>().join(", ");
		// Build and execute the SELECT statement
		let stm = format!("SELECT * FROM record WHERE id IN ({ids})");
		let res = self.client.query(&stm, &params).await?;
		assert_eq!(res.len(), keys.len());
		for row in res {
			black_box(self.consume(row, true).unwrap());
		}
		Ok(())
	}

	async fn batch_update<T>(&self, key_vals: Vec<(T, BenchValue)>) -> Result<()>
	where
		T: ToSql + Sync,
	{
		// Store the columns to update
		let columns = self
			.columns
			.0
			.iter()
			.map(|(name, _)| {
				format!("{name} = data.{name}", name = AnsiSqlDialect::escape_field(name.clone()),)
			})
			.collect::<Vec<String>>()
			.join(", ");
		// Store the columns to select
		let fields = format!("id, {}", AnsiSqlDialect::insert_columns(&self.columns));
		// Store the records to insert
		let mut inserts = Vec::with_capacity(key_vals.len());
		// Store the query parameters
		let mut params: Vec<&(dyn ToSql + Sync)> = Vec::new();
		// Store the column values
		let mut values: Vec<Box<dyn ToSql + Sync + Send>> = Vec::new();
		// Store the row index
		let mut index = 1;
		// Iterate over the key-value pairs
		for (_, val) in &key_vals {
			// Start with the key parameter placeholder
			let mut row = vec![format!("${index}::{}", get_key_type(&self.kt))];
			index += 1;
			// Process each column value in the record
			if let BenchValue::Object(obj) = val {
				for (column, column_type) in &self.columns.0 {
					// Add parameter placeholder for this column
					row.push(format!("${index}::{}", get_column_type(column_type)));
					index += 1;
					// Add the column value with proper type conversion
					if let Some(value) = obj.iter().find(|(k, _)| k == column).map(|(_, v)| v) {
						let value = bench_to_postgres_param(column, column_type, value)?;
						values.push(value);
					} else {
						return Err(anyhow::anyhow!("Missing value for column {column}"));
					}
				}
			}
			// Add the complete row to the VALUES construct
			inserts.push(format!("({})", row.join(", ")));
		}
		// Store the param index
		let mut index = 0;
		// Iterate over the key-value pairs
		for (key, val) in &key_vals {
			// Add the key as the first parameter
			params.push(key);
			// Add all column values as subsequent parameters
			if let BenchValue::Object(_) = val {
				for _ in &self.columns.0 {
					params.push(values[index].as_ref());
					index += 1;
				}
			}
		}
		// Build and execute the UPDATE statement
		let stm = format!(
			"UPDATE record SET {columns} FROM (VALUES {}) AS data({fields}) WHERE record.id = data.id",
			inserts.join(", "),
		);
		let res = self.client.execute(&stm, &params).await?;
		assert_eq!(res, key_vals.len() as u64);
		Ok(())
	}

	async fn batch_delete<T>(&self, keys: Vec<T>) -> Result<()>
	where
		T: ToSql + Sync,
	{
		// Store the record ids
		let params: Vec<&(dyn ToSql + Sync)> =
			keys.iter().map(|k| k as &(dyn ToSql + Sync)).collect();
		// Build the IN clause
		let ids = (1..=keys.len()).map(|i| format!("${i}")).collect::<Vec<String>>().join(", ");
		// Build and execute the DELETE statement
		let stm = format!("DELETE FROM record WHERE id IN ({ids})");
		let res = self.client.execute(&stm, &params).await?;
		assert_eq!(res as usize, keys.len());
		Ok(())
	}
}

/// Get PostgreSQL type name for explicit casting
fn get_key_type(key_type: &KeyType) -> &'static str {
	match key_type {
		KeyType::Integer => "INTEGER",
		KeyType::String26 => "TEXT",
		KeyType::String90 => "TEXT",
		KeyType::String250 => "TEXT",
		KeyType::String506 => "TEXT",
		KeyType::Uuid => "UUID",
	}
}

/// Get PostgreSQL type name for explicit casting
fn get_column_type(column_type: &ColumnType) -> &'static str {
	match column_type {
		ColumnType::Integer => "INTEGER",
		ColumnType::Float => "REAL",
		ColumnType::Bool => "BOOLEAN",
		ColumnType::String => "TEXT",
		ColumnType::Object => "JSONB",
		ColumnType::Array => "JSONB",
		ColumnType::DateTime => "TIMESTAMP",
		ColumnType::Uuid => "UUID",
		ColumnType::Decimal => "NUMERIC",
		ColumnType::Bytes => "BYTEA",
		ColumnType::FloatVector(_) => "vector",
	}
}