use std::sync::Arc;
use super::{KalamLinkClient, KalamLinkClientBuilder, QueryUploadFile};
#[cfg(feature = "consumer")]
use crate::consumer::ConsumerBuilder;
use crate::{
auth::{AuthProvider, ResolvedAuth},
error::{KalamLinkError, Result},
event_handlers::EventHandlers,
models::file_upload::uploads_to_owned,
models::{FileUpload, LoginResponse, QueryResponse, SubscriptionConfig, SubscriptionInfo},
query::models::query_param::params_to_json,
query::models::QueryParam,
query::UploadProgressCallback,
subscription::{LiveRowsConfig, LiveRowsSubscription, SubscriptionManager},
timeouts::KalamLinkTimeouts,
};
impl KalamLinkClient {
pub fn builder() -> KalamLinkClientBuilder {
KalamLinkClientBuilder::new()
}
pub async fn execute_query(
&self,
sql: &str,
files: Option<Vec<FileUpload>>,
params: Option<Vec<QueryParam>>,
namespace_id: Option<&str>,
) -> Result<QueryResponse> {
self.execute_query_with_progress(sql, files, params, namespace_id, None).await
}
pub async fn execute_query_with_progress(
&self,
sql: &str,
files: Option<Vec<FileUpload>>,
params: Option<Vec<QueryParam>>,
namespace_id: Option<&str>,
progress: Option<UploadProgressCallback>,
) -> Result<QueryResponse> {
self.query_executor
.execute_with_progress_ref(
sql,
uploads_to_owned(files),
params_to_json(params),
namespace_id,
progress,
)
.await
}
pub async fn execute_query_with_tuples(
&self,
sql: &str,
files: Option<Vec<QueryUploadFile<'_>>>,
params: Option<Vec<QueryParam>>,
namespace_id: Option<&str>,
) -> Result<QueryResponse> {
let files = files.map(|items| {
items
.into_iter()
.map(|(placeholder, filename, data, mime)| {
let mut upload = FileUpload::new(placeholder, filename, data);
if let Some(mime) = mime {
upload = upload.with_mime(mime);
}
upload
})
.collect()
});
self.execute_query(sql, files, params, namespace_id).await
}
#[cfg(feature = "file-uploads")]
pub async fn execute_with_files(
&self,
sql: &str,
files: Vec<FileUpload>,
params: Option<Vec<QueryParam>>,
namespace_id: Option<&str>,
) -> Result<QueryResponse> {
self.execute_query(sql, Some(files), params, namespace_id).await
}
#[cfg(feature = "file-uploads")]
pub async fn execute_with_files_with_progress(
&self,
sql: &str,
files: Vec<FileUpload>,
params: Option<Vec<QueryParam>>,
namespace_id: Option<&str>,
progress: Option<UploadProgressCallback>,
) -> Result<QueryResponse> {
self.execute_query_with_progress(sql, Some(files), params, namespace_id, progress)
.await
}
pub async fn live_events(&self, query: &str) -> Result<SubscriptionManager> {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
let subscription_id = format!("sub_{}", nanos);
self.live_events_with_config(SubscriptionConfig::new(subscription_id, query))
.await
}
pub async fn live_events_with_config(
&self,
config: SubscriptionConfig,
) -> Result<SubscriptionManager> {
if self.connection_options.ws_lazy_connect {
let conn_guard = self.connection.lock().await;
if conn_guard.is_none() {
drop(conn_guard);
self.connect().await?;
}
}
let conn = {
let conn_guard = self.connection.lock().await;
conn_guard.clone()
};
if let Some(conn) = conn {
let (event_rx, result_rx) =
conn.subscribe_send(config.id.clone(), config.sql, config.options).await?;
let shared_control = conn.subscription_control();
let (generation, resume_from) = result_rx.await.map_err(|_| {
KalamLinkError::WebSocketError(
"Connection task died before confirming subscribe".to_string(),
)
})??;
return Ok(SubscriptionManager::from_shared(
config.id,
event_rx,
shared_control,
generation,
resume_from,
&self.timeouts,
));
}
Err(KalamLinkError::WebSocketError(
"Not connected. Call connect() before opening live streams.".to_string(),
))
}
pub async fn live(&self, query: &str) -> Result<LiveRowsSubscription> {
self.live_with_config(
SubscriptionConfig::new(
format!(
"live_rows_{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos()
),
query,
),
LiveRowsConfig::default(),
)
.await
}
pub async fn live_with_config(
&self,
config: SubscriptionConfig,
live_rows_config: LiveRowsConfig,
) -> Result<LiveRowsSubscription> {
let subscription = self.live_events_with_config(config).await?;
Ok(LiveRowsSubscription::new(subscription, live_rows_config))
}
pub async fn connect(&self) -> Result<()> {
{
let conn_guard = self.connection.lock().await;
if conn_guard.is_some() {
return Ok(());
}
}
let resolved_auth = match self.fresh_auth().await? {
AuthProvider::BasicAuth(user, password) => {
let login_response = self.exchange_login_credentials(&user, &password).await?;
AuthProvider::jwt_token(login_response.access_token)
},
auth => auth,
};
self.update_shared_auth(resolved_auth);
let conn = Arc::new(
crate::connection::SharedConnection::connect(
self.base_url.clone(),
self.shared_resolved_auth.clone(),
self.timeouts.clone(),
self.connection_options.clone(),
self.event_handlers.clone(),
)
.await?,
);
let mut conn_guard = self.connection.lock().await;
if conn_guard.is_none() {
*conn_guard = Some(conn);
} else {
drop(conn_guard);
conn.disconnect().await;
}
Ok(())
}
pub async fn disconnect(&self) {
let conn = {
let mut guard = self.connection.lock().await;
guard.take()
};
if let Some(conn) = conn {
conn.disconnect().await;
}
}
pub async fn cancel_subscription(&self, id: &str) -> Result<()> {
let conn = {
let guard = self.connection.lock().await;
guard.clone()
};
if let Some(conn) = conn {
conn.unsubscribe(id).await?;
}
Ok(())
}
pub async fn is_connected(&self) -> bool {
let guard = self.connection.lock().await;
guard.as_ref().is_some_and(|conn| conn.is_connected())
}
pub async fn subscriptions(&self) -> Vec<SubscriptionInfo> {
let conn = {
let guard = self.connection.lock().await;
guard.clone()
};
match conn.as_ref() {
Some(conn) => conn.list_subscriptions().await,
None => Vec::new(),
}
}
pub fn event_handlers(&self) -> &EventHandlers {
&self.event_handlers
}
pub fn timeouts(&self) -> &KalamLinkTimeouts {
&self.timeouts
}
#[cfg(feature = "consumer")]
pub fn consumer(&self) -> ConsumerBuilder {
ConsumerBuilder::from_client(self.clone())
}
#[cfg(feature = "consumer")]
pub(crate) fn auth(&self) -> &AuthProvider {
&self.auth
}
pub fn resolved_auth(&self) -> &ResolvedAuth {
&self.resolved_auth
}
pub fn set_auth(&mut self, auth: AuthProvider) {
self.auth = auth.clone();
self.query_executor.set_auth(auth.clone());
let resolved = ResolvedAuth::Static(auth);
self.resolved_auth = resolved.clone();
*self.shared_resolved_auth.write().unwrap() = resolved;
}
pub fn update_shared_auth(&self, auth: AuthProvider) {
self.query_executor.set_auth(auth.clone());
let resolved = ResolvedAuth::Static(auth);
*self.shared_resolved_auth.write().unwrap() = resolved;
}
pub async fn fresh_auth(&self) -> Result<AuthProvider> {
self.resolved_auth.resolve().await
}
pub(crate) fn base_url(&self) -> &str {
&self.base_url
}
pub(crate) fn http_client(&self) -> reqwest::Client {
self.http_client.clone()
}
pub(crate) async fn jwt_for_http_request(&self) -> Result<AuthProvider> {
match self.fresh_auth().await? {
AuthProvider::BasicAuth(user, password) => {
let login = self.exchange_login_credentials(&user, &password).await?;
Ok(AuthProvider::jwt_token(login.access_token))
},
auth @ (AuthProvider::JwtToken(_) | AuthProvider::None) => Ok(auth),
}
}
async fn exchange_login_credentials(
&self,
user: &str,
password: &str,
) -> Result<LoginResponse> {
let url = format!("{}/v1/api/auth/login", self.base_url);
let body = serde_json::json!({
"user": user,
"password": password,
});
let response = self.http_client.post(&url).json(&body).send().await?;
let status = response.status();
if !status.is_success() {
let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
return Err(KalamLinkError::AuthenticationError(format!(
"Login failed during auth exchange ({}): {}",
status, error_text
)));
}
Ok(response.json::<LoginResponse>().await?)
}
}