genius_core_client/client/
mod.rs

1pub mod inference;
2use std::fmt;
3use std::path::Path;
4
5use crate::auth::auth_interceptor::AuthInterceptor;
6use crate::auth::jwt::decode_jwt;
7use crate::query::QueryEntitiesReturn;
8use crate::types::error::HstpError;
9use crate::upsert::upsert;
10use crate::utils::read_hsml_json;
11use crate::{query::query_t, types::entity::HSMLEntity};
12use kortex_gen_grpc::hstp::v1::hstp_service_client::HstpServiceClient;
13use kortex_gen_grpc::hstp::v1::CollisionStrategy;
14#[cfg(feature = "pyo3")]
15use pyo3::prelude::*;
16use serde_json::Value;
17use tonic::codegen::InterceptedService;
18use tonic::transport::{Channel, Endpoint};
19
20pub struct TimeoutAndRetries {
21    pub timeout: tokio::time::Duration,
22    pub retries: u32,
23}
24impl Default for TimeoutAndRetries {
25    fn default() -> Self {
26        Self {
27            timeout: tokio::time::Duration::from_secs(30),
28            retries: 3,
29        }
30    }
31}
32
33pub(crate) type InternalClient = HstpServiceClient<InterceptedService<Channel, AuthInterceptor>>;
34
35#[cfg(feature = "pyo3")]
36#[pyclass]
37pub struct Client {
38    client: InternalClient,
39    token: String,
40    retries: u32,
41}
42
43impl fmt::Debug for Client {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_struct("Client")
46            .field("token", &self.token)
47            .field("retries", &self.retries)
48            .finish()
49    }
50}
51
52#[derive(Default)]
53pub enum Protocol {
54    HTTP,
55    #[default]
56    HTTPS,
57}
58
59impl From<Protocol> for &str {
60    fn from(value: Protocol) -> Self {
61        match value {
62            Protocol::HTTP => "http",
63            Protocol::HTTPS => "https",
64        }
65    }
66}
67
68impl From<&str> for Protocol {
69    fn from(value: &str) -> Self {
70        match value.to_lowercase().as_str() {
71            "http" => Protocol::HTTP,
72            "https" => Protocol::HTTPS,
73            _ => panic!("Invalid protocol"),
74        }
75    }
76}
77
78#[cfg(not(feature = "napi"))]
79impl Client {
80    pub fn set_token<S: Into<String>>(&mut self, token: S) {
81        self.token = token.into();
82    }
83
84    pub fn get_token(&self) -> &str {
85        &self.token
86    }
87
88    pub async fn new_with_oauth2_token<S: Into<String>, T: Into<String>, U: Into<String>>(
89        protocol: Protocol,
90        host: S,
91        port: T,
92        token: U,
93        timeout_and_retries: Option<TimeoutAndRetries>,
94    ) -> Result<Self, HstpError> {
95        let timeout = timeout_and_retries.unwrap_or_default();
96        let token = token.into();
97        let client = Self::construct_internal_client(
98            protocol.into(),
99            host.into(),
100            port.into(),
101            token.clone(),
102            &timeout,
103        )
104        .await?;
105        //update this
106        Ok(Self {
107            client,
108            retries: timeout.retries,
109            token,
110        })
111    }
112
113    async fn construct_internal_client(
114        protocol: &str,
115        host: String,
116        port: String,
117        token: String,
118        timeout_and_retries: &TimeoutAndRetries,
119    ) -> Result<InternalClient, tonic::transport::Error> {
120        let genius_core_endpoint = format!("{}://{}:{}", protocol, host, port);
121        let connection_timeout_seconds = timeout_and_retries.timeout.as_secs();
122
123        let channel = Endpoint::from_shared(genius_core_endpoint)?
124            .tls_config(tonic::transport::ClientTlsConfig::default())?
125            .timeout(std::time::Duration::from_secs(connection_timeout_seconds))
126            .connect()
127            .await?;
128
129        let client = HstpServiceClient::with_interceptor(channel, AuthInterceptor { token });
130        let max_size = 2 * 1024 * 1024 * 1024; // 2GB
131        Ok(client
132            .max_decoding_message_size(max_size)
133            .max_encoding_message_size(max_size))
134    }
135
136    async fn refresh_token(&mut self) -> Result<(), HstpError> {
137        // TODO
138
139        Ok(())
140    }
141
142    pub async fn get_user_id(&mut self) -> Result<String, HstpError> {
143        let claims = decode_jwt(&self.token).await?;
144        let user_id = claims.sub;
145        Ok(user_id.clone())
146    }
147
148    pub async fn query_for_entity_array<S: Into<String>>(
149        &mut self,
150        query: S,
151    ) -> Result<QueryEntitiesReturn, HstpError> {
152        self.refresh_token().await?;
153        query_t::<QueryEntitiesReturn>(&mut self.client, query.into(), self.retries).await
154    }
155
156    pub async fn query_for_entity<S: Into<String>>(
157        &mut self,
158        query: S,
159    ) -> Result<HSMLEntity, HstpError> {
160        self.refresh_token().await?;
161        query_t::<HSMLEntity>(&mut self.client, query.into(), self.retries).await
162    }
163
164    pub async fn query_for_value_array<S: Into<String>>(
165        &mut self,
166        query: S,
167    ) -> Result<Vec<Value>, HstpError> {
168        self.refresh_token().await?;
169        query_t::<Vec<Value>>(&mut self.client, query.into(), self.retries).await
170    }
171
172    pub async fn query_for_value<S: Into<String>>(&mut self, query: S) -> Result<Value, HstpError> {
173        self.refresh_token().await?;
174        query_t::<Value>(&mut self.client, query.into(), self.retries).await
175    }
176
177    pub async fn query<S: Into<String>>(&mut self, query: S) -> Result<Value, HstpError> {
178        self.refresh_token().await?;
179        query_t::<Value>(&mut self.client, query.into(), self.retries).await
180    }
181
182    pub async fn upsert<V: AsRef<[HSMLEntity]>>(
183        &mut self,
184        entities: V,
185        collision_strategy: CollisionStrategy,
186    ) -> Result<Vec<HSMLEntity>, HstpError> {
187        self.refresh_token().await?;
188        upsert(
189            &mut self.client,
190            entities.as_ref(),
191            collision_strategy,
192            self.retries,
193        )
194        .await
195    }
196
197    pub async fn upsert_one(
198        &mut self,
199        entity: &HSMLEntity,
200        collision_strategy: CollisionStrategy,
201    ) -> Result<HSMLEntity, HstpError> {
202        self.refresh_token().await?;
203        let entities = vec![entity.clone()];
204        let entities = upsert(
205            &mut self.client,
206            &entities,
207            collision_strategy,
208            self.retries,
209        )
210        .await?;
211        Ok(entities[0].clone())
212    }
213
214    pub async fn upsert_hsml_json<P: AsRef<Path>>(
215        &mut self,
216        path: P,
217        collision_strategy: CollisionStrategy,
218    ) -> Result<Vec<HSMLEntity>, HstpError> {
219        let entities = read_hsml_json(path)?;
220        self.upsert(entities, collision_strategy).await
221    }
222
223    pub async fn create_listener(&mut self) -> Result<crate::listen::Listener, HstpError> {
224        self.refresh_token().await?;
225        crate::listen::Listener::new(&mut self.client).await
226    }
227}
genius_core_client/client/mod.rs

genius_core_client/client/
mod.rs
