uri_register/

postgres.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.

use crate::cache::{create_cache, Cache, CacheStrategy};
use crate::error::{ConfigurationError, Result};
use crate::service::UriService;
use async_trait::async_trait;
use deadpool_postgres::{ManagerConfig, Pool, RecyclingMethod, Runtime};
use rustls::RootCertStore;
use std::sync::Arc;
use tokio_postgres::{Config, NoTls};
use tokio_postgres_rustls::MakeRustlsConnect;
use tracing::{debug, info, instrument, trace};
use url::Url;

/// PostgreSQL-based URI register implementation with configurable caching
///
/// This implementation uses a PostgreSQL table to store URI-to-ID mappings
/// with an in-memory cache (W-TinyLFU by default, or LRU) to reduce database round-trips.
/// It's designed for high concurrency with connection pooling and batch operations.
///
/// ## Prerequisites
///
/// The database schema must be initialized before using this service.
/// See `schema.sql` for the DDL statements.
///
/// ## URI Validation
///
/// All URIs are validated before registration to ensure they conform to RFC 3986.
/// Invalid URIs will return an error.
///
/// ## Performance
///
/// With default logged tables on typical hardware:
/// - Batch insert: ~10K-50K URIs/sec
/// - Batch lookup (cached): ~100K-1M+ URIs/sec (no DB round-trip)
/// - Batch lookup (uncached): ~100K-200K URIs/sec
/// - Query overhead: ~2-10ms per query (2 round-trips)
///
/// The cache (W-TinyLFU or LRU) significantly improves performance for repeated URI lookups.
/// Cache strategy and size are configurable when creating the register instance.
///
/// For faster writes at the cost of durability, the table can be configured
/// as UNLOGGED (see `schema.sql` for options).
pub struct PostgresUriRegister {
    pool: Pool,
    /// Cache for URI-to-ID mappings (W-TinyLFU or LRU)
    cache: Arc<dyn Cache>,
    /// Name of the database table to use
    table_name: String,
}

