uqa-engine 0.1.9

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! File-format opening, connection binding, and independent session creation.

use super::{
    Arc, DeepModel, Engine, ManagedConnection, Path, PersistentStorageBackend,
    PersistentStorageProvider, PersistentStorageSession, SQLiteCompressedContainerAnchor,
    SQLiteCompressionOptions, SQLiteError, SQLiteStorageProvider, StorageBackendError,
    StorageBackendResult,
};

struct BackendSessionProvider {
    backend: Arc<dyn PersistentStorageBackend>,
}

impl PersistentStorageProvider for BackendSessionProvider {
    fn open_session(&self) -> StorageBackendResult<PersistentStorageSession> {
        self.backend.open_session()
    }

    fn storage_identity(
        &self,
    ) -> StorageBackendResult<Option<uqa_storage::PersistentStorageIdentity>> {
        self.backend.storage_identity()
    }
}

impl Engine {
    pub fn open(path: &Path) -> Result<Self, SQLiteError> {
        let conn = ManagedConnection::open(path)?;
        Self::open_with_connection(&conn)
    }

    /// Classify the on-disk format of `path` without opening it: plain
    /// `SQLite`, UQA compressed container (with its encryption flag), a
    /// missing/empty file, or an unrecognized header (`SQLCipher`
    /// databases fall here because `SQLCipher` encrypts the whole file).
    pub fn detect_database_file(path: &Path) -> std::io::Result<uqa_storage::DatabaseFileFormat> {
        uqa_storage::detect_database_file_format(path)
    }

    /// Open `path` with the variant its on-disk format calls for.
    ///
    /// - Missing/empty file: creates a new database, `SQLCipher`
    ///   encrypted when `key` is provided, plaintext otherwise.
    /// - Plain `SQLite`: opens plaintext; providing a key is an error
    ///   ([`SQLiteError::NotEncrypted`]) rather than a silent no-op so
    ///   callers never believe an unencrypted database is protected.
    /// - Compressed container: opens with the codec recorded in the
    ///   container header; the encryption flag decides whether `key`
    ///   is required ([`SQLiteError::EncryptionKeyRequired`]) or
    ///   rejected ([`SQLiteError::NotEncrypted`]).
    /// - Unrecognized header: treated as `SQLCipher` when `key` is
    ///   provided; without a key this fails with
    ///   [`SQLiteError::EncryptionKeyRequired`] because an encrypted
    ///   database cannot be told apart from a foreign file.
    ///
    /// New compressed containers are not created through this entry
    /// point; use [`Engine::open_compressed`] or
    /// [`Engine::open_compressed_encrypted`] to choose compression for
    /// a new database.
    pub fn open_auto(path: &Path, key: Option<&str>) -> Result<Self, SQLiteError> {
        use uqa_storage::DatabaseFileFormat;
        let key = match key {
            Some("") => return Err(SQLiteError::EmptyEncryptionKey),
            other => other,
        };
        match uqa_storage::detect_database_file_format(path)? {
            DatabaseFileFormat::Missing => match key {
                Some(key) => Self::open_encrypted(path, key),
                None => Self::open(path),
            },
            DatabaseFileFormat::PlainSQLite => match key {
                Some(_) => Err(SQLiteError::NotEncrypted),
                None => Self::open(path),
            },
            DatabaseFileFormat::CompressedContainer { encrypted: true } => match key {
                Some(key) => {
                    Self::open_compressed_encrypted(path, key, SQLiteCompressionOptions::default())
                }
                None => Err(SQLiteError::EncryptionKeyRequired),
            },
            DatabaseFileFormat::CompressedContainer { encrypted: false } => match key {
                Some(_) => Err(SQLiteError::NotEncrypted),
                None => Self::open_compressed(path, SQLiteCompressionOptions::default()),
            },
            DatabaseFileFormat::Unrecognized => match key {
                Some(key) => Self::open_encrypted(path, key),
                None => Err(SQLiteError::EncryptionKeyRequired),
            },
        }
    }

    /// SQLCipher-backed engine. Applies `key` before any catalog
    /// access, runs migrations, and rebuilds the in-memory table
    /// registry from the encrypted catalog.
    pub fn open_encrypted(path: &Path, key: &str) -> Result<Self, SQLiteError> {
        let conn = ManagedConnection::open_encrypted(path, key)?;
        Self::open_with_connection(&conn)
    }

