prefix_register/

lib.rs

// Copyright TELICENT LTD
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! # Prefix Register
//!
//! **Status: Beta** - API may change before 1.0 release.
//!
//! A PostgreSQL-backed namespace prefix registry for [CURIE](https://www.w3.org/TR/curie/)
//! expansion and prefix management.
//!
//! This library provides bidirectional mapping between namespace prefixes
//! (like "foaf", "rdf", "schema") and their full URI bases, optimised for
//! use in RDF/semantic web applications.
//!
//! **API:** Async-only, built on tokio and deadpool-postgres for high concurrency.
//!
//! ## Features
//!
//! - **Async-only** - Built on tokio for high concurrency
//! - **In-memory caching** - Prefixes loaded on startup for fast CURIE expansion
//! - **First-prefix-wins** - Each URI can only have one registered prefix
//! - **Batch operations** - Efficiently store multiple prefixes in a single round trip
//! - **PostgreSQL backend** - Durable, scalable storage with connection pooling
//! - **Startup resilience** - Optional retry with exponential backoff for container orchestration
//! - **Input validation** - Prevents DoS via length limits (prefix max 64, URI max 2048 chars)
//! - **Tracing instrumentation** - Built-in spans and events for observability
//!
//! ## Use Cases
//!
//! - CURIE expansion in RDF processing
//! - Namespace prefix management for semantic web applications
//! - Prefix discovery from Turtle, JSON-LD, XML documents
//!
//! ## Example
//!
//! ```rust,no_run
//! use prefix_register::PrefixRegistry;
//!
//! #[tokio::main]
//! async fn main() -> prefix_register::Result<()> {
//!     // Connect to PostgreSQL (schema must have namespaces table)
//!     let registry = PrefixRegistry::new(
//!         "postgres://localhost/mydb",
//!         10,  // max connections
//!     ).await?;
//!
//!     // Store a prefix (only if URI doesn't already have one)
//!     let stored = registry.store_prefix_if_new("foaf", "http://xmlns.com/foaf/0.1/").await?;
//!     println!("Prefix stored: {}", stored);
//!
//!     // Expand a CURIE
//!     if let Some(uri) = registry.expand_curie("foaf", "Person").await? {
//!         println!("foaf:Person = {}", uri);
//!     }
//!
//!     Ok(())
//! }
//! ```

mod error;

#[cfg(feature = "python")]
mod python;

pub use error::{ConfigurationError, Error, Result};

use deadpool_postgres::{Config as PoolConfig, Pool, Runtime};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tokio_postgres::NoTls;
use tracing::{debug, info, instrument, warn};

/// Maximum length for a prefix (64 characters).
pub const MAX_PREFIX_LENGTH: usize = 64;

/// Maximum length for a URI (2048 characters).
pub const MAX_URI_LENGTH: usize = 2048;

/// Configuration for connection retry behaviour.
///
/// Used with [`PrefixRegistry::new_with_retry`] to handle transient
/// database unavailability during startup.
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retry attempts (0 = no retries, just fail immediately).
    pub max_retries: u32,
    /// Initial delay before first retry. Doubles with each attempt (exponential backoff).
    pub initial_delay: Duration,
    /// Maximum delay between retries (caps the exponential growth).
    pub max_delay: Duration,
}

impl Default for RetryConfig {
    /// Default retry configuration: 5 retries, starting at 1 second, max 30 seconds.
    fn default() -> Self {
        Self {
            max_retries: 5,
            initial_delay: Duration::from_secs(1),
            max_delay: Duration::from_secs(30),
        }
    }
}

impl RetryConfig {
    /// Create a new retry configuration.
    pub fn new(max_retries: u32, initial_delay: Duration, max_delay: Duration) -> Self {
        Self {
            max_retries,
            initial_delay,
            max_delay,
        }
    }

    /// No retries - fail immediately on first error.
    pub fn none() -> Self {
        Self {
            max_retries: 0,
            initial_delay: Duration::ZERO,
            max_delay: Duration::ZERO,
        }
    }
}

/// Validate a prefix string.
///
/// Prefixes must be:
/// - Non-empty
/// - At most 64 characters
/// - Contain only alphanumeric characters, underscores, and hyphens
fn validate_prefix(prefix: &str) -> Result<()> {
    if prefix.is_empty() {
        return Err(Error::invalid_prefix("prefix cannot be empty"));
    }
    if prefix.len() > MAX_PREFIX_LENGTH {
        return Err(Error::invalid_prefix(format!(
            "prefix exceeds maximum length of {} characters",
            MAX_PREFIX_LENGTH
        )));
    }
    if !prefix
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
    {
        return Err(Error::invalid_prefix(
            "prefix must contain only alphanumeric characters, underscores, and hyphens",
        ));
    }
    Ok(())
}

/// Validate a URI string.
///
/// URIs must be:
/// - Non-empty
/// - At most 2048 characters
fn validate_uri(uri: &str) -> Result<()> {
    if uri.is_empty() {
        return Err(Error::invalid_uri("URI cannot be empty"));
    }
    if uri.len() > MAX_URI_LENGTH {
        return Err(Error::invalid_uri(format!(
            "URI exceeds maximum length of {} characters",
            MAX_URI_LENGTH
        )));
    }
    Ok(())
}

/// Result of a batch store operation.
///
/// Provides detailed information about what happened during
/// a batch prefix store, allowing callers to log appropriately.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct BatchStoreResult {
    /// Number of new prefixes successfully stored.
    pub stored: usize,
    /// Number of prefixes skipped (URI already had a prefix).
    pub skipped: usize,
}

impl BatchStoreResult {
    /// Total number of prefixes processed.
    pub fn total(&self) -> usize {
        self.stored + self.skipped
    }

