1use anyhow::{Context, Result, anyhow, bail};
2use serde_json::Value;
3use std::collections::{HashMap, HashSet};
4use std::fs;
5use std::io;
6use std::path::Path;
7use std::path::PathBuf;
8use std::sync::Arc;
9use tempfile::TempDir;
10#[cfg(feature = "tools")]
11use tokio::io::{AsyncWrite, AsyncWriteExt};
12#[cfg(feature = "tools")]
13use tokio::runtime::Runtime;
14#[cfg(feature = "tools")]
15use wasmer_wasix::virtual_net::VirtualTcpSocket;
16#[cfg(feature = "tools")]
17use wasmer_wasix::virtual_net::tcp_pair::TcpSocketHalfRx;
18
19use crate::oliphaunt::aot;
20#[cfg(feature = "extensions")]
21use crate::oliphaunt::assets;
22use crate::oliphaunt::backend::{BackendOpenKind, BackendSession};
23#[cfg(feature = "extensions")]
24use crate::oliphaunt::base::install_bundled_extension_bytes;
25use crate::oliphaunt::base::{InstallOutcome, OliphauntPaths, RootLock};
26use crate::oliphaunt::builder::OliphauntBuilder;
27use crate::oliphaunt::config::{PostgresConfig, StartupConfig};
28use crate::oliphaunt::data_dir::{DataDirArchiveFormat, dump_pgdata_archive};
29use crate::oliphaunt::engine::EngineCapabilities;
30use crate::oliphaunt::errors::OliphauntError;
31#[cfg(feature = "extensions")]
32use crate::oliphaunt::extensions::{
33 Extension, by_sql_name, extension_session_setup_sql, extension_setup_sql, resolve_extension_set,
34};
35use crate::oliphaunt::interface::{
36 DataTransferContainer, DescribeQueryParam, DescribeQueryResult, DescribeResultField,
37 ExecProtocolOptions, ExecProtocolResult, NoticeCallback, ParserMap, QueryOptions, Results,
38 SerializerMap,
39};
40use crate::oliphaunt::parse::{
41 command_tag_row_count, parse_describe_statement_results, parse_results,
42};
43#[cfg(feature = "tools")]
44use crate::oliphaunt::pg_dump::{PgDumpOptions, PgDumpVirtualSocket, dump_direct_sql};
45#[cfg(feature = "extensions")]
46use crate::oliphaunt::postgres_mod::PostgresMod;
47use crate::oliphaunt::timing;
48use crate::oliphaunt::types::{
49 ArrayTypeInfo, DEFAULT_PARSERS, DEFAULT_SERIALIZERS, TEXT, register_array_type,
50};
51#[cfg(feature = "tools")]
52use crate::oliphaunt::wire::{FrontendFrameKind, FrontendFrameReader, classify_frontend_message};
53use crate::protocol::messages::{BackendMessage, DatabaseError};
54use crate::protocol::parser::Parser as ProtocolParser;
55use crate::protocol::serializer::{BindConfig, BindValue, PortalTarget, Serialize};
56
57type ChannelCallback = Arc<dyn Fn(&str) + Send + Sync + 'static>;
58type GlobalCallback = Arc<dyn Fn(&str, &str) + Send + Sync + 'static>;
59
60#[derive(Debug, Clone, PartialEq, Eq, Hash)]
61pub struct ListenerHandle {
62 channel: String,
63 normalized_channel: String,
64 id: u64,
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
68pub struct GlobalListenerHandle {
69 id: u64,
70}
71
72impl ListenerHandle {
73 pub fn channel(&self) -> &str {
74 &self.channel
75 }
76
77 pub fn id(&self) -> u64 {
78 self.id
79 }
80}
81
82impl GlobalListenerHandle {
83 pub fn id(&self) -> u64 {
84 self.id
85 }
86}
87
88struct ChannelListener {
89 id: u64,
90 callback: ChannelCallback,
91}
92
93struct GlobalListener {
94 id: u64,
95 callback: GlobalCallback,
96}
97
98enum ExecTransportResult {
99 CommandOnly(Vec<usize>),
100 Raw(Vec<u8>),
101}
102
103pub struct Oliphaunt {
105 backend: BackendSession,
106 _temp_dir: Option<TempDir>,
107 _root_lock: Option<RootLock>,
108 parser: ProtocolParser,
109 serializers: SerializerMap,
110 parsers: ParserMap,
111 array_type_lookup_misses: HashSet<i32>,
112 in_transaction: bool,
113 ready: bool,
114 closing: bool,
115 closed: bool,
116 blob_input_provided: bool,
117 notify_listeners: HashMap<String, Vec<ChannelListener>>,
118 global_notify_listeners: Vec<GlobalListener>,
119 next_listener_id: u64,
120 next_global_listener_id: u64,
121}
122
123impl Oliphaunt {
124 pub fn builder() -> OliphauntBuilder {
126 OliphauntBuilder::new()
127 }
128
129 pub fn open(root: impl AsRef<Path>) -> Result<Self> {
131 Self::builder().path(root.as_ref().to_path_buf()).open()
132 }
133
134 pub fn open_app(app_id: (&str, &str, &str)) -> Result<Self> {
136 Self::builder().app_id(app_id).open()
137 }
138
139 pub fn temporary() -> Result<Self> {
141 Self::builder().temporary().open()
142 }
143
144 pub fn preload() -> Result<()> {
146 let (temp_dir, paths) = {
147 let _phase = timing::phase("preload.tempdir");
148 OliphauntPaths::with_temp_dir()?
149 };
150 {
151 let _phase = timing::phase("preload.runtime_module");
152 crate::oliphaunt::base::preload_runtime_module(&paths)?;
153 }
154 {
155 let _phase = timing::phase("preload.aot_runtime");
156 aot::preload_runtime_artifact()?;
157 }
158 drop(temp_dir);
159 Ok(())
160 }
161
162 #[cfg(feature = "extensions")]
164 pub fn preload_extensions(extensions: impl IntoIterator<Item = Extension>) -> Result<()> {
165 Self::preload()?;
166 let extensions = extensions.into_iter().collect::<Vec<_>>();
167 for extension in resolve_extension_set(&extensions)? {
168 let bytes = assets::extension_archive(extension.sql_name()).ok_or_else(|| {
169 anyhow!(
170 "extension asset '{}' is not bundled in this oliphaunt-wasix build",
171 extension.sql_name()
172 )
173 })?;
174 let (temp_dir, paths) = {
175 let _phase = timing::phase("preload.extension_tempdir");
176 OliphauntPaths::with_temp_dir()?
177 };
178 {
179 let _phase = timing::phase("preload.extension_runtime_module");
180 crate::oliphaunt::base::preload_runtime_module(&paths)?;
181 }
182 {
183 let _phase = timing::phase("preload.extension_archive_install");
184 install_bundled_extension_bytes(&paths, extension.sql_name(), bytes)?;
185 }
186 {
187 let _phase = timing::phase("preload.extension_side_module");
188 PostgresMod::preload_extension_module_from_paths(&paths, extension)?;
189 }
190 {
191 let _phase = timing::phase("preload.extension_aot");
192 aot::preload_extension_artifact(extension)?;
193 }
194 drop(temp_dir);
195 }
196 Ok(())
197 }
198
199 #[doc(hidden)]
201 pub fn new(paths: OliphauntPaths) -> Result<Self> {
202 let outcome = crate::oliphaunt::base::prepare_database_root(
203 paths,
204 crate::oliphaunt::base::RootPrepareOptions::template(),
205 )?;
206 Self::new_prepared(outcome)
207 }
208
209 pub(crate) fn new_prepared(outcome: InstallOutcome) -> Result<Self> {
210 Self::new_prepared_with_config(outcome, PostgresConfig::default(), StartupConfig::default())
211 }
212
213 pub(crate) fn new_prepared_with_config(
214 outcome: InstallOutcome,
215 postgres_config: PostgresConfig,
216 startup_config: StartupConfig,
217 ) -> Result<Self> {
218 let _phase = timing::phase("oliphaunt.open");
219 let session_startup_config = startup_config.clone();
220 let backend = BackendSession::open(
221 outcome,
222 postgres_config,
223 startup_config,
224 BackendOpenKind::Direct,
225 )?;
226
227 let mut instance = {
228 let _phase = timing::phase("oliphaunt.client_struct_init");
229 Self {
230 backend,
231 _temp_dir: None,
232 _root_lock: None,
233 parser: ProtocolParser::new(),
234 serializers: DEFAULT_SERIALIZERS.clone(),
235 parsers: DEFAULT_PARSERS.clone(),
236 array_type_lookup_misses: HashSet::new(),
237 in_transaction: false,
238 ready: true,
239 closing: false,
240 closed: false,
241 blob_input_provided: false,
242 notify_listeners: HashMap::new(),
243 global_notify_listeners: Vec::new(),
244 next_listener_id: 1,
245 next_global_listener_id: 1,
246 }
247 };
248
249 if session_startup_config.username != "postgres" {
250 let sql = format!(
251 "SET ROLE {}",
252 crate::oliphaunt::templating::quote_identifier(&session_startup_config.username)
253 );
254 instance
255 .exec(&sql, None)
256 .with_context(|| format!("set startup role {}", session_startup_config.username))?;
257 }
258
259 Ok(instance)
260 }
261
262 #[cfg(feature = "extensions")]
264 pub fn enable_extension(&mut self, extension: Extension) -> Result<()> {
265 let _phase = timing::phase("extension.enable");
266 let bytes = assets::extension_archive(extension.sql_name()).ok_or_else(|| {
267 anyhow!(
268 "extension asset '{}' is not bundled in this oliphaunt-wasix build",
269 extension.sql_name()
270 )
271 })?;
272 install_bundled_extension_bytes(self.paths(), extension.sql_name(), bytes)?;
273 self.backend.preload_extension_module(extension)?;
274 for sql in extension_setup_sql(extension) {
275 self.exec(&sql, None)?;
276 }
277 Ok(())
278 }
279
280 #[cfg(feature = "extensions")]
281 pub(crate) fn enable_preinstalled_extension(&mut self, extension: Extension) -> Result<()> {
282 let _phase = timing::phase("extension.enable_preinstalled");
283 self.backend.preload_installed_extension(extension)?;
284 for sql in extension_session_setup_sql(extension) {
285 self.exec(&sql, None)?;
286 }
287 Ok(())
288 }
289
290 pub fn refresh_array_types(&mut self) -> Result<()> {
296 self.check_ready()?;
297 self.refresh_array_types_internal()
298 }
299
300 pub fn query(
302 &mut self,
303 sql: &str,
304 params: &[Value],
305 options: Option<&QueryOptions>,
306 ) -> Result<Results> {
307 self.check_ready()?;
308
309 self.query_internal(sql, params, options)
310 }
311
312 fn query_internal(
313 &mut self,
314 sql: &str,
315 params: &[Value],
316 options: Option<&QueryOptions>,
317 ) -> Result<Results> {
318 let default_options = QueryOptions::default();
319 let query_opts = options.unwrap_or(&default_options);
320
321 self.handle_blob_input(query_opts.blob.as_ref())?;
322
323 let params_snapshot: Vec<Value> = params.to_vec();
324 let options_snapshot = options.cloned();
325 let mut collected_messages: Vec<BackendMessage> = Vec::new();
326
327 let mut exec_opts = ExecProtocolOptions::no_sync();
328 exec_opts.on_notice = query_opts.on_notice.clone();
329 exec_opts.data_transfer_container = query_opts.data_transfer_container;
330
331 let result: Result<()> = (|| {
332 let param_types = if query_opts.param_types.is_empty() {
333 &[] as &[i32]
334 } else {
335 &query_opts.param_types
336 };
337
338 let mut messages = {
339 let _phase = timing::phase("client.query.parse_describe");
340 self.parse_and_describe(sql, param_types, exec_opts.clone())?
341 };
342 let mut data_type_ids = parse_describe_statement_results(&messages);
343 if self.ensure_array_types_for_bind_values(params, &data_type_ids, query_opts)? {
344 messages = {
345 let _phase = timing::phase("client.query.parse_describe_after_array_register");
346 self.parse_and_describe(sql, param_types, exec_opts.clone())?
347 };
348 data_type_ids = parse_describe_statement_results(&messages);
349 }
350 collected_messages.extend(messages);
351 let bind_values = {
352 let _phase = timing::phase("client.query.prepare_bind_values");
353 self.prepare_bind_values(params, &data_type_ids, query_opts)?
354 };
355 let bind_config = BindConfig {
356 values: bind_values,
357 ..Default::default()
358 };
359 let execute_batch = {
360 let _phase = timing::phase("client.query.serialize_execute");
361 let mut execute_batch = Vec::new();
362 execute_batch.extend(Serialize::bind(&bind_config));
363 execute_batch.extend(Serialize::describe(&PortalTarget::new('P', None)));
364 execute_batch.extend(Serialize::execute(None));
365 execute_batch.extend(Serialize::sync());
366 execute_batch
367 };
368 let ExecProtocolResult { messages, .. } = {
369 let _phase = timing::phase("client.query.execute_roundtrip");
370 self.exec_protocol(&execute_batch, exec_opts.clone())?
371 };
372 collected_messages.extend(messages);
373
374 Ok(())
375 })();
376
377 if let Err(err) = result {
378 match err.downcast::<DatabaseError>() {
379 Ok(db_err) => {
380 let enriched =
381 OliphauntError::new(db_err, sql, params_snapshot, options_snapshot);
382 return Err(enriched.into());
383 }
384 Err(err) => {
385 return Err(err.context(format!("failed to execute extended query: {sql}")));
386 }
387 }
388 }
389
390 {
391 let _phase = timing::phase("client.query.finish");
392 self.finish_query(collected_messages, options)
393 }
394 }
395
396 pub fn is_ready(&self) -> bool {
398 self.ready && !self.closing && !self.closed
399 }
400
401 pub fn engine_capabilities(&self) -> EngineCapabilities {
403 self.backend.capabilities()
404 }
405
406 #[doc(hidden)]
408 pub fn paths(&self) -> &OliphauntPaths {
409 self.backend.paths()
410 }
411
412 #[doc(hidden)]
414 #[cfg(debug_assertions)]
415 pub fn guest_bridge_allocation_counts(&self) -> (u64, u64) {
416 self.backend.guest_bridge_allocation_counts()
417 }
418
419 pub fn dump_data_dir(&mut self) -> Result<Vec<u8>> {
425 self.dump_data_dir_with_format(DataDirArchiveFormat::TarGz)
426 }
427
428 pub fn dump_data_dir_with_format(&mut self, format: DataDirArchiveFormat) -> Result<Vec<u8>> {
430 self.check_ready()?;
431 self.archive_quiesced_pgdata("dump PGDATA archive", format)
432 }
433
434 pub fn try_clone(&mut self) -> Result<Self> {
436 #[cfg(feature = "extensions")]
437 let extensions = self.bundled_extensions_in_database()?;
438 let archive = self.dump_data_dir_with_format(DataDirArchiveFormat::Tar)?;
439 let builder = Self::builder().temporary().load_data_dir_archive(archive);
440 #[cfg(feature = "extensions")]
441 let builder = builder.extensions(extensions);
442 builder.open()
443 }
444
445 #[cfg(feature = "tools")]
447 pub fn dump_sql(&mut self, options: PgDumpOptions) -> Result<String> {
448 self.check_ready()?;
449 options.validate()?;
450 self.checkpoint_backend_for_physical_snapshot("direct pg_dump")?;
451 self.dump_sql_via_direct_protocol(&options)
452 }
453
454 #[cfg(feature = "tools")]
456 pub fn dump_bytes(&mut self, options: PgDumpOptions) -> Result<Vec<u8>> {
457 Ok(self.dump_sql(options)?.into_bytes())
458 }
459
460 fn checkpoint_backend_for_physical_snapshot(&mut self, operation: &'static str) -> Result<()> {
461 if self.in_transaction {
462 bail!("{operation} cannot run while a direct transaction is active");
463 }
464 self.exec("CHECKPOINT", None)
465 .with_context(|| format!("checkpoint before {operation}"))?;
466 Ok(())
467 }
468
469 fn archive_quiesced_pgdata(
470 &mut self,
471 operation: &'static str,
472 format: DataDirArchiveFormat,
473 ) -> Result<Vec<u8>> {
474 self.checkpoint_backend_for_physical_snapshot(operation)?;
475 self.backend
476 .shutdown()
477 .with_context(|| format!("quiesce backend before {operation}"))?;
478
479 let archive = dump_pgdata_archive(
480 &self.backend.paths().pgdata,
481 self.backend.pgdata_template_root(),
482 format,
483 )
484 .with_context(|| format!("materialize physical PGDATA archive for {operation}"));
485 let restart = self
486 .backend
487 .restart()
488 .and_then(|_| self.restore_session_state_after_backend_restart())
489 .with_context(|| format!("restart backend after {operation}"));
490
491 match (archive, restart) {
492 (Ok(archive), Ok(())) => Ok(archive),
493 (Err(err), Ok(())) => Err(err),
494 (Ok(_), Err(err)) => {
495 self.ready = false;
496 self.closed = true;
497 Err(err)
498 }
499 (Err(err), Err(restart_err)) => {
500 self.ready = false;
501 self.closed = true;
502 Err(err.context(format!(
503 "backend restart after failed {operation} also failed: {restart_err:#}"
504 )))
505 }
506 }
507 }
508
509 fn restore_session_state_after_backend_restart(&mut self) -> Result<()> {
510 let username = self.backend.startup_config().username.clone();
511 if username != "postgres" {
512 let sql = format!(
513 "SET ROLE {}",
514 crate::oliphaunt::templating::quote_identifier(&username)
515 );
516 self.exec(&sql, None).with_context(|| {
517 format!("restore startup role {username} after backend restart")
518 })?;
519 }
520
521 let channels = self
522 .notify_listeners
523 .iter()
524 .filter(|(_, listeners)| !listeners.is_empty())
525 .map(|(channel, _)| channel.clone())
526 .collect::<Vec<_>>();
527 for channel in channels {
528 let quoted_channel = crate::oliphaunt::templating::quote_identifier(&channel);
529 self.exec_internal(&format!("LISTEN {quoted_channel}"), None)
530 .with_context(|| format!("restore LISTEN {channel} after backend restart"))?;
531 }
532 Ok(())
533 }
534
535 #[cfg(feature = "tools")]
536 fn dump_sql_via_direct_protocol(&mut self, options: &PgDumpOptions) -> Result<String> {
537 ensure_direct_pg_dump_options_match_session(self.backend.startup_config(), options)?;
538 let result = dump_direct_sql(options, |socket| self.serve_direct_pg_dump_protocol(socket));
539 let cleanup_result = self.cleanup_after_direct_pg_dump_session();
540
541 match (result, cleanup_result) {
542 (Ok(sql), Ok(())) => Ok(sql),
543 (Err(err), Ok(())) => Err(err),
544 (Ok(_), Err(err)) => Err(err),
545 (Err(err), Err(cleanup_err)) => Err(err.context(format!(
546 "direct pg_dump cleanup also failed: {cleanup_err:#}"
547 ))),
548 }
549 }
550
551 #[cfg(feature = "tools")]
552 fn cleanup_after_direct_pg_dump_session(&mut self) -> Result<()> {
553 self.exec("DEALLOCATE ALL; SET search_path TO DEFAULT;", None)
554 .context("reset direct pg_dump session state")?;
555 Ok(())
556 }
557
558 #[cfg(feature = "tools")]
559 fn serve_direct_pg_dump_protocol(&mut self, mut socket: PgDumpVirtualSocket) -> Result<()> {
560 let _ = socket.set_nodelay(true);
561 let (mut socket_tx, mut socket_rx) = socket.split();
562 let runtime = tokio::runtime::Builder::new_current_thread()
563 .enable_all()
564 .build()
565 .context("create direct pg_dump virtual socket runtime")?;
566 let mut reader = FrontendFrameReader::default();
567 let mut buffer = [0u8; 64 * 1024];
568 loop {
569 let read = read_direct_pg_dump_socket(&runtime, &mut socket_rx, &mut buffer)
570 .context("read direct pg_dump protocol socket")?;
571 if read == 0 {
572 return Ok(());
573 }
574 for message in reader.push(&buffer[..read])? {
575 match classify_frontend_message(&message)? {
576 FrontendFrameKind::SslOrGssRequest => {
577 write_direct_pg_dump_socket(&runtime, &mut socket_tx, b"N")
578 .context("write direct pg_dump SSL refusal")?;
579 }
580 FrontendFrameKind::CancelRequest | FrontendFrameKind::Terminate => {
581 return Ok(());
582 }
583 FrontendFrameKind::Startup => {
584 if let Some(response) = self.backend.existing_startup_response() {
585 write_direct_pg_dump_socket(&runtime, &mut socket_tx, &response)
586 .context("write direct pg_dump existing startup response")?;
587 } else {
588 let response = self.backend.startup_with_packet(&message)?;
589 write_direct_pg_dump_socket(&runtime, &mut socket_tx, &response.output)
590 .context("write direct pg_dump startup response")?;
591 if !response.accepted {
592 return Ok(());
593 }
594 }
595 }
596 FrontendFrameKind::Protocol => {
597 self.exec_protocol_raw_stream(
598 &message,
599 ExecProtocolOptions::no_sync(),
600 |chunk| {
601 write_direct_pg_dump_socket(&runtime, &mut socket_tx, chunk)
602 .context("write direct pg_dump backend protocol chunk")?;
603 Ok(())
604 },
605 )?;
606 }
607 }
608 }
609 flush_direct_pg_dump_socket(&runtime, &mut socket_tx)
610 .context("flush direct pg_dump socket")?;
611 }
612 }
613
614 #[cfg(feature = "extensions")]
615 fn bundled_extensions_in_database(&mut self) -> Result<Vec<Extension>> {
616 let results = self.query(
617 "SELECT extname FROM pg_catalog.pg_extension ORDER BY extname",
618 &[],
619 None,
620 )?;
621 let extensions = results
622 .rows
623 .iter()
624 .filter_map(|row| row.get("extname"))
625 .filter_map(|value| value.as_str())
626 .filter_map(by_sql_name)
627 .collect();
628 Ok(extensions)
629 }
630
631 pub(crate) fn attach_temp_dir(&mut self, temp_dir: TempDir) {
632 self._temp_dir = Some(temp_dir);
633 }
634
635 pub(crate) fn attach_root_lock(&mut self, root_lock: RootLock) {
636 self._root_lock = Some(root_lock);
637 }
638
639 pub fn is_closed(&self) -> bool {
641 self.closed
642 }
643
644 pub fn close(&mut self) -> Result<()> {
646 self.close_backend()
647 }
648
649 fn close_backend(&mut self) -> Result<()> {
650 if self.closed {
651 return Ok(());
652 }
653 if self.closing {
654 bail!("Oliphaunt is closing");
655 }
656
657 self.closing = true;
658 let result = (|| {
659 self.backend.shutdown()?;
660 self.sync_to_fs()
661 })();
662
663 self.closing = false;
664 if result.is_ok() {
665 self.closed = true;
666 self.ready = false;
667 self.notify_listeners.clear();
668 self.global_notify_listeners.clear();
669 self._root_lock = None;
670 }
671 result
672 }
673
674 #[cfg(feature = "extensions")]
675 pub(crate) fn close_for_template_cache(&mut self) -> Result<()> {
676 self.close_backend()
677 }
678
679 pub fn exec(&mut self, sql: &str, options: Option<&QueryOptions>) -> Result<Vec<Results>> {
681 self.check_ready()?;
682
683 self.exec_internal(sql, options)
684 }
685
686 fn exec_internal(&mut self, sql: &str, options: Option<&QueryOptions>) -> Result<Vec<Results>> {
687 let options_snapshot = options.cloned();
688 let default_options = QueryOptions::default();
689 let exec_opts_ref = options.unwrap_or(&default_options);
690 let mut exec_opts = ExecProtocolOptions::no_sync();
691 exec_opts.on_notice = exec_opts_ref.on_notice.clone();
692 exec_opts.data_transfer_container = exec_opts_ref.data_transfer_container;
693
694 self.handle_blob_input(exec_opts_ref.blob.as_ref())?;
695
696 let mut collected_messages: Vec<BackendMessage> = Vec::new();
697
698 let message = Serialize::query(sql);
699 let transport_result = {
700 let _phase = timing::phase("client.protocol_transport_send");
701 self.backend
702 .with_buffered(&message, exec_opts.data_transfer_container, |data| {
703 if let Some(affected_rows) = parse_command_only_result_counts(data) {
704 Ok(ExecTransportResult::CommandOnly(affected_rows))
705 } else {
706 Ok(ExecTransportResult::Raw(data.to_vec()))
707 }
708 })
709 };
710 let transport_result = match transport_result {
711 Ok(data) => data,
712 Err(err) => match err.downcast::<DatabaseError>() {
713 Ok(db_err) => {
714 let enriched = OliphauntError::new(db_err, sql, Vec::new(), options_snapshot);
715 return Err(enriched.into());
716 }
717 Err(err) => {
718 return Err(err.context(format!("failed to execute simple query: {sql}")));
719 }
720 },
721 };
722
723 let data = match transport_result {
724 ExecTransportResult::CommandOnly(affected_rows) => {
725 return self.finish_exec_command_only(affected_rows, options);
726 }
727 ExecTransportResult::Raw(data) => data,
728 };
729 let ExecProtocolResult { messages, .. } =
730 match self.parse_protocol_data(data, exec_opts.throw_on_error, exec_opts.on_notice) {
731 Ok(result) => result,
732 Err(err) => match err.downcast::<DatabaseError>() {
733 Ok(db_err) => {
734 let enriched =
735 OliphauntError::new(db_err, sql, Vec::new(), options_snapshot);
736 return Err(enriched.into());
737 }
738 Err(err) => {
739 return Err(err.context(format!("failed to execute simple query: {sql}")));
740 }
741 },
742 };
743 let has_row_description = messages
744 .iter()
745 .any(|message| matches!(message, BackendMessage::RowDescription(_)));
746 collected_messages.extend(messages);
747
748 self.finish_exec(collected_messages, options, has_row_description)
749 }
750
751 pub fn listen<F>(&mut self, channel: &str, callback: F) -> Result<ListenerHandle>
753 where
754 F: Fn(&str) + Send + Sync + 'static,
755 {
756 self.check_ready()?;
757
758 let quoted_channel = crate::oliphaunt::templating::quote_identifier(channel);
759 let normalized = channel.to_string();
760 let should_listen = match self.notify_listeners.get(&normalized) {
761 Some(existing) => existing.is_empty(),
762 None => true,
763 };
764
765 if should_listen {
766 self.exec_internal(&format!("LISTEN {quoted_channel}"), None)?;
767 }
768
769 let callback: ChannelCallback = Arc::new(callback);
770 let entry = self.notify_listeners.entry(normalized.clone()).or_default();
771 let id = self.next_listener_id;
772 self.next_listener_id = self.next_listener_id.wrapping_add(1);
773 entry.push(ChannelListener { id, callback });
774
775 Ok(ListenerHandle {
776 channel: channel.to_string(),
777 normalized_channel: normalized,
778 id,
779 })
780 }
781
782 pub fn unlisten(&mut self, handle: ListenerHandle) -> Result<()> {
784 if let Some(listeners) = self.notify_listeners.get_mut(&handle.normalized_channel) {
785 listeners.retain(|listener| listener.id != handle.id);
786 if listeners.is_empty() {
787 self.notify_listeners.remove(&handle.normalized_channel);
788 let quoted_channel =
789 crate::oliphaunt::templating::quote_identifier(&handle.channel);
790 self.exec_internal(&format!("UNLISTEN {quoted_channel}"), None)?;
791 }
792 }
793 Ok(())
794 }
795
796 pub fn unlisten_channel(&mut self, channel: &str) -> Result<()> {
798 let quoted_channel = crate::oliphaunt::templating::quote_identifier(channel);
799 let normalized = channel.to_string();
800 if self.notify_listeners.remove(&normalized).is_some() {
801 self.exec_internal(&format!("UNLISTEN {quoted_channel}"), None)?;
802 }
803 Ok(())
804 }
805
806 pub fn on_notification<F>(&mut self, callback: F) -> GlobalListenerHandle
808 where
809 F: Fn(&str, &str) + Send + Sync + 'static,
810 {
811 let id = self.next_global_listener_id;
812 self.next_global_listener_id = self.next_global_listener_id.wrapping_add(1);
813 let callback: GlobalCallback = Arc::new(callback);
814 self.global_notify_listeners
815 .push(GlobalListener { id, callback });
816 GlobalListenerHandle { id }
817 }
818
819 pub fn off_notification(&mut self, handle: GlobalListenerHandle) {
821 self.global_notify_listeners
822 .retain(|listener| listener.id != handle.id);
823 }
824
825 pub fn describe_query(
827 &mut self,
828 sql: &str,
829 options: Option<&QueryOptions>,
830 ) -> Result<DescribeQueryResult> {
831 self.check_ready()?;
832
833 let default_options = QueryOptions::default();
834 let query_opts = options.unwrap_or(&default_options);
835
836 let options_snapshot = options.cloned();
837 let mut exec_opts = ExecProtocolOptions::no_sync();
838 exec_opts.on_notice = query_opts.on_notice.clone();
839 exec_opts.data_transfer_container = query_opts.data_transfer_container;
840
841 let mut describe_messages: Vec<BackendMessage> = Vec::new();
842
843 let result: Result<()> = (|| {
844 let param_types = if query_opts.param_types.is_empty() {
845 &[] as &[i32]
846 } else {
847 &query_opts.param_types
848 };
849
850 let mut describe_batch = Vec::new();
851 describe_batch.extend(Serialize::parse(None, sql, param_types));
852 describe_batch.extend(Serialize::describe(&PortalTarget::new('S', None)));
853 describe_batch.extend(Serialize::sync());
854 let ExecProtocolResult { messages, .. } =
855 self.exec_protocol(&describe_batch, exec_opts.clone())?;
856 if !messages
857 .iter()
858 .any(|message| matches!(message, BackendMessage::ParseComplete { .. }))
859 {
860 bail!("extended query parse did not complete");
861 }
862 describe_messages.extend(messages);
863
864 Ok(())
865 })();
866
867 if let Err(err) = result {
868 match err.downcast::<DatabaseError>() {
869 Ok(db_err) => {
870 let enriched = OliphauntError::new(db_err, sql, Vec::new(), options_snapshot);
871 return Err(enriched.into());
872 }
873 Err(err) => {
874 return Err(err.context(format!("failed to describe query: {sql}")));
875 }
876 }
877 }
878
879 let param_type_ids = parse_describe_statement_results(&describe_messages);
880 self.ensure_array_types_for_oids(param_type_ids.iter().copied(), Some(query_opts))?;
881 let result_type_ids = describe_messages
882 .iter()
883 .filter_map(|msg| match msg {
884 BackendMessage::RowDescription(desc) => Some(desc),
885 _ => None,
886 })
887 .flat_map(|desc| desc.fields.iter().map(|field| field.data_type_id))
888 .collect::<Vec<_>>();
889 self.ensure_array_types_for_oids(result_type_ids.iter().copied(), Some(query_opts))?;
890
891 let query_params = param_type_ids
892 .into_iter()
893 .map(|oid| DescribeQueryParam {
894 data_type_id: oid,
895 serializer: self.serializers.get(&oid).cloned(),
896 })
897 .collect();
898
899 let result_fields = describe_messages
900 .iter()
901 .find_map(|msg| match msg {
902 BackendMessage::RowDescription(desc) => Some(
903 desc.fields
904 .iter()
905 .map(|field| DescribeResultField {
906 name: field.name.clone(),
907 data_type_id: field.data_type_id,
908 parser: self.parsers.get(&field.data_type_id).cloned(),
909 })
910 .collect::<Vec<_>>(),
911 ),
912 _ => None,
913 })
914 .unwrap_or_default();
915
916 Ok(DescribeQueryResult {
917 query_params,
918 result_fields,
919 })
920 }
921
922 pub fn transaction<F, T>(&mut self, mut callback: F) -> Result<T>
924 where
925 F: FnMut(&mut Transaction<'_>) -> Result<T>,
926 {
927 self.check_ready()?;
928
929 self.run_exec_command("BEGIN")?;
931 self.in_transaction = true;
932
933 let mut tx = Transaction::new(self);
934 let callback_result = callback(&mut tx);
935
936 let txn_result = match callback_result {
937 Ok(value) => {
938 if !tx.closed {
939 tx.commit_internal()?;
940 }
941 Ok(value)
942 }
943 Err(err) => {
944 if !tx.closed {
945 tx.rollback_internal()?;
946 }
947 Err(err)
948 }
949 };
950
951 self.in_transaction = false;
952 txn_result
953 }
954
955 pub fn sync_to_fs(&mut self) -> Result<()> {
962 Ok(())
963 }
964
965 fn prepare_bind_values(
966 &self,
967 params: &[Value],
968 data_type_ids: &[i32],
969 options: &QueryOptions,
970 ) -> Result<Vec<BindValue>> {
971 if params.is_empty() {
972 return Ok(Vec::new());
973 }
974
975 let mut values = Vec::with_capacity(params.len());
976 let overrides = if options.serializers.is_empty() {
977 None
978 } else {
979 Some(&options.serializers)
980 };
981
982 for (idx, value) in params.iter().enumerate() {
983 if value.is_null() {
984 values.push(BindValue::Null);
985 continue;
986 }
987
988 let oid = data_type_ids.get(idx).copied().unwrap_or(TEXT);
989 let serializer = overrides
990 .and_then(|map| map.get(&oid))
991 .or_else(|| self.serializers.get(&oid));
992
993 let serialized = match serializer {
994 Some(func) => func(value).with_context(|| {
995 format!("failed to serialize parameter {idx} using OID {oid}")
996 })?,
997 None => self.default_serialize_value(value),
998 };
999
1000 values.push(BindValue::Text(serialized));
1001 }
1002
1003 Ok(values)
1004 }
1005
1006 fn parse_and_describe(
1007 &mut self,
1008 sql: &str,
1009 param_types: &[i32],
1010 exec_opts: ExecProtocolOptions,
1011 ) -> Result<Vec<BackendMessage>> {
1012 let mut prepare_batch = Vec::new();
1013 prepare_batch.extend(Serialize::parse(None, sql, param_types));
1014 prepare_batch.extend(Serialize::describe(&PortalTarget::new('S', None)));
1015 prepare_batch.extend(Serialize::sync());
1016 let ExecProtocolResult { messages, .. } = self.exec_protocol(&prepare_batch, exec_opts)?;
1017 if !messages
1018 .iter()
1019 .any(|message| matches!(message, BackendMessage::ParseComplete { .. }))
1020 {
1021 bail!("extended query parse did not complete");
1022 }
1023 Ok(messages)
1024 }
1025
1026 fn default_serialize_value(&self, value: &Value) -> String {
1027 Self::default_serialize_value_static(value)
1028 }
1029
1030 pub(crate) fn default_serialize_value_static(value: &Value) -> String {
1031 match value {
1032 Value::String(s) => s.clone(),
1033 Value::Number(num) => num.to_string(),
1034 Value::Bool(flag) => {
1035 if *flag {
1036 "t".to_string()
1037 } else {
1038 "f".to_string()
1039 }
1040 }
1041 _ => value.to_string(),
1042 }
1043 }
1044
1045 fn finish_query(
1046 &mut self,
1047 messages: Vec<BackendMessage>,
1048 options: Option<&QueryOptions>,
1049 ) -> Result<Results> {
1050 let blob = {
1051 let _phase = timing::phase("client.finish.blob_read");
1052 self.get_written_blob()?
1053 };
1054 {
1055 let _phase = timing::phase("client.finish.blob_cleanup");
1056 self.cleanup_blob()?;
1057 }
1058 if !self.in_transaction {
1059 let _phase = timing::phase("client.finish.sync_to_fs");
1060 self.sync_to_fs()?;
1061 }
1062 {
1063 let _phase = timing::phase("client.finish.ensure_array_types");
1064 self.ensure_array_types_for_result_messages(&messages, options)?;
1065 }
1066 let parsed = {
1067 let _phase = timing::phase("client.finish.parse_results");
1068 parse_results(&messages, &self.parsers, options, blob)
1069 };
1070 parsed
1071 .into_iter()
1072 .next()
1073 .ok_or_else(|| anyhow!("query returned no result sets"))
1074 }
1075
1076 fn finish_exec(
1077 &mut self,
1078 messages: Vec<BackendMessage>,
1079 options: Option<&QueryOptions>,
1080 has_row_description: bool,
1081 ) -> Result<Vec<Results>> {
1082 let blob = {
1083 let _phase = timing::phase("client.finish.blob_read");
1084 self.get_written_blob()?
1085 };
1086 {
1087 let _phase = timing::phase("client.finish.blob_cleanup");
1088 self.cleanup_blob()?;
1089 }
1090 if !self.in_transaction {
1091 let _phase = timing::phase("client.finish.sync_to_fs");
1092 self.sync_to_fs()?;
1093 }
1094 if has_row_description {
1095 let _phase = timing::phase("client.finish.ensure_array_types");
1096 self.ensure_array_types_for_result_messages(&messages, options)?;
1097 }
1098 let parsed = {
1099 let _phase = timing::phase("client.finish.parse_results");
1100 parse_results(&messages, &self.parsers, options, blob)
1101 };
1102 Ok(parsed)
1103 }
1104
1105 fn finish_exec_command_only(
1106 &mut self,
1107 affected_rows: Vec<usize>,
1108 options: Option<&QueryOptions>,
1109 ) -> Result<Vec<Results>> {
1110 let blob = {
1111 let _phase = timing::phase("client.finish.blob_read");
1112 self.get_written_blob()?
1113 };
1114 {
1115 let _phase = timing::phase("client.finish.blob_cleanup");
1116 self.cleanup_blob()?;
1117 }
1118 if !self.in_transaction {
1119 let _phase = timing::phase("client.finish.sync_to_fs");
1120 self.sync_to_fs()?;
1121 }
1122
1123 let _ = options;
1124 let mut results = Vec::with_capacity(affected_rows.len().max(1));
1125 for count in affected_rows {
1126 results.push(Results {
1127 rows: Vec::new(),
1128 fields: Vec::new(),
1129 affected_rows: Some(count),
1130 blob: blob.clone(),
1131 });
1132 }
1133 if results.is_empty() {
1134 results.push(Results {
1135 rows: Vec::new(),
1136 fields: Vec::new(),
1137 affected_rows: Some(0),
1138 blob,
1139 });
1140 }
1141 Ok(results)
1142 }
1143
1144 pub fn exec_protocol(
1147 &mut self,
1148 message: &[u8],
1149 options: ExecProtocolOptions,
1150 ) -> Result<ExecProtocolResult> {
1151 let ExecProtocolOptions {
1152 sync_to_fs,
1153 throw_on_error,
1154 on_notice,
1155 data_transfer_container,
1156 } = options;
1157
1158 let data = {
1159 let _phase = timing::phase("client.protocol_roundtrip");
1160 self.exec_protocol_raw_inner(message, sync_to_fs, data_transfer_container)?
1161 };
1162 self.parse_protocol_data(data, throw_on_error, on_notice)
1163 }
1164
1165 fn parse_protocol_data(
1166 &mut self,
1167 data: Vec<u8>,
1168 throw_on_error: bool,
1169 on_notice: Option<NoticeCallback>,
1170 ) -> Result<ExecProtocolResult> {
1171 let mut messages = Vec::new();
1172 let on_notice_cb = on_notice.clone();
1173 let parse_result = {
1174 let _phase = timing::phase("client.protocol_parse");
1175 self.parser.parse(&data, |msg| {
1176 if let BackendMessage::Error(db_err) = &msg
1177 && throw_on_error
1178 {
1179 return Err(anyhow!(db_err.clone()));
1180 }
1181 if let Some(callback) = on_notice_cb.as_ref()
1182 && let BackendMessage::Notice(notice) = &msg
1183 {
1184 callback(notice);
1185 }
1186 messages.push(msg);
1187 Ok(())
1188 })
1189 };
1190 if let Err(err) = parse_result {
1191 match err.downcast::<DatabaseError>() {
1192 Ok(db_err) => {
1193 self.parser = ProtocolParser::new();
1194 return Err(anyhow!(db_err));
1195 }
1196 Err(err) => return Err(err),
1197 }
1198 }
1199
1200 for message in &messages {
1201 if let BackendMessage::Notification(note) = message {
1202 if let Some(listeners) = self.notify_listeners.get(¬e.channel) {
1203 for listener in listeners {
1204 (listener.callback)(¬e.payload);
1205 }
1206 }
1207 for listener in &self.global_notify_listeners {
1208 (listener.callback)(¬e.channel, ¬e.payload);
1209 }
1210 }
1211 }
1212
1213 Ok(ExecProtocolResult { data, messages })
1214 }
1215
1216 pub fn exec_protocol_raw(
1219 &mut self,
1220 message: &[u8],
1221 options: ExecProtocolOptions,
1222 ) -> Result<Vec<u8>> {
1223 self.exec_protocol_raw_inner(message, options.sync_to_fs, options.data_transfer_container)
1224 }
1225
1226 pub fn exec_protocol_raw_stream<F>(
1229 &mut self,
1230 message: &[u8],
1231 options: ExecProtocolOptions,
1232 mut on_data: F,
1233 ) -> Result<()>
1234 where
1235 F: FnMut(&[u8]) -> Result<()>,
1236 {
1237 self.backend.send_framed_raw_stream(
1238 message,
1239 options.data_transfer_container,
1240 &mut on_data,
1241 )?;
1242 if options.sync_to_fs {
1243 let _phase = timing::phase("client.protocol_stream_sync_to_fs");
1244 self.sync_to_fs()?;
1245 }
1246 Ok(())
1247 }
1248
1249 fn exec_protocol_raw_inner(
1250 &mut self,
1251 message: &[u8],
1252 sync_to_fs: bool,
1253 data_transfer_container: Option<DataTransferContainer>,
1254 ) -> Result<Vec<u8>> {
1255 let data = {
1256 let _phase = timing::phase("client.protocol_transport_send");
1257 self.backend
1258 .send_buffered(message, data_transfer_container)?
1259 };
1260 if sync_to_fs {
1261 let _phase = timing::phase("client.protocol_sync_to_fs");
1262 self.sync_to_fs()?;
1263 }
1264 Ok(data)
1265 }
1266
1267 fn ensure_array_types_for_bind_values(
1268 &mut self,
1269 params: &[Value],
1270 data_type_ids: &[i32],
1271 options: &QueryOptions,
1272 ) -> Result<bool> {
1273 let mut registered = false;
1274 for (idx, value) in params.iter().enumerate() {
1275 if !value.is_array() {
1276 continue;
1277 }
1278 let oid = data_type_ids.get(idx).copied().unwrap_or(TEXT);
1279 if options.serializers.contains_key(&oid) || self.serializers.contains_key(&oid) {
1280 continue;
1281 }
1282 registered |= self.try_register_array_type_by_array_oid(oid)?;
1283 }
1284 Ok(registered)
1285 }
1286
1287 fn ensure_array_types_for_result_messages(
1288 &mut self,
1289 messages: &[BackendMessage],
1290 options: Option<&QueryOptions>,
1291 ) -> Result<()> {
1292 let oids = messages
1293 .iter()
1294 .filter_map(|msg| match msg {
1295 BackendMessage::RowDescription(desc) => Some(desc),
1296 _ => None,
1297 })
1298 .flat_map(|desc| desc.fields.iter().map(|field| field.data_type_id))
1299 .collect::<Vec<_>>();
1300 self.ensure_array_types_for_oids(oids, options)
1301 }
1302
1303 fn ensure_array_types_for_oids(
1304 &mut self,
1305 oids: impl IntoIterator<Item = i32>,
1306 options: Option<&QueryOptions>,
1307 ) -> Result<()> {
1308 for oid in oids {
1309 if oid <= 0 || self.parsers.contains_key(&oid) {
1310 continue;
1311 }
1312 if options.is_some_and(|options| options.parsers.contains_key(&oid)) {
1313 continue;
1314 }
1315 self.try_register_array_type_by_array_oid(oid)?;
1316 }
1317 Ok(())
1318 }
1319
1320 fn refresh_array_types_internal(&mut self) -> Result<()> {
1321 let sql = "
1322 SELECT e.oid, a.oid AS typarray, e.typdelim::text AS typdelim
1323 FROM pg_catalog.pg_type a
1324 JOIN pg_catalog.pg_type e ON e.oid = a.typelem
1325 WHERE a.typcategory = 'A'
1326 AND a.typelem <> 0
1327 ORDER BY e.oid
1328 ";
1329 let results = {
1330 let _phase = timing::phase("oliphaunt.array_type_catalog_query");
1331 self.exec_internal(sql, None)?
1332 };
1333 let result_set = results
1334 .into_iter()
1335 .next()
1336 .ok_or_else(|| anyhow!("array type discovery returned no results"))?;
1337
1338 {
1339 let _phase = timing::phase("oliphaunt.array_type_register");
1340 for row in result_set.rows {
1341 if let Some(info) = array_type_info_from_row(&row) {
1342 self.register_array_type(info);
1343 }
1344 }
1345 }
1346 Ok(())
1347 }
1348
1349 fn try_register_array_type_by_array_oid(&mut self, array_oid: i32) -> Result<bool> {
1350 if array_oid <= 0
1351 || self.parsers.contains_key(&array_oid)
1352 || self.array_type_lookup_misses.contains(&array_oid)
1353 {
1354 return Ok(false);
1355 }
1356
1357 let sql = format!(
1358 "SELECT e.oid, a.oid AS typarray, e.typdelim::text AS typdelim \
1359 FROM pg_catalog.pg_type a \
1360 JOIN pg_catalog.pg_type e ON e.oid = a.typelem \
1361 WHERE a.oid = {array_oid}::oid \
1362 AND a.typcategory = 'A' \
1363 AND a.typelem <> 0"
1364 );
1365 let results = {
1366 let _phase = timing::phase("oliphaunt.array_type_targeted_lookup");
1367 self.exec_internal(&sql, None)?
1368 };
1369 let Some(result_set) = results.into_iter().next() else {
1370 self.array_type_lookup_misses.insert(array_oid);
1371 return Ok(false);
1372 };
1373 let Some(row) = result_set.rows.into_iter().next() else {
1374 self.array_type_lookup_misses.insert(array_oid);
1375 return Ok(false);
1376 };
1377 let Some(info) = array_type_info_from_row(&row) else {
1378 self.array_type_lookup_misses.insert(array_oid);
1379 return Ok(false);
1380 };
1381
1382 self.register_array_type(info);
1383 Ok(true)
1384 }
1385
1386 fn register_array_type(&mut self, info: ArrayTypeInfo) {
1387 register_array_type(&mut self.parsers, &mut self.serializers, info);
1388 self.array_type_lookup_misses.remove(&info.array_oid);
1389 }
1390
1391 fn run_exec_command(&mut self, sql: &str) -> Result<()> {
1392 self.exec_internal(sql, None).map(|_| ())
1393 }
1394
1395 fn handle_blob_input(&mut self, blob: Option<&Vec<u8>>) -> Result<()> {
1396 let path = self.dev_blob_path();
1397 if let Some(bytes) = blob {
1398 if let Some(parent) = path.parent() {
1399 fs::create_dir_all(parent).with_context(|| {
1400 format!("failed to create blob directory {}", parent.display())
1401 })?;
1402 }
1403 fs::write(&path, bytes)
1404 .with_context(|| format!("write blob input to {}", path.display()))?;
1405 self.blob_input_provided = true;
1406 } else {
1407 self.blob_input_provided = false;
1408 let _ = fs::remove_file(&path);
1409 }
1410 Ok(())
1411 }
1412
1413 fn dev_blob_path(&self) -> PathBuf {
1414 self.backend.paths().runtime_root().join("dev/blob")
1415 }
1416
1417 fn cleanup_blob(&mut self) -> Result<()> {
1418 Ok(())
1419 }
1420
1421 fn get_written_blob(&mut self) -> Result<Option<Vec<u8>>> {
1422 let path = self.dev_blob_path();
1423
1424 if self.blob_input_provided {
1425 self.blob_input_provided = false;
1426 let _ = fs::remove_file(&path);
1427 return Ok(None);
1428 }
1429
1430 match fs::read(&path) {
1431 Ok(data) => {
1432 self.blob_input_provided = false;
1433 let _ = fs::remove_file(&path);
1434 if data.is_empty() {
1435 Ok(None)
1436 } else {
1437 Ok(Some(data))
1438 }
1439 }
1440 Err(err) => {
1441 if err.kind() == io::ErrorKind::NotFound {
1442 self.blob_input_provided = false;
1443 Ok(None)
1444 } else {
1445 Err(err).with_context(|| format!("read blob output from {}", path.display()))
1446 }
1447 }
1448 }
1449 }
1450
1451 fn check_ready(&self) -> Result<()> {
1452 if self.closing {
1453 bail!("Oliphaunt instance is closing");
1454 }
1455 if self.closed {
1456 bail!("Oliphaunt instance is closed");
1457 }
1458 if !self.ready {
1459 bail!("Oliphaunt instance is not ready");
1460 }
1461 Ok(())
1462 }
1463}
1464
1465impl Drop for Oliphaunt {
1466 fn drop(&mut self) {
1467 if !self.closed {
1468 let _ = self.close();
1469 }
1470 }
1471}
1472
1473#[cfg(feature = "tools")]
1474fn ensure_direct_pg_dump_options_match_session(
1475 startup_config: &StartupConfig,
1476 options: &PgDumpOptions,
1477) -> Result<()> {
1478 if options.database_ref() != startup_config.database {
1479 bail!(
1480 "direct pg_dump runs against the already-open embedded backend database '{}'; requested database '{}' would require a separate server connection",
1481 startup_config.database,
1482 options.database_ref()
1483 );
1484 }
1485 if options.username_ref() != startup_config.username {
1486 bail!(
1487 "direct pg_dump runs through the already-open embedded backend user '{}'; requested user '{}' would require a separate server connection",
1488 startup_config.username,
1489 options.username_ref()
1490 );
1491 }
1492 Ok(())
1493}
1494
1495#[cfg(feature = "tools")]
1496fn read_direct_pg_dump_socket(
1497 runtime: &Runtime,
1498 reader: &mut TcpSocketHalfRx,
1499 buffer: &mut [u8],
1500) -> Result<usize> {
1501 runtime
1502 .block_on(async {
1503 std::future::poll_fn(|cx| {
1504 let read = match reader.poll_fill_buf(cx) {
1505 std::task::Poll::Ready(Ok(available)) => {
1506 let read = available.len().min(buffer.len());
1507 buffer[..read].copy_from_slice(&available[..read]);
1508 read
1509 }
1510 std::task::Poll::Ready(Err(err)) => return std::task::Poll::Ready(Err(err)),
1511 std::task::Poll::Pending => return std::task::Poll::Pending,
1512 };
1513 reader.consume(read);
1514 std::task::Poll::Ready(Ok(read))
1515 })
1516 .await
1517 })
1518 .context("read direct pg_dump virtual socket")
1519}
1520
1521#[cfg(feature = "tools")]
1522fn write_direct_pg_dump_socket(
1523 runtime: &Runtime,
1524 writer: &mut (impl AsyncWrite + Unpin),
1525 bytes: &[u8],
1526) -> Result<()> {
1527 runtime
1528 .block_on(writer.write_all(bytes))
1529 .context("write direct pg_dump virtual socket")
1530}
1531
1532#[cfg(feature = "tools")]
1533fn flush_direct_pg_dump_socket(
1534 runtime: &Runtime,
1535 writer: &mut (impl AsyncWrite + Unpin),
1536) -> Result<()> {
1537 runtime
1538 .block_on(writer.flush())
1539 .context("flush direct pg_dump virtual socket")
1540}
1541
1542fn parse_command_only_result_counts(data: &[u8]) -> Option<Vec<usize>> {
1543 let mut offset = 0usize;
1544 let mut affected_total = 0usize;
1545 let mut affected_rows = Vec::new();
1546 while offset + 5 <= data.len() {
1547 let tag = data[offset];
1548 let length = u32::from_be_bytes([
1549 data[offset + 1],
1550 data[offset + 2],
1551 data[offset + 3],
1552 data[offset + 4],
1553 ]) as usize;
1554 if length < 4 {
1555 return None;
1556 }
1557 let frame_len = 1 + length;
1558 if frame_len > data.len() - offset {
1559 return None;
1560 }
1561 let body_start = offset + 5;
1562 let body_end = offset + frame_len;
1563 match tag {
1564 b'C' => {
1565 let command_tag = data[body_start..body_end]
1566 .strip_suffix(&[0])
1567 .unwrap_or(&data[body_start..body_end]);
1568 affected_total = affected_total.saturating_add(command_tag_row_count(command_tag));
1569 affected_rows.push(affected_total);
1570 }
1571 b'Z' => {}
1572 _ => return None,
1573 }
1574 offset += frame_len;
1575 }
1576 (offset == data.len()).then_some(affected_rows)
1577}
1578
1579fn value_to_i32(value: Option<&Value>) -> Option<i32> {
1580 match value? {
1581 Value::Number(number) => number.as_i64().map(|value| value as i32),
1582 Value::String(string) => string.parse::<i32>().ok(),
1583 _ => None,
1584 }
1585}
1586
1587fn value_to_char(value: Option<&Value>) -> Option<char> {
1588 match value? {
1589 Value::String(string) => string.chars().next(),
1590 _ => None,
1591 }
1592}
1593
1594fn array_type_info_from_row(row: &Value) -> Option<ArrayTypeInfo> {
1595 let Value::Object(map) = row else {
1596 return None;
1597 };
1598 let element_oid = value_to_i32(map.get("oid"))?;
1599 let array_oid = value_to_i32(map.get("typarray"))?;
1600 if element_oid == 0 || array_oid == 0 {
1601 return None;
1602 }
1603 let delimiter = value_to_char(map.get("typdelim")).unwrap_or(',');
1604 Some(ArrayTypeInfo::new(element_oid, array_oid, delimiter))
1605}
1606
1607pub struct Transaction<'a> {
1609 client: &'a mut Oliphaunt,
1610 closed: bool,
1611}
1612
1613impl<'a> Transaction<'a> {
1614 fn new(client: &'a mut Oliphaunt) -> Self {
1615 Self {
1616 client,
1617 closed: false,
1618 }
1619 }
1620
1621 fn commit_internal(&mut self) -> Result<()> {
1622 self.ensure_open()?;
1623 self.client.exec_internal("COMMIT", None)?;
1624 self.closed = true;
1625 Ok(())
1626 }
1627
1628 fn rollback_internal(&mut self) -> Result<()> {
1629 self.ensure_open()?;
1630 self.client.exec_internal("ROLLBACK", None)?;
1631 self.closed = true;
1632 Ok(())
1633 }
1634
1635 fn ensure_open(&self) -> Result<()> {
1636 if self.closed {
1637 bail!("transaction is already closed");
1638 }
1639 Ok(())
1640 }
1641
1642 pub fn query(
1643 &mut self,
1644 sql: &str,
1645 params: &[Value],
1646 options: Option<&QueryOptions>,
1647 ) -> Result<Results> {
1648 self.ensure_open()?;
1649 self.client.query_internal(sql, params, options)
1650 }
1651
1652 pub fn exec(&mut self, sql: &str, options: Option<&QueryOptions>) -> Result<Vec<Results>> {
1653 self.ensure_open()?;
1654 self.client.exec_internal(sql, options)
1655 }
1656
1657 pub fn refresh_array_types(&mut self) -> Result<()> {
1658 self.ensure_open()?;
1659 self.client.refresh_array_types_internal()
1660 }
1661
1662 pub fn commit(&mut self) -> Result<()> {
1663 self.commit_internal()
1664 }
1665
1666 pub fn rollback(&mut self) -> Result<()> {
1667 self.rollback_internal()
1668 }
1669
1670 pub fn is_closed(&self) -> bool {
1671 self.closed
1672 }
1673
1674 pub fn closed(&self) -> bool {
1675 self.closed
1676 }
1677}