impl PostgresUriRegister {
    /// Create a new PostgreSQL URI register service with configurable cache
    ///
    /// # Arguments
    ///
    /// * `database_url` - PostgreSQL connection string (e.g., "postgres://user:password@host:port/database")
    /// * `table_name` - Name of the database table to use (must be a valid SQL identifier, default: "uri_register")
    /// * `max_connections` - Maximum number of connections in the pool (recommended: 10-50)
    /// * `cache_size` - Number of URI-to-ID mappings to cache in memory (recommended: 1,000-100,000)
    /// * `cache_strategy` - Cache strategy to use (Moka/W-TinyLFU is default and recommended for most workloads)
    ///
    /// # Prerequisites
    ///
    /// The database schema must be initialized before using this service.
    /// See the `schema.sql` file and README.md for setup instructions.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use uri_register::PostgresUriRegister;
    ///
    /// #[tokio::main]
    /// async fn main() -> uri_register::Result<()> {
    ///     let register = PostgresUriRegister::new(
    ///         "postgres://localhost/mydb",
    ///         "uri_register",  // table name
    ///         20,              // max connections
    ///         10_000           // cache size (defaults to Moka/W-TinyLFU)
    ///     ).await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn new(
        database_url: &str,
        table_name: &str,
        max_connections: u32,
        cache_size: usize,
    ) -> Result<Self> {
        Self::new_with_cache_strategy(
            database_url,
            table_name,
            max_connections,
            cache_size,
            None, // Default to Moka
            None, // Default to no TLS
        )
        .await
    }

    /// Create a new PostgreSQL URI register with a specific cache strategy and TLS
    ///
    /// This is identical to `new()` but allows specifying a cache strategy and TLS option.
    /// Most users should use `new()` which defaults to the recommended Moka (W-TinyLFU) cache and no TLS.
    ///
    /// # Arguments
    ///
    /// * `cache_strategy` - Optional cache strategy (None = Moka default, or specify CacheStrategy::Lru)
    /// * `use_tls` - Optional TLS flag (None/false = no TLS, true = TLS with webpki root certificates)
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use uri_register::{CacheStrategy, PostgresUriRegister};
    ///
    /// #[tokio::main]
    /// async fn main() -> uri_register::Result<()> {
    ///     // Use LRU instead of default Moka, with TLS enabled
    ///     let register = PostgresUriRegister::new_with_cache_strategy(
    ///         "postgres://localhost/mydb",
    ///         "uri_register",
    ///         20,
    ///         10_000,
    ///         Some(CacheStrategy::Lru),
    ///         Some(true)  // Enable TLS
    ///     ).await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn new_with_cache_strategy(
        database_url: &str,
        table_name: &str,
        max_connections: u32,
        cache_size: usize,
        cache_strategy: Option<CacheStrategy>,
        use_tls: Option<bool>,
    ) -> Result<Self> {
        // Validate inputs
        if cache_size == 0 {
            return Err(ConfigurationError::InvalidCacheSize(cache_size).into());
        }

        if max_connections == 0 {
            return Err(ConfigurationError::InvalidMaxConnections(max_connections).into());
        }

        // Validate table name as SQL identifier
        Self::validate_table_name(table_name)?;

        // Parse the database URL into tokio-postgres Config
        let pg_config: Config = database_url.parse().map_err(|e| {
            ConfigurationError::InvalidBackoff(format!("Failed to parse database URL: {}", e))
        })?;

        // Create deadpool configuration
        let mut cfg = deadpool_postgres::Config::new();
        cfg.dbname = pg_config.get_dbname().map(|s| s.to_string());
        cfg.host = pg_config.get_hosts().first().map(|h| match h {
            tokio_postgres::config::Host::Tcp(s) => s.to_string(),
            #[cfg(unix)]
            tokio_postgres::config::Host::Unix(p) => p.to_str().unwrap_or_default().to_string(),
        });
        cfg.port = pg_config.get_ports().first().copied();
        cfg.user = pg_config.get_user().map(|s| s.to_string());
        cfg.password = pg_config
            .get_password()
            .map(|p| std::str::from_utf8(p).unwrap_or_default().to_string());
        cfg.manager = Some(ManagerConfig {
            recycling_method: RecyclingMethod::Fast,
        });
        cfg.pool = Some(deadpool_postgres::PoolConfig {
            max_size: max_connections as usize,
            timeouts: deadpool_postgres::Timeouts {
                wait: Some(std::time::Duration::from_secs(10)),
                create: Some(std::time::Duration::from_secs(10)),
                recycle: Some(std::time::Duration::from_secs(10)),
            },
            ..Default::default()
        });

        // Create the pool with or without TLS based on use_tls parameter
        let pool = if use_tls.unwrap_or(false) {
            // Configure TLS with webpki root certificates
            let mut root_store = RootCertStore::empty();
            root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());

            let tls_config = rustls::ClientConfig::builder()
                .with_root_certificates(root_store)
                .with_no_client_auth();

            let tls = MakeRustlsConnect::new(tls_config);

            cfg.create_pool(Some(Runtime::Tokio1), tls).map_err(|e| {
                ConfigurationError::InvalidBackoff(format!(
                    "Failed to create connection pool with TLS: {}",
                    e
                ))
            })?
        } else {
            // No TLS
            cfg.create_pool(Some(Runtime::Tokio1), NoTls).map_err(|e| {
                ConfigurationError::InvalidBackoff(format!(
                    "Failed to create connection pool: {}",
                    e
                ))
            })?
        };

        let cache = create_cache(cache_strategy.unwrap_or_default(), cache_size);

        info!(
            table = table_name,
            max_connections,
            cache_size,
            tls = use_tls.unwrap_or(false),
            "URI register connected"
        );

        Ok(Self {
            pool,
            cache,
            table_name: table_name.to_string(),
        })
    }

    /// Validate that a table name is a valid SQL identifier
    ///
    /// Prevents SQL injection by ensuring the table name only contains
    /// alphanumeric characters and underscores, and doesn't start with a digit.
    fn validate_table_name(name: &str) -> Result<()> {
        if name.is_empty() {
            return Err(ConfigurationError::InvalidTableName(
                "table name cannot be empty".to_string(),
            )
            .into());
        }

        if name.len() > 63 {
            return Err(ConfigurationError::InvalidTableName(format!(
                "table name too long (max 63 characters): '{}'",
                name
            ))
            .into());
        }

        // First character must be a letter or underscore
        let first_char = name.chars().next().unwrap();
        if !first_char.is_ascii_alphabetic() && first_char != '_' {
            return Err(ConfigurationError::InvalidTableName(format!(
                "table name must start with a letter or underscore: '{}'",
                name
            ))
            .into());
        }

        // All characters must be alphanumeric or underscore
        if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
            return Err(ConfigurationError::InvalidTableName(format!(
                "table name can only contain letters, numbers, and underscores: '{}'",
                name
            ))
            .into());
        }

        Ok(())
    }

    /// Get statistics about the URI register
    ///
    /// Returns the total number of URIs and the storage size.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use uri_register::PostgresUriRegister;
    ///
    /// #[tokio::main]
    /// async fn main() -> uri_register::Result<()> {
    ///     let register = PostgresUriRegister::new(
    ///         "postgres://localhost/mydb",
    ///         "uri_register",
    ///         20,
    ///         10_000
    ///     ).await?;
    ///     let stats = register.stats().await?;
    ///     println!("Total URIs: {}", stats.total_uris);
    ///     println!("Size: {} bytes", stats.size_bytes);
    ///     Ok(())
    /// }
    /// ```
    pub async fn stats(&self) -> Result<RegisterStats> {
        // Build query with validated table name (safe from SQL injection)
        let query = format!(
            r#"
            SELECT
                COUNT(*)::bigint as count,
                pg_total_relation_size('{}')::bigint as size_bytes
            FROM {}
            "#,
            self.table_name, self.table_name
        );

        // Execute with retry logic
        let client = self.pool.get().await.map_err(|e| {
            crate::error::Error::Database(format!("Failed to get database connection: {}", e))
        })?;

        let rows = client
            .query(&query, &[])
            .await
            .map_err(|e| crate::error::Error::Database(e.to_string()))?;

        let row = rows.into_iter().next().ok_or_else(|| {
            crate::error::Error::Database("No rows returned from stats query".to_string())
        })?;

        // Get cache statistics
        let cache_stats = self.cache.stats();

        // Get connection pool statistics
        let status = self.pool.status();
        let pool_stats = PoolStats {
            connections_active: (status.size - status.available) as u32,
            connections_idle: status.available as u32,
            connections_max: status.max_size as u32,
        };

        Ok(RegisterStats {
            total_uris: row.get::<_, i64>("count") as u64,
            size_bytes: row.get::<_, i64>("size_bytes") as u64,
            cache: cache_stats,
            pool: pool_stats,
        })
    }

    /// Clone the register instance (shares pool and cache)
    ///
    /// This is a shallow clone that shares both the connection pool and cache.
    /// Both pool and cache use Arc internally, so this clone is cheap and shares
    /// the underlying resources.
    ///
    /// This method is primarily used for Python bindings where we need to move
    /// data into async closures.
    #[cfg(feature = "python")]
    pub(crate) fn clone_inner(&self) -> Self {
        PostgresUriRegister {
            pool: self.pool.clone(),
            cache: self.cache.clone(), // Clone the Arc, shares the same cache
            table_name: self.table_name.clone(),
        }
    }

    /// Validate that a string is a valid URI according to RFC 3986
    fn validate_uri(uri: &str) -> Result<()> {
        Url::parse(uri).map_err(|e| {
            crate::error::Error::InvalidUri(format!("Invalid URI '{}': {}", uri, e))
        })?;
        Ok(())
    }
}

