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, 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
104/// Primary entry point for interacting with the embedded Postgres runtime.
105pub 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    /// Create a builder for opening persistent or temporary Oliphaunt databases.
126    pub fn builder() -> OliphauntBuilder {
127        OliphauntBuilder::new()
128    }
129
130    /// Open a persistent Oliphaunt database rooted at `root`, installing and initializing it if needed.
131    pub fn open(root: impl AsRef<Path>) -> Result<Self> {
132        Self::builder().path(root.as_ref().to_path_buf()).open()
133    }
134
135    /// Open a persistent Oliphaunt database under the platform data directory for `app_id`.
136    pub fn open_app(app_id: (&str, &str, &str)) -> Result<Self> {
137        Self::builder().app_id(app_id).open()
138    }
139
140    /// Create an ephemeral Oliphaunt database whose files are removed when the instance is dropped.
141    pub fn temporary() -> Result<Self> {
142        Self::builder().temporary().open()
143    }
144
145    /// Warm the runtime module and bundled AOT artifact cache without opening a database.
146    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    /// Warm bundled extension artifacts without permanently opening a database.
164    #[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    /// Create a new Oliphaunt instance backed by the provided runtime paths.
201    #[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    /// Install and enable a bundled Postgres extension.
307    ///
308    /// Extensions that require startup configuration must instead be selected
309    /// on [`OliphauntBuilder`](crate::OliphauntBuilder) before the database is
310    /// opened. This method returns an actionable error when the running backend
311    /// does not already satisfy such a requirement.
312    #[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    /// Refresh direct API array parser and serializer registrations.
337    ///
338    /// This mirrors upstream Oliphaunt's `refreshArrayTypes()` escape hatch. Most
339    /// applications should not need it because built-in arrays are registered
340    /// statically and runtime custom arrays are discovered lazily when possible.
341    pub fn refresh_array_types(&mut self) -> Result<()> {
342        self.check_ready()?;
343        self.refresh_array_types_internal()
344    }
345
346    /// Execute a SQL query using the extended protocol.
347    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    /// Return `true` if the instance is ready for new work.
443    pub fn is_ready(&self) -> bool {
444        self.ready && !self.closing && !self.closed
445    }
446
447    /// Return the capabilities of the active embedded engine.
448    pub fn engine_capabilities(&self) -> EngineCapabilities {
449        self.backend.capabilities()
450    }
451
452    /// Return the host-side runtime and data-directory paths backing this instance.
453    #[doc(hidden)]
454    pub fn paths(&self) -> &OliphauntPaths {
455        self.backend.paths()
456    }
457
458    /// Return debug-build bridge allocation/free counters for ownership tests.
459    #[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    /// Dump the physical PGDATA directory to a gzipped tar archive.
466    ///
467    /// The archive is intended to be loaded back into oliphaunt-wasix/Oliphaunt with
468    /// the same PostgreSQL/Oliphaunt version. Use [`dump_sql`](Self::dump_sql) for
469    /// logical backups across versions.
470    pub fn dump_data_dir(&mut self) -> Result<Vec<u8>> {
471        self.dump_data_dir_with_format(DataDirArchiveFormat::TarGz)
472    }
473
474    /// Dump the physical PGDATA directory with the selected archive format.
475    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    /// Clone this database into a new temporary [`Oliphaunt`] instance.
481    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    /// Run the bundled WASIX `pg_dump` against this database and return SQL text.
492    #[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    /// Run the bundled WASIX `pg_dump` and return UTF-8 SQL bytes.
501    #[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    /// Return `true` if the instance has already been closed.
686    pub fn is_closed(&self) -> bool {
687        self.closed
688    }
689
690    /// Shut down the embedded Postgres runtime.
691    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    /// Execute a simple SQL statement that may contain multiple commands.
726    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    /// Register a listener for `LISTEN channel`. Returns a handle that can be used to unlisten.
798    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    /// Remove a listener corresponding to the provided handle.
829    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    /// Remove all listeners for the specified channel.
843    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    /// Register a global notification callback.
853    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    /// Deregister a previously registered global notification callback.
866    pub fn off_notification(&mut self, handle: GlobalListenerHandle) {
867        self.global_notify_listeners
868            .retain(|listener| listener.id != handle.id);
869    }
870
871    /// Describe the parameter and result metadata for a SQL query.
872    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    /// Run a closure within an SQL transaction (`BEGIN .. COMMIT/ROLLBACK`).
969    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        // Begin transaction
976        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    /// Flush runtime writes to the underlying filesystem.
1002    ///
1003    /// The WASIX backend uses host-mounted files and PostgreSQL's own fsync/WAL
1004    /// behavior for durability. Adding an unconditional host directory
1005    /// `sync_all` after every direct query is both expensive and weaker than the
1006    /// database's file-level fsyncs, so the Rust-level hook remains a no-op.
1007    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    /// Execute raw PostgreSQL frontend protocol bytes and parse backend
1191    /// protocol messages.
1192    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(&note.channel) {
1249                    for listener in listeners {
1250                        (listener.callback)(&note.payload);
1251                    }
1252                }
1253                for listener in &self.global_notify_listeners {
1254                    (listener.callback)(&note.channel, &note.payload);
1255                }
1256            }
1257        }
1258
1259        Ok(ExecProtocolResult { data, messages })
1260    }
1261
1262    /// Execute raw PostgreSQL frontend protocol bytes and return raw backend
1263    /// protocol bytes.
1264    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    /// Execute raw protocol bytes and pass the returned backend bytes to
1273    /// `on_data`.
1274    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
1653/// Transaction handle used within [`Oliphaunt::transaction`].
1654pub 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}