Skip to main content

golem_ai_graph_janusgraph/
lib.rs

1mod client;
2mod connection;
3mod conversions;
4mod helpers;
5mod query_utils;
6mod schema;
7mod transaction;
8
9pub mod config;
10
11pub use config::JanusGraphConfig;
12#[cfg(feature = "golem")]
13pub use config::JanusGraphHostConfig;
14
15use client::JanusGraphApi;
16use golem_ai_graph::config::with_config_key;
17use golem_ai_graph::durability::{DurableGraph, ExtendedGuest};
18use golem_ai_graph::model::{connection::ConnectionConfig, errors::GraphError};
19use golem_ai_graph::TransactionProvider;
20use std::sync::Arc;
21
22pub struct JanusGraph;
23
24pub struct Graph {
25    pub api: Arc<JanusGraphApi>,
26}
27
28pub struct Transaction {
29    api: Arc<JanusGraphApi>,
30    state: std::sync::RwLock<TransactionState>,
31}
32
33#[derive(Debug, Clone, PartialEq)]
34enum TransactionState {
35    Active,
36    Committed,
37    RolledBack,
38}
39
40pub struct SchemaManager {
41    pub graph: Arc<Graph>,
42}
43
44impl ExtendedGuest for JanusGraph {
45    type Graph = Graph;
46    fn connect_internal(config: &ConnectionConfig) -> Result<Graph, GraphError> {
47        let host = with_config_key(config, "JANUSGRAPH_HOST")
48            .or_else(|| {
49                config
50                    .hosts
51                    .as_ref()
52                    .and_then(|hosts| hosts.first())
53                    .cloned()
54            })
55            .ok_or_else(|| GraphError::ConnectionFailed("Missing host".to_string()))?;
56
57        let port = with_config_key(config, "JANUSGRAPH_PORT")
58            .and_then(|p| p.parse().ok())
59            .or(config.port)
60            .unwrap_or(8182);
61        let username =
62            with_config_key(config, "JANUSGRAPH_USER").or_else(|| config.username.clone());
63
64        let password =
65            with_config_key(config, "JANUSGRAPH_PASSWORD").or_else(|| config.password.clone());
66
67        let api = JanusGraphApi::new(&host, port, username.as_deref(), password.as_deref())?;
68        api.execute("g.tx().open()", None)?;
69        Ok(Graph::new(api))
70    }
71}
72
73impl TransactionProvider for JanusGraph {
74    type Transaction = Transaction;
75}
76
77impl Graph {
78    fn new(api: JanusGraphApi) -> Self {
79        Self { api: Arc::new(api) }
80    }
81}
82
83impl Transaction {
84    fn new(api: Arc<JanusGraphApi>) -> Self {
85        Self {
86            api,
87            state: std::sync::RwLock::new(TransactionState::Active),
88        }
89    }
90}
91
92pub type DurableJanusGraph = DurableGraph<JanusGraph>;