awaken-stores 0.6.0

Storage backends (memory, file, PostgreSQL, SQLite mailbox) for Awaken agent state
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

//! PostgreSQL storage backend using `sqlx`.
//!
//! Tables are auto-created on first access via `ensure_schema()`.

use sqlx::PgPool;
use tokio::sync::Mutex;

/// PostgreSQL storage backend.
pub struct PostgresStore {
    pub(crate) pool: PgPool,
    pub(crate) threads_table: String,
    pub(crate) runs_table: String,
    pub(crate) messages_table: String,
    pub(crate) configs_table: String,
    config_notify_channel: String,
    schema_ready: Mutex<bool>,
}

impl Clone for PostgresStore {
    fn clone(&self) -> Self {
        Self {
            pool: self.pool.clone(),
            threads_table: self.threads_table.clone(),
            runs_table: self.runs_table.clone(),
            messages_table: self.messages_table.clone(),
            configs_table: self.configs_table.clone(),
            config_notify_channel: self.config_notify_channel.clone(),
            schema_ready: Mutex::new(false),
        }
    }
}

impl PostgresStore {
    /// Create a new store with default table names.
    pub fn new(pool: PgPool) -> Self {
        Self {
            pool,
            threads_table: "awaken_threads".to_string(),
            runs_table: "awaken_runs".to_string(),
            messages_table: "awaken_messages".to_string(),
            configs_table: "awaken_configs".to_string(),
            config_notify_channel: "awaken_config_changes".to_string(),
            schema_ready: Mutex::new(false),
        }
    }

    /// Create a new store with a custom table prefix.
    pub fn with_prefix(pool: PgPool, prefix: impl Into<String>) -> Self {
        let prefix = prefix.into();
        Self {
            pool,
            threads_table: format!("{prefix}_threads"),
            runs_table: format!("{prefix}_runs"),
            messages_table: format!("{prefix}_messages"),
            configs_table: format!("{prefix}_configs"),
            config_notify_channel: format!("{prefix}_config_changes"),
            schema_ready: Mutex::new(false),
        }
    }

    /// Name of the thread-scoped state table, derived from the threads
    /// table so a custom prefix carries through without an extra field.
    pub(crate) fn thread_states_table(&self) -> String {
        format!("{}_state", self.threads_table)
    }

    pub(crate) fn transaction_scope_descriptor(&self) -> String {
        let options = self.pool.connect_options();
        let socket = options
            .get_socket()
            .map(|path| path.display().to_string())
            .unwrap_or_default();
        format!(
            "pg::host={}::port={}::socket={}::user={}::database={}::options={}::threads={}::runs={}::messages={}::configs={}",
            options.get_host(),
            options.get_port(),
            socket,
            options.get_username(),
            options.get_database().unwrap_or_default(),
            options.get_options().unwrap_or_default(),
            self.threads_table,
            self.runs_table,
            self.messages_table,
            self.configs_table,
        )
    }

    pub(crate) fn thread_run_storage_identity_descriptor(&self) -> String {
        format!("pg-thread-run::{}", self.transaction_scope_descriptor())
    }
}

mod config;
mod pending;
mod run;
mod schema;
mod thread;

#[cfg(test)]
mod tests;