#[async_trait]
impl UriService for PostgresUriRegister {
    #[instrument(skip(self), fields(table = %self.table_name))]
    async fn register_uri(&self, uri: &str) -> Result<u64> {
        // Validate URI first
        Self::validate_uri(uri)?;

        // Check cache first
        if let Some(id) = self.cache.get(uri) {
            trace!(id, "cache hit");
            return Ok(id);
        }
        trace!("cache miss, querying database");

        // Insert and return ID (ON CONFLICT handles race conditions and existing URIs)
        // Build query with validated table name (safe from SQL injection)
        let query = format!(
            r#"
            INSERT INTO {} (uri)
            VALUES ($1)
            ON CONFLICT (uri_hash) DO UPDATE SET uri = EXCLUDED.uri
            RETURNING id
            "#,
            self.table_name
        );

        // Execute with retry logic
        let client = self.pool.get().await.map_err(|e| {
            crate::error::Error::Database(format!("Failed to get database connection: {}", e))
        })?;

        let rows = client
            .query(&query, &[&uri])
            .await
            .map_err(|e| crate::error::Error::Database(e.to_string()))?;

        let result = rows.into_iter().next().ok_or_else(|| {
            crate::error::Error::Database("No rows returned from register_uri query".to_string())
        })?;

        let id = result.get::<_, i64>("id") as u64;

        // Update cache
        self.cache.put(uri.to_string(), id);

        Ok(id)
    }

