Skip to main content

mnemo_pgwire/
lib.rs

1//! PostgreSQL wire protocol server for Mnemo.
2//!
3//! Exposes Mnemo's memory database through the PostgreSQL wire protocol,
4//! allowing SQL-native clients (psql, pgAdmin, any PostgreSQL driver) to
5//! query memories using familiar SQL syntax.
6//!
7//! # Supported SQL subset
8//!
9//! - `SELECT * FROM memories WHERE agent_id = '...' LIMIT n`
10//! - `INSERT INTO memories (content, importance, ...) VALUES (...)`
11//! - `DELETE FROM memories WHERE id = '...'`
12//!
13//! # Architecture
14//!
15//! The server accepts TCP connections and speaks the PostgreSQL wire protocol
16//! (startup, query, parse/bind/execute extended protocol). Queries are parsed
17//! and mapped to Mnemo engine operations:
18//!
19//! - `SELECT` → `engine.recall()`
20//! - `INSERT` → `engine.remember()`
21//! - `DELETE` → `engine.forget()`
22
23pub mod parser;
24pub mod server;
25
26use mnemo_core::query::MnemoEngine;
27use std::sync::Arc;
28
29/// Configuration for the pgwire server.
30#[derive(Debug, Clone)]
31pub struct PgWireConfig {
32    /// TCP bind address (defaults to localhost for security)
33    pub bind_addr: String,
34    /// Maximum concurrent connections
35    pub max_connections: usize,
36    /// Default agent ID for connections without explicit agent context
37    pub default_agent_id: String,
38    /// Optional password for cleartext password authentication.
39    /// When `None`, the server uses trust mode (no auth) — only safe on localhost.
40    pub password: Option<String>,
41}
42
43impl Default for PgWireConfig {
44    fn default() -> Self {
45        Self {
46            bind_addr: "127.0.0.1:5433".to_string(),
47            max_connections: 100,
48            default_agent_id: "default".to_string(),
49            password: None,
50        }
51    }
52}
53
54/// Start the pgwire server.
55///
56/// Listens on the configured address and accepts PostgreSQL wire protocol
57/// connections. Each connection is handled in a separate tokio task.
58///
59/// # Example
60///
61/// ```no_run
62/// # use std::sync::Arc;
63/// # use mnemo_pgwire::{PgWireConfig, start_server};
64/// # use mnemo_core::query::MnemoEngine;
65/// # async fn run(engine: Arc<MnemoEngine>) {
66/// let config = PgWireConfig::default();
67/// start_server(engine, config).await.unwrap();
68/// # }
69/// ```
70pub async fn start_server(
71    engine: Arc<MnemoEngine>,
72    config: PgWireConfig,
73) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
74    let listener = tokio::net::TcpListener::bind(&config.bind_addr).await?;
75    tracing::info!("pgwire server listening on {}", config.bind_addr);
76
77    let semaphore = Arc::new(tokio::sync::Semaphore::new(config.max_connections));
78
79    loop {
80        let (stream, addr) = listener.accept().await?;
81        tracing::debug!("pgwire connection from {addr}");
82
83        let engine = engine.clone();
84        let config = config.clone();
85        let permit = semaphore.clone().acquire_owned().await?;
86
87        tokio::spawn(async move {
88            if let Err(e) = server::handle_connection(stream, engine, &config).await {
89                tracing::warn!("pgwire connection error from {addr}: {e}");
90            }
91            drop(permit);
92        });
93    }
94}