Skip to main content

oliphaunt_wasix/oliphaunt/
client.rs

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
103/// Primary entry point for interacting with the embedded Postgres runtime.
104pub 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    /// Create a builder for opening persistent or temporary Oliphaunt databases.
125    pub fn builder() -> OliphauntBuilder {
126        OliphauntBuilder::new()
127    }
128
129    /// Open a persistent Oliphaunt database rooted at `root`, installing and initializing it if needed.
130    pub fn open(root: impl AsRef<Path>) -> Result<Self> {
131        Self::builder().path(root.as_ref().to_path_buf()).open()
132    }
133
134    /// Open a persistent Oliphaunt database under the platform data directory for `app_id`.
135    pub fn open_app(app_id: (&str, &str, &str)) -> Result<Self> {
136        Self::builder().app_id(app_id).open()
137    }
138
139    /// Create an ephemeral Oliphaunt database whose files are removed when the instance is dropped.
140    pub fn temporary() -> Result<Self> {
141        Self::builder().temporary().open()
142    }
143
144    /// Warm the runtime module and bundled AOT artifact cache without opening a database.
145    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    /// Warm bundled extension artifacts without permanently opening a database.
163    #[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    /// Create a new Oliphaunt instance backed by the provided runtime paths.
200    #[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    /// Install and enable a bundled Postgres extension.
263    #[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    /// Refresh direct API array parser and serializer registrations.
291    ///
292    /// This mirrors upstream Oliphaunt's `refreshArrayTypes()` escape hatch. Most
293    /// applications should not need it because built-in arrays are registered
294    /// statically and runtime custom arrays are discovered lazily when possible.
295    pub fn refresh_array_types(&mut self) -> Result<()> {
296        self.check_ready()?;
297        self.refresh_array_types_internal()
298    }
299
300    /// Execute a SQL query using the extended protocol.
301    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    /// Return `true` if the instance is ready for new work.
397    pub fn is_ready(&self) -> bool {
398        self.ready && !self.closing && !self.closed
399    }
400
401    /// Return the capabilities of the active embedded engine.
402    pub fn engine_capabilities(&self) -> EngineCapabilities {
403        self.backend.capabilities()
404    }
405
406    /// Return the host-side runtime and data-directory paths backing this instance.
407    #[doc(hidden)]
408    pub fn paths(&self) -> &OliphauntPaths {
409        self.backend.paths()
410    }
411
412    /// Return debug-build bridge allocation/free counters for ownership tests.
413    #[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    /// Dump the physical PGDATA directory to a gzipped tar archive.
420    ///
421    /// The archive is intended to be loaded back into oliphaunt-wasix/Oliphaunt with
422    /// the same PostgreSQL/Oliphaunt version. Use [`dump_sql`](Self::dump_sql) for
423    /// logical backups across versions.
424    pub fn dump_data_dir(&mut self) -> Result<Vec<u8>> {
425        self.dump_data_dir_with_format(DataDirArchiveFormat::TarGz)
426    }
427
428    /// Dump the physical PGDATA directory with the selected archive format.
429    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    /// Clone this database into a new temporary [`Oliphaunt`] instance.
435    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    /// Run the bundled WASIX `pg_dump` against this database and return SQL text.
446    #[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    /// Run the bundled WASIX `pg_dump` and return UTF-8 SQL bytes.
455    #[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    /// Return `true` if the instance has already been closed.
640    pub fn is_closed(&self) -> bool {
641        self.closed
642    }
643
644    /// Shut down the embedded Postgres runtime.
645    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    /// Execute a simple SQL statement that may contain multiple commands.
680    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    /// Register a listener for `LISTEN channel`. Returns a handle that can be used to unlisten.
752    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    /// Remove a listener corresponding to the provided handle.
783    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    /// Remove all listeners for the specified channel.
797    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    /// Register a global notification callback.
807    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    /// Deregister a previously registered global notification callback.
820    pub fn off_notification(&mut self, handle: GlobalListenerHandle) {
821        self.global_notify_listeners
822            .retain(|listener| listener.id != handle.id);
823    }
824
825    /// Describe the parameter and result metadata for a SQL query.
826    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    /// Run a closure within an SQL transaction (`BEGIN .. COMMIT/ROLLBACK`).
923    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        // Begin transaction
930        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    /// Flush runtime writes to the underlying filesystem.
956    ///
957    /// The WASIX backend uses host-mounted files and PostgreSQL's own fsync/WAL
958    /// behavior for durability. Adding an unconditional host directory
959    /// `sync_all` after every direct query is both expensive and weaker than the
960    /// database's file-level fsyncs, so the Rust-level hook remains a no-op.
961    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    /// Execute raw PostgreSQL frontend protocol bytes and parse backend
1145    /// protocol messages.
1146    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(&note.channel) {
1203                    for listener in listeners {
1204                        (listener.callback)(&note.payload);
1205                    }
1206                }
1207                for listener in &self.global_notify_listeners {
1208                    (listener.callback)(&note.channel, &note.payload);
1209                }
1210            }
1211        }
1212
1213        Ok(ExecProtocolResult { data, messages })
1214    }
1215
1216    /// Execute raw PostgreSQL frontend protocol bytes and return raw backend
1217    /// protocol bytes.
1218    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    /// Execute raw protocol bytes and pass the returned backend bytes to
1227    /// `on_data`.
1228    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
1607/// Transaction handle used within [`Oliphaunt::transaction`].
1608pub 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}