zippa-db 0.1.0

A fast, lightweight, cross-platform database client for PostgreSQL, MySQL, and SQLite.
//! What the session knows about the server: the database list, the sidebar's
//! objects, the schema-search catalog, and switching database.

use std::sync::Arc;

use gpui_kit::component::button::Button;
use gpui_kit::component::button::ButtonVariants as _;
use gpui_kit::component::menu::{DropdownMenu, PopupMenuItem};
use gpui_kit::component::{Disableable, Sizable};
use gpui_kit::prelude::*;
use gpui_kit::{Context, Entity, px};

use crate::db::{Catalog, CatalogEntry, Connection, runtime};

use super::{Session, SessionEvent, Status};

impl Session {
    /// Read the database list and the current database's tables and views.
    pub(crate) fn reload_metadata(&mut self, cx: &mut Context<Self>) {
        let connection = self.connection.clone();
        let started_on = connection.clone();
        let task = runtime::spawn(async move {
            (
                connection.databases().await,
                connection.objects().await,
                connection.stored_objects().await,
                connection.catalog().await,
            )
        });

        cx.spawn(async move |this, cx| {
            let loaded = task.await;
            this.update_in(cx, |this, window, cx| {
                // `switch_database` replaces `self.connection` and closes the
                // one this task started on, then kicks off its own reload. A
                // read still in flight against the old connection at that
                // point loses its pool mid-await and comes back an error;
                // that error is stale (the fresh reload already has, or will
                // have, the right answer), so it is dropped rather than
                // shown.
                if !Arc::ptr_eq(&this.connection, &started_on) {
                    return;
                }
                this.metadata_error = None;
                match loaded {
                    Ok((databases, objects, stored, catalog)) => {
                        match databases {
                            Ok(databases) => this.databases = databases,
                            Err(error) => this.metadata_error = Some(format!("{error:#}")),
                        }
                        match objects {
                            Ok(objects) => this.objects = objects,
                            Err(error) => this.metadata_error = Some(format!("{error:#}")),
                        }
                        match stored {
                            Ok(stored) => this.stored = stored,
                            Err(error) => this.metadata_error = Some(format!("{error:#}")),
                        }
                        this.store_catalog(catalog);
                    }
                    Err(_) => {
                        this.metadata_error = Some("reading the schema was cancelled".into());
                        this.catalog_loading = false;
                    }
                }
                if let Some(error) = &this.metadata_error {
                    crate::ui::notify_error(window, cx, format!("Error: {error}"));
                }
                this.rebuild_tree(cx);
                cx.notify();
            })
            .ok();
        })
        .detach();
    }

    /// Keep the catalog read's answer, falling back to the names the sidebar
    /// already has when the full read failed — a database whose column query
    /// times out should still be searchable by name.
    fn store_catalog(&mut self, catalog: Result<Catalog, anyhow::Error>) {
        self.catalog_loading = false;
        match catalog {
            Ok(catalog) => {
                self.catalog = Arc::new(catalog);
                self.catalog_error = None;
            }
            Err(error) => {
                let mut fallback = Catalog::default();
                fallback.entries = self
                    .objects
                    .iter()
                    .cloned()
                    .map(CatalogEntry::object)
                    .chain(self.stored.iter().cloned().map(CatalogEntry::routine))
                    .collect();
                fallback.total = fallback.entries.len();
                self.catalog = Arc::new(fallback);
                self.catalog_error = Some(format!("{error:#}"));
            }
        }
    }

    /// The whole schema, for [`SearchSchema`].
    pub(crate) fn catalog(&self) -> Arc<Catalog> {
        self.catalog.clone()
    }

    pub(crate) fn catalog_loading(&self) -> bool {
        self.catalog_loading
    }

    pub(crate) fn catalog_error(&self) -> Option<&str> {
        self.catalog_error.as_deref()
    }

