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