    /// Compressed SQLite-backed engine. The compression VFS is
    /// schema-neutral: it compresses `SQLite` byte ranges in chunks
    /// without knowledge of UQA catalog tables or columns.
    pub fn open_compressed(
        path: &Path,
        compression: SQLiteCompressionOptions,
    ) -> Result<Self, SQLiteError> {
        let conn = ManagedConnection::open_compressed(path, compression)?;
        Self::open_with_connection(&conn)
    }

    /// Compressed and encrypted SQLite-backed engine. Chunk payloads
    /// are compressed first, then encrypted by the compressed VFS. The v2
    /// format authenticates container metadata, chunk placement, and commit
    /// records, but cannot distinguish replacement by an internally valid
    /// snapshot or fork without an external trusted state anchor.
    /// Security-sensitive deployments that do not require compression should
    /// prefer [`Engine::open_encrypted`] and `SQLCipher`.
    pub fn open_compressed_encrypted(
        path: &Path,
        key: &str,
        compression: SQLiteCompressionOptions,
    ) -> Result<Self, SQLiteError> {
        let conn = ManagedConnection::open_compressed_encrypted(path, key, compression)?;
        Self::open_with_connection(&conn)
    }

    /// Open an encrypted compressed database and reject a different file or
    /// any state other than `trusted_anchor` before `SQLite` reads the main
    /// database. Refresh the trusted anchor after every committed write.
    pub fn open_compressed_encrypted_with_anchor(
        path: &Path,
        key: &str,
        compression: SQLiteCompressionOptions,
        trusted_anchor: SQLiteCompressedContainerAnchor,
    ) -> Result<Self, SQLiteError> {
        let conn = ManagedConnection::open_compressed_encrypted_with_anchor(
            path,
            key,
            compression,
            trusted_anchor,
        )?;
        Self::open_with_connection(&conn)
    }

    /// Authenticate and return the anchor to persist in a trusted store after
    /// committed writes to an encrypted compressed database.
    pub fn compressed_container_anchor(
        path: &Path,
        key: &str,
    ) -> Result<SQLiteCompressedContainerAnchor, SQLiteError> {
        Ok(uqa_storage::read_authenticated_anchor(path, key)?)
    }

    fn open_with_connection(conn: &ManagedConnection) -> Result<Self, SQLiteError> {
        let provider: Arc<dyn PersistentStorageProvider> =
            Arc::new(SQLiteStorageProvider::new(conn.clone()));
        Self::from_persistent_provider(provider).map_err(Self::sqlite_open_error)
    }

    /// Create an independent SQL session over this engine's durable database.
    ///
    /// The new session gets its own catalog/backend pair, transaction stack,
    /// runtime variables, prepared statements, statement cache, and
    /// cancellation token. Durable registry caches remain session-private and
    /// synchronize through shared epochs; runtime-only Rust extensions are
    /// shared. The provider must return catalog and data handles bound to one
    /// session transaction so every durable mutation commits atomically.
    pub fn new_session(&self) -> StorageBackendResult<Self> {
        let provider = self.storage.provider.as_ref().ok_or_else(|| {
            StorageBackendError::Other(
                "independent sessions require a PersistentStorageProvider".into(),
            )
        })?;
        let storage_session = provider.open_session()?;
        let mut session =
            Self::from_persistent_session(storage_session, Some(Arc::clone(provider)))?;
        session.row_locks = Arc::clone(&self.row_locks);
        session.session_id = self.row_locks.allocate_session();
        session.epochs.share_published_from(&self.epochs);
        // Force one catalog rebind after attaching the shared generation.
        // Otherwise a DDL commit racing the initial restore could leave this
        // session with the old table snapshot but the new generation marked
        // as already observed.
        // Durable registries remain session-local. Sharing these maps
        // would expose a writer's uncommitted graph/schema/view/FDW changes
        // to sibling sessions before storage COMMIT. Runtime-only registries
        // may remain shared.
        session.extensions = super::RuntimeExtensions::shared_from(&self.extensions);
        session.synchronize_table_catalog()?;
        session.synchronize_table_data()?;
        session.synchronize_catalog_registries()?;
        Ok(session)
    }