    /// Returns true if all prefixes were stored (none skipped).
    pub fn all_stored(&self) -> bool {
        self.skipped == 0
    }

    /// Returns true if no prefixes were stored (all skipped or empty input).
    pub fn none_stored(&self) -> bool {
        self.stored == 0
    }
}

/// Registry for namespace prefixes.
///
/// Provides async access to the namespaces table in PostgreSQL.
/// Prefixes are stored with their corresponding URIs, following
/// the rule that each URI can only have one prefix (first one wins).
///
/// The registry maintains an in-memory cache of all prefixes, which
/// is populated on startup and updated as new prefixes are stored.
/// This ensures fast CURIE expansion without database round-trips.
#[derive(Clone)]
pub struct PrefixRegistry {
    /// Connection pool for the prefix database.
    /// Pool is internally Arc-wrapped, so cloning is cheap.
    pool: Pool,
    /// In-memory cache of prefix -> URI mappings for fast lookups.
    /// This cache is populated on startup and updated as new prefixes are stored.
    /// Wrapped in Arc to allow cheap cloning for async Python bindings.
    prefix_cache: Arc<RwLock<HashMap<String, String>>>,
    /// Reverse cache of URI -> prefix mappings for URI shortening.
    /// Used by shorten_uri() to find the longest matching namespace.
    uri_to_prefix: Arc<RwLock<HashMap<String, String>>>,
}