    #[instrument(skip(self, uris), fields(table = %self.table_name, batch_size = uris.len()))]
    async fn register_uri_batch(&self, uris: &[String]) -> Result<Vec<u64>> {
        if uris.is_empty() {
            trace!("empty batch, returning early");
            return Ok(Vec::new());
        }

        // Validate all URIs first
        for uri in uris {
            Self::validate_uri(uri)?;
        }

        // CORRECTNESS GUARANTEE: Order preservation
        // We maintain strict correspondence between input URIs and output IDs
        // by tracking the original index of each URI and using URI strings
        // (not SQL result order) to map IDs back to their positions.

        let mut result_ids = vec![None; uris.len()];
        let mut uncached_indices = Vec::new();
        let mut uncached_uris_dedup = Vec::new();
        let mut seen_uncached = std::collections::HashMap::new();

        // Step 1: Check cache for all URIs
        for (idx, uri) in uris.iter().enumerate() {
            if let Some(id) = self.cache.get(uri) {
                result_ids[idx] = Some(id);
            } else {
                uncached_indices.push(idx);
                // Deduplicate uncached URIs for DB query
                if !seen_uncached.contains_key(uri) {
                    seen_uncached.insert(uri.clone(), uncached_uris_dedup.len());
                    uncached_uris_dedup.push(uri.clone());
                }
            }
        }

        // If everything was cached, return early
        if uncached_uris_dedup.is_empty() {
            debug!(cached = uris.len(), "all URIs found in cache");
            return Ok(result_ids.into_iter().map(|id| id.unwrap()).collect());
        }

        let cached_count = uris.len() - uncached_indices.len();
        debug!(
            cached = cached_count,
            uncached = uncached_uris_dedup.len(),
            "cache lookup complete, querying database"
        );

        // Step 2: Register deduplicated uncached URIs in batch
        // IMPORTANT: SQL may return results in ANY order (not guaranteed to match input order)
        // We use "RETURNING id, uri" to get BOTH values together, then map by URI string
        // Build query with validated table name (safe from SQL injection)
        let query = format!(
            r#"
            INSERT INTO {} (uri)
            SELECT unnest($1::text[])
            ON CONFLICT (uri_hash) DO UPDATE SET uri = EXCLUDED.uri
            RETURNING id, uri
            "#,
            self.table_name
        );

        // Execute with retry logic
        let client = self.pool.get().await.map_err(|e| {
            crate::error::Error::Database(format!("Failed to get database connection: {}", e))
        })?;

        let rows = client
            .query(&query, &[&uncached_uris_dedup])
            .await
            .map_err(|e| crate::error::Error::Database(e.to_string()))?;

        // Build a map of URI -> ID from database results
        // This allows us to look up IDs by URI string (order-independent)
        let mut uri_to_id = std::collections::HashMap::new();
        for row in rows {
            let uri: String = row.get("uri");
            let id: i64 = row.get("id");
            uri_to_id.insert(uri, id as u64);
        }

        // Step 3: Fill in the result vector and update cache
        // CORRECTNESS: We use the saved indices and look up by URI string,
        // guaranteeing that result_ids[i] corresponds to uris[i]
        for idx in uncached_indices {
            let uri = &uris[idx]; // Get URI from original position
            if let Some(&id) = uri_to_id.get(uri) {
                // Look up ID by URI string
                result_ids[idx] = Some(id); // Store at original index
                self.cache.put(uri.clone(), id);
            }
        }

        // Convert Option<u64> to u64 (all should be Some at this point)
        Ok(result_ids
            .into_iter()
            .map(|id| id.expect("All URIs should have IDs"))
            .collect())
    }