    /// Reopen the pool against `database` and reload the object list.
    ///
    /// No engine can move an open pool to another database, so this replaces
    /// the connection and drains the old one in the background.
    pub(crate) fn switch_database(&mut self, database: String, cx: &mut Context<Self>) {
        // A file-based engine has a single database; "switching" would try to
        // open a file named after it.
        if self.connection.config.engine.is_file_based() {
            return;
        }
        if self.switching || database == self.connection.database() {
            return;
        }

        self.switching = true;
        cx.notify();

        let connection = self.connection.clone();
        let task = runtime::spawn(async move { connection.with_database(&database).await });

        cx.spawn(async move |this, cx| {
            let opened = task.await;
            this.update_in(cx, |this, window, cx| {
                this.switching = false;
                match opened {
                    Ok(Ok(connection)) => {
                        this.adopt_connection(connection, cx);
                        cx.emit(SessionEvent::Changed);
                    }
                    Ok(Err(error)) => {
                        let message = format!("{error:#}");
                        if let Some(panel) = this.active_panel() {
                            panel.update(cx, |panel, _| {
                                panel.set_status(Status::Error(message.clone()))
                            });
                        }
                        crate::ui::notify_error(window, cx, format!("Error: {message}"));
                    }
                    Err(_) => {
                        let message = "switching database was cancelled";
                        if let Some(panel) = this.active_panel() {
                            panel.update(cx, |panel, _| {
                                panel.set_status(Status::Error(message.into()))
                            });
                        }
                        crate::ui::notify_error(window, cx, format!("Error: {message}"));
                    }
                }
                cx.notify();
            })
            .ok();
        })
        .detach();
    }

    /// Swap `self.connection` for a freshly opened one on another database:
    /// close what it replaces, carry its console history forward, reset the
    /// schema read, and point every open panel at the new connection.
    ///
    /// Shared by [`Self::switch_database`] and a restore that switches before
    /// building any tabs, so both keep the panels, the catalog, and the
    /// console in step the same way.
    pub(super) fn adopt_connection(&mut self, connection: Connection, cx: &mut Context<Self>) {
        let previous = std::mem::replace(&mut self.connection, Arc::new(connection));

        // The console's history belongs to the tab, not the pool about to
        // close: an internal read that failed right before the switch (or
        // any read from before it) should still be there to look at
        // afterward rather than vanish with the connection that logged it.
        for entry in previous.query_log().snapshot() {
            self.connection.query_log().record(entry);
        }
        runtime::spawn(async move { previous.close().await });

        self.objects.clear();
        self.stored.clear();
        self.catalog = Arc::new(Catalog::default());
        self.catalog_loading = true;
        self.catalog_error = None;
        self.rebuild_tree(cx);
        let connection = self.connection.clone();
        for panel in self.panels.clone() {
            let connection = connection.clone();
            panel.update(cx, |panel, cx| panel.set_connection(connection, cx));
        }
        self.reload_metadata(cx);
    }

    /// The database dropdown, rendered by whoever owns the toolbar.
    ///
    /// Sized and styled here rather than by the caller: `dropdown_menu`
    /// wraps the button in a popover that does not itself implement
    /// [`Sizable`]/[`ButtonVariants`].
    pub(crate) fn render_database_picker(
        session: &Entity<Session>,
        cx: &mut gpui_kit::App,
    ) -> impl IntoElement {
        let this = session.read(cx);
        let current = this.connection.database().to_string();

        // A SQLite connection is one file: there is nothing to switch to, and
        // its "databases" (main, plus attachments) are not separate files.
        if this.connection.config.engine.is_file_based() {
            return Button::new("database")
                .custom(crate::ui::subtle_button(cx))
                .xsmall()
                .max_w(px(160.))
                .icon(gpui_kit::assets::IconName::Database)
                .label(crate::db::file_name(&current))
                .disabled(true)
                .into_any_element();
        }

        let databases = this.databases.clone();
        let weak = session.downgrade();

        let label = if this.switching {
            "Switching…".to_string()
        } else if current.is_empty() {
            "No database".to_string()
        } else {
            current.clone()
        };

        Button::new("database")
            .custom(crate::ui::subtle_button(cx))
            .xsmall()
            .max_w(px(160.))
            .icon(gpui_kit::assets::IconName::Database)
            .label(label)
            .dropdown_menu(move |mut menu, _window, _cx| {
                if databases.is_empty() {
                    return menu.label("No databases");
                }

                for database in &databases {
                    let name = database.clone();
                    let weak = weak.clone();

                    menu = menu.item(
                        PopupMenuItem::new(database.clone())
                            .checked(*database == current)
                            .on_click(move |_, _window, cx| {
                                let name = name.clone();
                                if let Some(session) = weak.upgrade() {
                                    session.update(cx, |session, cx| {
                                        session.switch_database(name, cx)
                                    });
                                }
                            }),
                    );
                }

                menu.scrollable(true).max_h(px(420.))
            })
            .into_any_element()
    }
}