    /// Build an engine and retain the provider used to create future
    /// independent sessions.
    #[allow(clippy::needless_pass_by_value)]
    pub fn from_persistent_provider(
        provider: Arc<dyn PersistentStorageProvider>,
    ) -> StorageBackendResult<Self> {
        let identity = provider.storage_identity()?;
        let session = provider.open_session()?;
        let mut engine = Self::from_persistent_session(session, Some(Arc::clone(&provider)))?;
        let row_locks = crate::row_locks::shared_provider_manager(identity, &provider);
        engine.session_id = row_locks.allocate_session();
        engine.row_locks = row_locks;
        Ok(engine)
    }

    /// Build an engine from already-open persistent metadata and data backends. The backend's session factory is retained for independent SQL sessions and latest-committed row-lock rechecks. Prefer [`Self::from_persistent_provider`] when a database-level owner is already available.
    pub fn from_persistent_backends(
        catalog: Arc<dyn uqa_storage::CatalogFacade>,
        backend: Arc<dyn PersistentStorageBackend>,
    ) -> StorageBackendResult<Self> {
        let identity = backend.storage_identity()?;
        let row_locks = crate::row_locks::shared_backend_manager(identity, &backend);
        let provider: Arc<dyn PersistentStorageProvider> = Arc::new(BackendSessionProvider {
            backend: Arc::clone(&backend),
        });
        let mut engine = Self::from_persistent_session(
            PersistentStorageSession::new(catalog, backend),
            Some(provider),
        )?;
        engine.session_id = row_locks.allocate_session();
        engine.row_locks = row_locks;
        Ok(engine)
    }

    fn from_persistent_session(
        storage_session: PersistentStorageSession,
        provider: Option<Arc<dyn PersistentStorageProvider>>,
    ) -> StorageBackendResult<Self> {
        let PersistentStorageSession { catalog, backend } = storage_session;
        let restore_catalog = Arc::clone(&catalog);
        let restore_backend = Arc::clone(&backend);
        let row_locks = Arc::new(crate::row_locks::RowLockManager::new());
        let session_id = row_locks.allocate_session();
        let mut engine = Self {
            storage: super::StorageContext::persistent(catalog, backend, provider),
            durable: Arc::new(super::DurableCatalogState::new()),
            session: Arc::new(super::SessionContext::new(super::initial_random_state())),
            extensions: super::RuntimeExtensions::new(),
            epochs: super::EpochCoordinator::new(),
            runtime: super::QueryRuntime::new(super::SQL_FUNCTION_DEPTH_LIMIT),
            row_locks,
            session_id,
            owns_session_registration: true,
            query_table_snapshots: None,
            query_view_snapshots: None,
            query_sql_function_snapshots: None,
            query_catalog_snapshot: None,
            query_transaction_overlay: None,
            query_transaction_origin: None,
        };
        Self::prepare_catalog_for_initial_restore(restore_catalog.as_ref())?;
        restore_backend.migrate_inverted_index_storage()?;
        engine.restore_from_catalog(restore_catalog.as_ref(), restore_backend.as_ref())?;
        engine.repair_reset_fts_storage(restore_catalog.as_ref())?;
        engine.repair_persistent_value_indexes_on_open()?;
        // Eagerly and fallibly populate read caches. Once open succeeds,
        // cache misses mean absence rather than a swallowed catalog error.
        for (name, json) in restore_catalog.load_models()? {
            let model = serde_json::from_str::<DeepModel>(&json)?;
            engine.durable.models.write().insert(name, model);
        }
        for (name, json) in restore_catalog.load_all_scoring_params()? {
            engine.durable.scoring_params.write().insert(name, json);
        }
        // Initial catalog migrations and physical-index repairs above may
        // commit. Establish the backend monitor baseline only after every
        // one-time write has completed.
        if let Some(version) = restore_backend.change_version()? {
            engine
                .epochs
                .seen_storage_change_version
                .store(version, std::sync::atomic::Ordering::Release);
        }
        Ok(engine)
    }

    fn sqlite_open_error(err: StorageBackendError) -> SQLiteError {
        match err {
            StorageBackendError::Analysis(err) => SQLiteError::Analysis(err),
            StorageBackendError::SQLite(err) => err,
            StorageBackendError::Serde(err) => SQLiteError::Serde(err),
            StorageBackendError::Backend { backend, source } => {
                SQLiteError::StorageBackend(format!("{backend} storage failed: {source}"))
            }
            StorageBackendError::Other(msg) => SQLiteError::StorageBackend(msg),
        }
    }
}