    #[instrument(skip(self, uris), fields(table = %self.table_name, batch_size = uris.len()))]
    async fn register_uri_batch_hashmap(
        &self,
        uris: &[String],
    ) -> Result<std::collections::HashMap<String, u64>> {
        if uris.is_empty() {
            trace!("empty batch, returning early");
            return Ok(std::collections::HashMap::new());
        }

        // Validate all URIs first
        for uri in uris {
            Self::validate_uri(uri)?;
        }

        // CORRECTNESS GUARANTEE: URI-to-ID mapping accuracy
        // Each URI in the result HashMap is guaranteed to map to its correct ID
        // because SQL returns both 'id' and 'uri' together in each row (RETURNING id, uri).
        // We never rely on positional correspondence, eliminating ordering errors.

        let mut result = std::collections::HashMap::new();
        let mut uncached_uris = Vec::new();

        // Step 1: Deduplicate input and check cache
        let unique_uris: std::collections::HashSet<_> = uris.iter().collect();

        for uri in unique_uris {
            if let Some(id) = self.cache.get(uri) {
                result.insert(uri.clone(), id);
            } else {
                uncached_uris.push(uri.clone());
            }
        }

        // If everything was cached, return early
        if uncached_uris.is_empty() {
            debug!(cached = result.len(), "all URIs found in cache");
            return Ok(result);
        }

        debug!(
            cached = result.len(),
            uncached = uncached_uris.len(),
            "cache lookup complete, querying database"
        );

        // Step 2: Register uncached URIs in batch
        // Build query with validated table name (safe from SQL injection)
        let query = format!(
            r#"
            INSERT INTO {} (uri)
            SELECT unnest($1::text[])
            ON CONFLICT (uri_hash) DO UPDATE SET uri = EXCLUDED.uri
            RETURNING id, uri
            "#,
            self.table_name
        );

        // Execute with retry logic
        let client = self.pool.get().await.map_err(|e| {
            crate::error::Error::Database(format!("Failed to get database connection: {}", e))
        })?;

        let rows = client
            .query(&query, &[&uncached_uris])
            .await
            .map_err(|e| crate::error::Error::Database(e.to_string()))?;

        // Step 3: Add database results to result map and update cache
        // CORRECTNESS: Each row contains both URI and ID from same DB row,
        // guaranteeing correct mapping (no opportunity for misalignment)
        for row in rows {
            let uri: String = row.get("uri");
            let id: i64 = row.get("id");
            let id_u64 = id as u64;

            result.insert(uri.clone(), id_u64); // URI and ID are from same row
            self.cache.put(uri, id_u64);
        }

        Ok(result)
    }
}

/// Statistics about the URI register for observability and OpenTelemetry
#[derive(Debug, Clone)]
pub struct RegisterStats {
    /// Total number of URIs in the register
    pub total_uris: u64,
    /// Total storage size in bytes (includes indexes)
    pub size_bytes: u64,
    /// Cache performance metrics
    pub cache: crate::cache::CacheStats,
    /// Connection pool metrics
    pub pool: PoolStats,
}

/// Connection pool statistics for observability
#[derive(Debug, Clone)]
pub struct PoolStats {
    /// Number of connections currently being used
    pub connections_active: u32,
    /// Number of idle connections in the pool
    pub connections_idle: u32,
    /// Maximum number of connections allowed in the pool
    pub connections_max: u32,
}