impl PrefixRegistry {
    /// Create a new prefix registry connected to the given PostgreSQL database.
    ///
    /// The registry will connect to the database and pre-populate its
    /// in-memory cache with existing prefixes for fast CURIE expansion.
    ///
    /// # Arguments
    ///
    /// * `database_url` - PostgreSQL connection URL (e.g., "postgres://user:pass@host:port/db")
    /// * `max_connections` - Maximum number of connections in the pool
    ///
    /// # Errors
    ///
    /// Returns an error if the database connection cannot be established
    /// or if the namespaces table does not exist.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use prefix_register::PrefixRegistry;
    ///
    /// # async fn example() -> prefix_register::Result<()> {
    /// let registry = PrefixRegistry::new(
    ///     "postgres://localhost/mydb",
    ///     10,
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(database_url), fields(max_connections))]
    pub async fn new(database_url: &str, max_connections: usize) -> Result<Self> {
        if max_connections == 0 {
            return Err(ConfigurationError::InvalidMaxConnections(max_connections).into());
        }
        if database_url.is_empty() {
            return Err(ConfigurationError::InvalidDatabaseUrl("empty URL".to_string()).into());
        }

        // Parse the database URL and create pool configuration
        let mut cfg = PoolConfig::new();
        cfg.url = Some(database_url.to_string());
        cfg.pool = Some(deadpool_postgres::PoolConfig::new(max_connections));

        // Create the connection pool
        let pool = cfg.create_pool(Some(Runtime::Tokio1), NoTls)?;

        debug!("created connection pool");

        // Verify connection by getting a client
        let client = pool.get().await?;

        // Pre-populate the cache with existing prefixes
        let rows = client
            .query("SELECT prefix, uri FROM namespaces", &[])
            .await?;

        let prefix_count = rows.len();
        let mut cache = HashMap::new();
        let mut reverse_cache = HashMap::new();
        for row in rows {
            let prefix: String = row.get(0);
            let uri: String = row.get(1);
            reverse_cache.insert(uri.clone(), prefix.clone());
            cache.insert(prefix, uri);
        }

        info!(prefix_count, "connected and loaded prefix cache");

        Ok(Self {
            pool,
            prefix_cache: Arc::new(RwLock::new(cache)),
            uri_to_prefix: Arc::new(RwLock::new(reverse_cache)),
        })
    }

    /// Create a new prefix registry with retry logic for transient failures.
    ///
    /// This variant of [`new`](Self::new) implements exponential backoff retry
    /// to handle transient database unavailability during startup (e.g., during
    /// container orchestration where the database may start after this service).
    ///
    /// # Arguments
    ///
    /// * `database_url` - PostgreSQL connection URL
    /// * `max_connections` - Maximum number of connections in the pool
    /// * `retry_config` - Configuration for retry behaviour
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use prefix_register::{PrefixRegistry, RetryConfig};
    /// use std::time::Duration;
    ///
    /// # async fn example() -> prefix_register::Result<()> {
    /// // Use default retry config (5 retries, 1s initial, 30s max)
    /// let registry = PrefixRegistry::new_with_retry(
    ///     "postgres://localhost/mydb",
    ///     10,
    ///     RetryConfig::default(),
    /// ).await?;
    ///
    /// // Or customize retry behaviour
    /// let registry = PrefixRegistry::new_with_retry(
    ///     "postgres://localhost/mydb",
    ///     10,
    ///     RetryConfig::new(
    ///         3,                           // max 3 retries
    ///         Duration::from_millis(500),  // start at 500ms
    ///         Duration::from_secs(10),     // cap at 10s
    ///     ),
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(database_url, retry_config), fields(max_connections, max_retries = retry_config.max_retries))]
    pub async fn new_with_retry(
        database_url: &str,
        max_connections: usize,
        retry_config: RetryConfig,
    ) -> Result<Self> {
        let mut last_error: Option<Error> = None;
        let mut delay = retry_config.initial_delay;

        for attempt in 0..=retry_config.max_retries {
            match Self::new(database_url, max_connections).await {
                Ok(registry) => return Ok(registry),
                Err(e) => {
                    // Configuration errors should not be retried
                    if matches!(e, Error::Configuration(_)) {
                        return Err(e);
                    }

                    last_error = Some(e);

                    // Don't sleep after the last attempt
                    if attempt < retry_config.max_retries {
                        warn!(
                            attempt = attempt + 1,
                            max_retries = retry_config.max_retries,
                            delay_ms = delay.as_millis() as u64,
                            "connection failed, retrying"
                        );
                        tokio::time::sleep(delay).await;
                        // Exponential backoff with cap
                        delay = std::cmp::min(delay * 2, retry_config.max_delay);
                    }
                }
            }
        }

        // All retries exhausted - return the last error
        Err(last_error.expect("should have at least one error after retries"))
    }

    /// Get the URI for a given prefix.
    ///
    /// First checks the in-memory cache, then falls back to the database.
    /// This is the primary method used for CURIE expansion.
    ///
    /// # Arguments
    ///
    /// * `prefix` - The namespace prefix (e.g., "foaf", "rdf")
    ///
    /// # Returns
    ///
    /// The URI if the prefix is known, None otherwise.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// if let Some(uri) = registry.get_uri_for_prefix("foaf").await? {
    ///     println!("foaf = {}", uri);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self), level = "debug")]
    pub async fn get_uri_for_prefix(&self, prefix: &str) -> Result<Option<String>> {
        // Check cache first (fast path)
        {
            let cache = self.prefix_cache.read().await;
            if let Some(uri) = cache.get(prefix) {
                debug!("cache hit");
                return Ok(Some(uri.clone()));
            }
        }

        // Cache miss - check database (handles concurrent updates)
        debug!("cache miss, querying database");
        let client = self.pool.get().await?;
        let row = client
            .query_opt("SELECT uri FROM namespaces WHERE prefix = $1", &[&prefix])
            .await?;

        if let Some(row) = row {
            let uri: String = row.get(0);
            // Update cache for future lookups
            {
                let mut cache = self.prefix_cache.write().await;
                cache.insert(prefix.to_string(), uri.clone());
            }
            debug!("found in database, cached");
            Ok(Some(uri))
        } else {
            debug!("not found");
            Ok(None)
        }
    }

    /// Get the prefix for a given URI.
    ///
    /// Used to check if a URI already has a registered prefix.
    ///
    /// # Arguments
    ///
    /// * `uri` - The full namespace URI
    ///
    /// # Returns
    ///
    /// The prefix if the URI is registered, None otherwise.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// if let Some(prefix) = registry.get_prefix_for_uri("http://xmlns.com/foaf/0.1/").await? {
    ///     println!("URI has prefix: {}", prefix);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self), level = "debug")]
    pub async fn get_prefix_for_uri(&self, uri: &str) -> Result<Option<String>> {
        let client = self.pool.get().await?;
        let row = client
            .query_opt("SELECT prefix FROM namespaces WHERE uri = $1", &[&uri])
            .await?;

        Ok(row.map(|r| r.get(0)))
    }

    /// Store a new prefix if the URI doesn't already have one.
    ///
    /// This follows the "first prefix wins" rule - if a URI already
    /// has a prefix registered, the new prefix is ignored.
    ///
    /// # Arguments
    ///
    /// * `prefix` - The namespace prefix to store
    /// * `uri` - The full namespace URI
    ///
    /// # Returns
    ///
    /// `true` if the prefix was stored, `false` if the URI already had a prefix.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let stored = registry.store_prefix_if_new("schema", "https://schema.org/").await?;
    /// if stored {
    ///     println!("New prefix stored");
    /// } else {
    ///     println!("URI already has a prefix");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self))]
    pub async fn store_prefix_if_new(&self, prefix: &str, uri: &str) -> Result<bool> {
        // Validate inputs to prevent DoS via memory exhaustion
        validate_prefix(prefix)?;
        validate_uri(uri)?;

        let client = self.pool.get().await?;

        // Use INSERT ... ON CONFLICT to atomically check and insert
        // This handles race conditions between multiple consumers
        let result = client
            .execute(
                "INSERT INTO namespaces (uri, prefix) VALUES ($1, $2)
                 ON CONFLICT (uri) DO NOTHING",
                &[&uri, &prefix],
            )
            .await?;

        if result > 0 {
            // Successfully inserted - update our caches
            {
                let mut cache = self.prefix_cache.write().await;
                cache.insert(prefix.to_string(), uri.to_string());
            }
            {
                let mut reverse = self.uri_to_prefix.write().await;
                reverse.insert(uri.to_string(), prefix.to_string());
            }
            debug!("stored new prefix");
            Ok(true)
        } else {
            // URI already has a prefix
            debug!("skipped, URI already has prefix");
            Ok(false)
        }
    }

    /// Store multiple prefixes, skipping any where the URI already has a prefix.
    ///
    /// More efficient than calling store_prefix_if_new repeatedly.
    ///
    /// # Arguments
    ///
    /// * `prefixes` - Iterator of (prefix, uri) pairs to store
    ///
    /// # Returns
    ///
    /// A [`BatchStoreResult`] with counts of stored and skipped prefixes.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let prefixes = vec![
    ///     ("foaf", "http://xmlns.com/foaf/0.1/"),
    ///     ("rdf", "http://www.w3.org/1999/02/22-rdf-syntax-ns#"),
    ///     ("schema", "https://schema.org/"),
    /// ];
    /// let result = registry.store_prefixes_if_new(prefixes).await?;
    /// println!("Stored {}, skipped {}", result.stored, result.skipped);
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self, prefixes))]
    pub async fn store_prefixes_if_new<'a, I>(&self, prefixes: I) -> Result<BatchStoreResult>
    where
        I: IntoIterator<Item = (&'a str, &'a str)>,
    {
        let prefixes: Vec<_> = prefixes.into_iter().collect();
        if prefixes.is_empty() {
            return Ok(BatchStoreResult::default());
        }

        // Validate all inputs upfront to prevent DoS via memory exhaustion
        for (prefix, uri) in &prefixes {
            validate_prefix(prefix)?;
            validate_uri(uri)?;
        }

        let total = prefixes.len();

        // Collect into separate arrays for UNNEST (single round trip)
        let uris: Vec<&str> = prefixes.iter().map(|(_, uri)| *uri).collect();
        let prefix_names: Vec<&str> = prefixes.iter().map(|(prefix, _)| *prefix).collect();

        let client = self.pool.get().await?;

        // Single batch INSERT using UNNEST, returning the rows that were inserted
        let rows = client
            .query(
                "INSERT INTO namespaces (uri, prefix)
                 SELECT * FROM UNNEST($1::text[], $2::text[])
                 ON CONFLICT (uri) DO NOTHING
                 RETURNING prefix, uri",
                &[&uris, &prefix_names],
            )
            .await?;

        let stored = rows.len();
        let skipped = total - stored;

        // Update caches with the rows that were actually inserted
        if !rows.is_empty() {
            let mut cache = self.prefix_cache.write().await;
            let mut reverse = self.uri_to_prefix.write().await;
            for row in &rows {
                let prefix: String = row.get(0);
                let uri: String = row.get(1);
                reverse.insert(uri.clone(), prefix.clone());
                cache.insert(prefix, uri);
            }
        }

        info!(total, stored, skipped, "batch store complete");

        Ok(BatchStoreResult { stored, skipped })
    }

    /// Expand a CURIE (Compact URI) to a full URI.
    ///
    /// Given a prefix and local name, returns the expanded URI
    /// if the prefix is known.
    ///
    /// # Arguments
    ///
    /// * `prefix` - The namespace prefix (e.g., "foaf")
    /// * `local_name` - The local part (e.g., "Person")
    ///
    /// # Returns
    ///
    /// The full URI (e.g., "http://xmlns.com/foaf/0.1/Person")
    /// or None if the prefix is unknown.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// if let Some(uri) = registry.expand_curie("foaf", "Person").await? {
    ///     println!("foaf:Person = {}", uri);
    ///     // Output: foaf:Person = http://xmlns.com/foaf/0.1/Person
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self), level = "debug")]
    pub async fn expand_curie(&self, prefix: &str, local_name: &str) -> Result<Option<String>> {
        if let Some(base_uri) = self.get_uri_for_prefix(prefix).await? {
            Ok(Some(format!("{}{}", base_uri, local_name)))
        } else {
            // Unknown prefix - caller can decide how to handle
            Ok(None)
        }
    }

    /// Get all registered prefixes.
    ///
    /// Returns a copy of the in-memory cache containing all prefix -> URI mappings.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let prefixes = registry.get_all_prefixes().await;
    /// for (prefix, uri) in prefixes {
    ///     println!("{}: {}", prefix, uri);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn get_all_prefixes(&self) -> HashMap<String, String> {
        self.prefix_cache.read().await.clone()
    }

    /// Get the number of registered prefixes.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let count = registry.prefix_count().await;
    /// println!("Registered prefixes: {}", count);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn prefix_count(&self) -> usize {
        self.prefix_cache.read().await.len()
    }

    /// Shorten a URI to a (prefix, local_name) tuple.
    ///
    /// Uses longest-match semantics: if multiple namespaces match the URI,
    /// the longest one wins. For example, with namespaces:
    /// - A = `http://a/`
    /// - B = `http://a/b#`
    ///
    /// The URI `http://a/b#thing` would shorten to `("B", "thing")`, not `("A", "b#thing")`.
    ///
    /// # Arguments
    ///
    /// * `uri` - The full URI to shorten
    ///
    /// # Returns
    ///
    /// `Some((prefix, local_name))` if a matching namespace is found, `None` otherwise.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// if let Some((prefix, local)) = registry.shorten_uri("http://xmlns.com/foaf/0.1/Person").await? {
    ///     println!("{}:{}", prefix, local);
    ///     // Output: foaf:Person
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self), level = "debug")]
    pub async fn shorten_uri(&self, uri: &str) -> Result<Option<(String, String)>> {
        let reverse = self.uri_to_prefix.read().await;

        // Find the longest matching namespace.
        // NOTE: This is O(n) where n = number of registered namespaces. Fine for typical
        // registries (thousands of entries), but if this becomes a bottleneck with very
        // large registries, consider replacing HashMap with a Trie for O(m) lookup
        // where m = URI length.
        let mut best_match: Option<(&str, &str)> = None;
        let mut best_len = 0;

        for (namespace_uri, prefix) in reverse.iter() {
            if uri.starts_with(namespace_uri) && namespace_uri.len() > best_len {
                best_match = Some((prefix.as_str(), namespace_uri.as_str()));
                best_len = namespace_uri.len();
            }
        }

        if let Some((prefix, namespace_uri)) = best_match {
            let local_name = &uri[namespace_uri.len()..];
            debug!(prefix, local_name, "shortened URI");
            Ok(Some((prefix.to_string(), local_name.to_string())))
        } else {
            debug!("no matching namespace found");
            Ok(None)
        }
    }

    /// Shorten a URI to a CURIE string, or return the original URI if no match.
    ///
    /// This is a convenience method that returns a ready-to-use string.
    ///
    /// # Arguments
    ///
    /// * `uri` - The full URI to shorten
    ///
    /// # Returns
    ///
    /// A CURIE string like `"prefix:local"`, or the original URI if no namespace matches.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let result = registry.shorten_uri_or_full("http://xmlns.com/foaf/0.1/Person").await?;
    /// println!("{}", result);
    /// // Output: foaf:Person (or the full URI if foaf is not registered)
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self), level = "debug")]
    pub async fn shorten_uri_or_full(&self, uri: &str) -> Result<String> {
        if let Some((prefix, local_name)) = self.shorten_uri(uri).await? {
            Ok(format!("{}:{}", prefix, local_name))
        } else {
            Ok(uri.to_string())
        }
    }

    /// Shorten multiple URIs in batch.
    ///
    /// Preserves order. Returns `None` for URIs that don't match any namespace.
    ///
    /// # Arguments
    ///
    /// * `uris` - Iterator of URIs to shorten
    ///
    /// # Returns
    ///
    /// A vector with `Some((prefix, local_name))` for matched URIs, `None` for unmatched.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let uris = vec![
    ///     "http://xmlns.com/foaf/0.1/Person",
    ///     "http://unknown.org/thing",
    ///     "https://schema.org/Organization",
    /// ];
    /// let results = registry.shorten_uri_batch(uris).await?;
    /// for (uri, result) in ["Person", "thing", "Organization"].iter().zip(results) {
    ///     match result {
    ///         Some((prefix, local)) => println!("{}:{}", prefix, local),
    ///         None => println!("No match"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self, uris), level = "debug")]
    pub async fn shorten_uri_batch<'a, I>(&self, uris: I) -> Result<Vec<Option<(String, String)>>>
    where
        I: IntoIterator<Item = &'a str>,
    {
        let reverse = self.uri_to_prefix.read().await;
        let uris: Vec<_> = uris.into_iter().collect();
        let mut results = Vec::with_capacity(uris.len());

        for uri in uris {
            // Find the longest matching namespace
            let mut best_match: Option<(&str, &str)> = None;
            let mut best_len = 0;

            for (namespace_uri, prefix) in reverse.iter() {
                if uri.starts_with(namespace_uri) && namespace_uri.len() > best_len {
                    best_match = Some((prefix.as_str(), namespace_uri.as_str()));
                    best_len = namespace_uri.len();
                }
            }

            if let Some((prefix, namespace_uri)) = best_match {
                let local_name = &uri[namespace_uri.len()..];
                results.push(Some((prefix.to_string(), local_name.to_string())));
            } else {
                results.push(None);
            }
        }

        debug!(count = results.len(), "batch shorten complete");
        Ok(results)
    }

    /// Expand multiple CURIEs in batch.
    ///
    /// Preserves order. Returns `None` for CURIEs with unknown prefixes.
    ///
    /// # Arguments
    ///
    /// * `curies` - Iterator of (prefix, local_name) pairs
    ///
    /// # Returns
    ///
    /// A vector with `Some(uri)` for known prefixes, `None` for unknown.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use prefix_register::PrefixRegistry;
    /// # async fn example(registry: &PrefixRegistry) -> prefix_register::Result<()> {
    /// let curies = vec![
    ///     ("foaf", "Person"),
    ///     ("unknown", "Thing"),
    ///     ("schema", "Organization"),
    /// ];
    /// let results = registry.expand_curie_batch(curies).await?;
    /// for result in results {
    ///     match result {
    ///         Some(uri) => println!("{}", uri),
    ///         None => println!("Unknown prefix"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[instrument(skip(self, curies), level = "debug")]
    pub async fn expand_curie_batch<'a, I>(&self, curies: I) -> Result<Vec<Option<String>>>
    where
        I: IntoIterator<Item = (&'a str, &'a str)>,
    {
        let cache = self.prefix_cache.read().await;
        let curies: Vec<_> = curies.into_iter().collect();
        let mut results = Vec::with_capacity(curies.len());

        for (prefix, local_name) in curies {
            if let Some(base_uri) = cache.get(prefix) {
                results.push(Some(format!("{}{}", base_uri, local_name)));
            } else {
                results.push(None);
            }
        }

        debug!(count = results.len(), "batch expand complete");
        Ok(results)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ==================== Unit Tests ====================
    // These run without a database

    #[test]
    fn test_configuration_error_max_connections() {
        let err = ConfigurationError::InvalidMaxConnections(0);
        assert!(err.to_string().contains("max_connections"));
    }

    #[test]
    fn test_configuration_error_database_url() {
        let err = ConfigurationError::InvalidDatabaseUrl("empty".to_string());
        assert!(err.to_string().contains("database_url"));
    }

    #[test]
    fn test_batch_store_result_default() {
        let result = BatchStoreResult::default();
        assert_eq!(result.stored, 0);
        assert_eq!(result.skipped, 0);
        assert_eq!(result.total(), 0);
        assert!(result.all_stored());
        assert!(result.none_stored());
    }

    #[test]
    fn test_batch_store_result_all_stored() {
        let result = BatchStoreResult {
            stored: 5,
            skipped: 0,
        };
        assert_eq!(result.total(), 5);
        assert!(result.all_stored());
        assert!(!result.none_stored());
    }

    #[test]
    fn test_batch_store_result_mixed() {
        let result = BatchStoreResult {
            stored: 3,
            skipped: 2,
        };
        assert_eq!(result.total(), 5);
        assert!(!result.all_stored());
        assert!(!result.none_stored());
    }

    #[test]
    fn test_batch_store_result_all_skipped() {
        let result = BatchStoreResult {
            stored: 0,
            skipped: 5,
        };
        assert_eq!(result.total(), 5);
        assert!(!result.all_stored());
        assert!(result.none_stored());
    }

    // ==================== Validation Tests ====================

    #[test]
    fn test_validate_prefix_valid() {
        assert!(validate_prefix("foaf").is_ok());
        assert!(validate_prefix("rdf").is_ok());
        assert!(validate_prefix("schema_org").is_ok());
        assert!(validate_prefix("my-prefix").is_ok());
        assert!(validate_prefix("prefix123").is_ok());
        assert!(validate_prefix("a").is_ok()); // Single char is valid
    }

    #[test]
    fn test_validate_prefix_empty() {
        let result = validate_prefix("");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("empty"));
    }

    #[test]
    fn test_validate_prefix_too_long() {
        let long_prefix = "a".repeat(MAX_PREFIX_LENGTH + 1);
        let result = validate_prefix(&long_prefix);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("maximum length"));
    }

    #[test]
    fn test_validate_prefix_max_length_ok() {
        let max_prefix = "a".repeat(MAX_PREFIX_LENGTH);
        assert!(validate_prefix(&max_prefix).is_ok());
    }

    #[test]
    fn test_validate_prefix_invalid_chars() {
        // Spaces not allowed
        let result = validate_prefix("foo bar");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("alphanumeric"));

        // Colons not allowed
        let result = validate_prefix("foo:bar");
        assert!(result.is_err());

        // Slashes not allowed
        let result = validate_prefix("foo/bar");
        assert!(result.is_err());

        // Unicode not allowed
        let result = validate_prefix("préfix");
        assert!(result.is_err());
    }

    #[test]
    fn test_validate_uri_valid() {
        assert!(validate_uri("http://example.org/").is_ok());
        assert!(validate_uri("https://schema.org/Person").is_ok());
        assert!(validate_uri("urn:isbn:0451450523").is_ok());
        assert!(validate_uri("http://example.org/path?query=1#fragment").is_ok());
    }

    #[test]
    fn test_validate_uri_empty() {
        let result = validate_uri("");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("empty"));
    }

    #[test]
    fn test_validate_uri_too_long() {
        let long_uri = format!("http://example.org/{}", "a".repeat(MAX_URI_LENGTH));
        let result = validate_uri(&long_uri);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("maximum length"));
    }

    #[test]
    fn test_validate_uri_max_length_ok() {
        // Create a URI exactly at the max length
        let base = "http://example.org/";
        let padding = "a".repeat(MAX_URI_LENGTH - base.len());
        let max_uri = format!("{}{}", base, padding);
        assert_eq!(max_uri.len(), MAX_URI_LENGTH);
        assert!(validate_uri(&max_uri).is_ok());
    }

    // ==================== RetryConfig Tests ====================

    #[test]
    fn test_retry_config_default() {
        let config = RetryConfig::default();
        assert_eq!(config.max_retries, 5);
        assert_eq!(config.initial_delay, Duration::from_secs(1));
        assert_eq!(config.max_delay, Duration::from_secs(30));
    }

    #[test]
    fn test_retry_config_none() {
        let config = RetryConfig::none();
        assert_eq!(config.max_retries, 0);
        assert_eq!(config.initial_delay, Duration::ZERO);
        assert_eq!(config.max_delay, Duration::ZERO);
    }

    #[test]
    fn test_retry_config_custom() {
        let config = RetryConfig::new(3, Duration::from_millis(100), Duration::from_secs(5));
        assert_eq!(config.max_retries, 3);
        assert_eq!(config.initial_delay, Duration::from_millis(100));
        assert_eq!(config.max_delay, Duration::from_secs(5));
    }

    #[tokio::test]
    async fn test_new_with_retry_config_error_not_retried() {
        // Configuration errors should fail immediately, not retry
        let result = PrefixRegistry::new_with_retry(
            "", // Empty URL is a config error
            5,
            RetryConfig::default(),
        )
        .await;

        assert!(matches!(
            result,
            Err(Error::Configuration(
                ConfigurationError::InvalidDatabaseUrl(_)
            ))
        ));
    }

    // ==================== Integration Tests ====================
    // These require DATABASE_URL to be set (provided by CI)

    /// Helper to get database URL from environment.
    /// Returns None if not set (skips integration tests locally).
    fn get_test_database_url() -> Option<String> {
        std::env::var("DATABASE_URL").ok()
    }

    /// Helper to clean up test data. Uses a unique prefix to avoid conflicts.
    async fn cleanup_test_data(registry: &PrefixRegistry, test_prefix: &str) {
        let client = registry.pool.get().await.unwrap();
        client
            .execute(
                "DELETE FROM namespaces WHERE prefix LIKE $1",
                &[&format!("{}%", test_prefix)],
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_new_with_invalid_max_connections() {
        let result = PrefixRegistry::new("postgres://localhost/test", 0).await;
        assert!(matches!(
            result,
            Err(Error::Configuration(
                ConfigurationError::InvalidMaxConnections(0)
            ))
        ));
    }

    #[tokio::test]
    async fn test_new_with_empty_url() {
        let result = PrefixRegistry::new("", 5).await;
        assert!(matches!(
            result,
            Err(Error::Configuration(
                ConfigurationError::InvalidDatabaseUrl(_)
            ))
        ));
    }

    #[tokio::test]
    async fn test_store_and_retrieve_prefix() {
        let Some(db_url) = get_test_database_url() else {
            return; // Skip if no database
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_sr_";
        cleanup_test_data(&registry, test_prefix).await;

        // Store a new prefix (use unique URI to avoid conflicts)
        let prefix = format!("{test_prefix}ns");
        let uri = format!("http://example.org/{test_prefix}ns/");
        let stored = registry.store_prefix_if_new(&prefix, &uri).await.unwrap();
        assert!(stored, "First store should succeed");

        // Retrieve by prefix
        let retrieved = registry.get_uri_for_prefix(&prefix).await.unwrap();
        assert_eq!(retrieved, Some(uri.to_string()));

        // Retrieve by URI
        let retrieved_prefix = registry.get_prefix_for_uri(&uri).await.unwrap();
        assert_eq!(retrieved_prefix, Some(prefix.clone()));

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_first_prefix_wins() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_fpw_";
        cleanup_test_data(&registry, test_prefix).await;

        let uri = "http://example.org/test/first-wins/";
        let first_prefix = format!("{test_prefix}first");
        let second_prefix = format!("{test_prefix}second");

        // Store first prefix
        let stored1 = registry
            .store_prefix_if_new(&first_prefix, uri)
            .await
            .unwrap();
        assert!(stored1, "First prefix should be stored");

        // Try to store second prefix for same URI
        let stored2 = registry
            .store_prefix_if_new(&second_prefix, uri)
            .await
            .unwrap();
        assert!(!stored2, "Second prefix should be rejected");

        // Verify the first prefix is still there
        let retrieved = registry.get_prefix_for_uri(uri).await.unwrap();
        assert_eq!(retrieved, Some(first_prefix));

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_expand_curie() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_ec_";
        cleanup_test_data(&registry, test_prefix).await;

        // Use unique URI to avoid conflicts
        let prefix = format!("{test_prefix}ns");
        let uri = format!("http://example.org/{test_prefix}ns/");
        registry.store_prefix_if_new(&prefix, &uri).await.unwrap();

        // Expand known prefix
        let expanded = registry.expand_curie(&prefix, "Person").await.unwrap();
        assert_eq!(expanded, Some(format!("{uri}Person")));

        // Unknown prefix returns None
        let unknown = registry
            .expand_curie(&format!("{test_prefix}unknown"), "Thing")
            .await
            .unwrap();
        assert_eq!(unknown, None);

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_batch_store_prefixes() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_bs_";
        cleanup_test_data(&registry, test_prefix).await;

        // Use unique URIs to avoid conflicts with pre-existing data
        let prefixes = [
            (
                format!("{test_prefix}ns1"),
                format!("http://example.org/{test_prefix}ns1/"),
            ),
            (
                format!("{test_prefix}ns2"),
                format!("http://example.org/{test_prefix}ns2/"),
            ),
            (
                format!("{test_prefix}ns3"),
                format!("http://example.org/{test_prefix}ns3/"),
            ),
        ];

        let initial_count = registry.prefix_count().await;

        let prefix_refs: Vec<(&str, &str)> = prefixes
            .iter()
            .map(|(p, u)| (p.as_str(), u.as_str()))
            .collect();

        let result = registry.store_prefixes_if_new(prefix_refs).await.unwrap();
        assert_eq!(result.stored, 3);
        assert_eq!(result.skipped, 0);
        assert!(result.all_stored());

        // Verify all were stored (count increased by 3)
        assert_eq!(registry.prefix_count().await, initial_count + 3);

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_batch_store_with_duplicates() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_bsd_";
        cleanup_test_data(&registry, test_prefix).await;

        // Pre-store one prefix
        let existing_uri = "http://example.org/existing/";
        registry
            .store_prefix_if_new(&format!("{test_prefix}existing"), existing_uri)
            .await
            .unwrap();

        // Batch store including the existing URI with a different prefix
        let prefixes = [
            (
                format!("{test_prefix}new1"),
                "http://example.org/new1/".to_string(),
            ),
            (format!("{test_prefix}duplicate"), existing_uri.to_string()), // Should be skipped
            (
                format!("{test_prefix}new2"),
                "http://example.org/new2/".to_string(),
            ),
        ];

        let prefix_refs: Vec<(&str, &str)> = prefixes
            .iter()
            .map(|(p, u)| (p.as_str(), u.as_str()))
            .collect();

        let result = registry.store_prefixes_if_new(prefix_refs).await.unwrap();
        assert_eq!(result.stored, 2);
        assert_eq!(result.skipped, 1);
        assert!(!result.all_stored());
        assert!(!result.none_stored());

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_batch_store_empty() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let empty: Vec<(&str, &str)> = vec![];

        let result = registry.store_prefixes_if_new(empty).await.unwrap();
        assert_eq!(result.stored, 0);
        assert_eq!(result.skipped, 0);
        assert!(result.all_stored()); // Vacuously true
        assert!(result.none_stored());
    }

    #[tokio::test]
    async fn test_get_all_prefixes() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_gap_";
        cleanup_test_data(&registry, test_prefix).await;

        // Store some prefixes
        let prefix1 = format!("{test_prefix}a");
        let prefix2 = format!("{test_prefix}b");
        registry
            .store_prefix_if_new(&prefix1, "http://example.org/a/")
            .await
            .unwrap();
        registry
            .store_prefix_if_new(&prefix2, "http://example.org/b/")
            .await
            .unwrap();

        let all = registry.get_all_prefixes().await;
        assert!(all.contains_key(&prefix1));
        assert!(all.contains_key(&prefix2));
        assert_eq!(
            all.get(&prefix1),
            Some(&"http://example.org/a/".to_string())
        );

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_cache_populated_on_startup() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let test_prefix = "test_cache_";

        // Create first registry and store a prefix
        {
            let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
            cleanup_test_data(&registry, test_prefix).await;
            registry
                .store_prefix_if_new(
                    &format!("{test_prefix}cached"),
                    "http://example.org/cached/",
                )
                .await
                .unwrap();
        }

        // Create new registry - should have prefix in cache from startup
        let registry2 = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let cached = registry2.get_all_prefixes().await;
        assert!(cached.contains_key(&format!("{test_prefix}cached")));

        cleanup_test_data(&registry2, test_prefix).await;
    }

    #[tokio::test]
    async fn test_unknown_prefix_returns_none() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();

        let result = registry
            .get_uri_for_prefix("definitely_not_a_real_prefix_xyz123")
            .await
            .unwrap();
        assert_eq!(result, None);
    }

    #[tokio::test]
    async fn test_unknown_uri_returns_none() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();

        let result = registry
            .get_prefix_for_uri("http://definitely-not-registered.example.org/")
            .await
            .unwrap();
        assert_eq!(result, None);
    }

    #[tokio::test]
    async fn test_new_with_retry_succeeds() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        // Should succeed on first try with valid database
        let registry = PrefixRegistry::new_with_retry(
            &db_url,
            5,
            RetryConfig::none(), // No retries needed for working DB
        )
        .await
        .unwrap();

        // Verify it works
        let test_prefix = "test_nwr_";
        cleanup_test_data(&registry, test_prefix).await;

        let stored = registry
            .store_prefix_if_new(&format!("{test_prefix}retry"), "http://example.org/retry/")
            .await
            .unwrap();
        assert!(stored);

        cleanup_test_data(&registry, test_prefix).await;
    }

    // ==================== Shorten URI Tests ====================

    #[tokio::test]
    async fn test_shorten_uri() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_su_";
        cleanup_test_data(&registry, test_prefix).await;

        // Use unique URI to avoid conflicts
        let prefix = format!("{test_prefix}ns");
        let uri = format!("http://example.org/{test_prefix}ns/");
        registry.store_prefix_if_new(&prefix, &uri).await.unwrap();

        // Shorten a known URI
        let result = registry.shorten_uri(&format!("{uri}Person")).await.unwrap();
        assert_eq!(result, Some((prefix.clone(), "Person".to_string())));

        // Unknown URI returns None
        let unknown = registry
            .shorten_uri("http://unknown.org/thing")
            .await
            .unwrap();
        assert_eq!(unknown, None);

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_shorten_uri_longest_match() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_sulm_";
        cleanup_test_data(&registry, test_prefix).await;

        // Register overlapping namespaces
        let short_prefix = format!("{test_prefix}short");
        let long_prefix = format!("{test_prefix}long");
        registry
            .store_prefix_if_new(&short_prefix, "http://example.org/")
            .await
            .unwrap();
        registry
            .store_prefix_if_new(&long_prefix, "http://example.org/nested/")
            .await
            .unwrap();

        // URI that matches both - should use the longer namespace
        let result = registry
            .shorten_uri("http://example.org/nested/thing")
            .await
            .unwrap();
        assert_eq!(result, Some((long_prefix.clone(), "thing".to_string())));

        // URI that only matches the short namespace
        let result = registry
            .shorten_uri("http://example.org/other")
            .await
            .unwrap();
        assert_eq!(result, Some((short_prefix.clone(), "other".to_string())));

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_shorten_uri_or_full() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_suof_";
        cleanup_test_data(&registry, test_prefix).await;

        // Use unique URIs to avoid conflicts with pre-existing data
        let prefix = format!("{test_prefix}ns");
        let uri = format!("http://example.org/{test_prefix}ns/");
        registry.store_prefix_if_new(&prefix, &uri).await.unwrap();

        // Known URI returns CURIE
        let result = registry
            .shorten_uri_or_full(&format!("{uri}Person"))
            .await
            .unwrap();
        assert_eq!(result, format!("{}:Person", prefix));

        // Unknown URI returns original
        let unknown_uri = "http://unknown.org/thing";
        let result = registry.shorten_uri_or_full(unknown_uri).await.unwrap();
        assert_eq!(result, unknown_uri);

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_shorten_uri_batch() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_sub_";
        cleanup_test_data(&registry, test_prefix).await;

        // Use unique URIs to avoid conflicts with pre-existing data
        let prefix1 = format!("{test_prefix}ns1");
        let prefix2 = format!("{test_prefix}ns2");
        let uri1 = format!("http://example.org/{test_prefix}ns1/");
        let uri2 = format!("http://example.org/{test_prefix}ns2/");
        registry.store_prefix_if_new(&prefix1, &uri1).await.unwrap();
        registry.store_prefix_if_new(&prefix2, &uri2).await.unwrap();

        let uris = [
            format!("{uri1}Person"),
            "http://unknown.org/thing".to_string(),
            format!("{uri2}Organization"),
        ];

        let results = registry
            .shorten_uri_batch(uris.iter().map(|s| s.as_str()))
            .await
            .unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0], Some((prefix1.clone(), "Person".to_string())));
        assert_eq!(results[1], None); // Unknown
        assert_eq!(
            results[2],
            Some((prefix2.clone(), "Organization".to_string()))
        );

        cleanup_test_data(&registry, test_prefix).await;
    }

    #[tokio::test]
    async fn test_expand_curie_batch() {
        let Some(db_url) = get_test_database_url() else {
            return;
        };

        let registry = PrefixRegistry::new(&db_url, 5).await.unwrap();
        let test_prefix = "test_ecb_";
        cleanup_test_data(&registry, test_prefix).await;

        // Use unique URIs to avoid conflicts
        let prefix1 = format!("{test_prefix}ns1");
        let prefix2 = format!("{test_prefix}ns2");
        let uri1 = format!("http://example.org/{test_prefix}ns1/");
        let uri2 = format!("http://example.org/{test_prefix}ns2/");
        registry.store_prefix_if_new(&prefix1, &uri1).await.unwrap();
        registry.store_prefix_if_new(&prefix2, &uri2).await.unwrap();

        let curies = vec![
            (prefix1.as_str(), "Person"),
            ("unknown_prefix_xyz", "Thing"),
            (prefix2.as_str(), "Organization"),
        ];

        let results = registry.expand_curie_batch(curies).await.unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0], Some(format!("{uri1}Person")));
        assert_eq!(results[1], None); // Unknown prefix
        assert_eq!(results[2], Some(format!("{uri2}Organization")));

        cleanup_test_data(&registry, test_prefix).await;
    }
}