use std::{fmt::Debug, path::Path, sync::Arc};
use super::migrate::MigrationBuilder;
use super::model::Model;
use super::operations::{CharybdisModelBatch, Delete, Insert, ModelBatch, Update};
use super::query::{CharybdisQuery, ModelMutation, ModelRow, ModelStream, QueryExecutor};
use super::stream::CharybdisModelStream;
use super::Result;
use super::{ConnectionParams, CrudParams};
use charybdis::query::OptionalModelRow;
use charybdis::scylla::response::query_result::QueryResult;
use charybdis::scylla::serialize::row::SerializeRow;
use futures::future::join_all;
use futures::StreamExt;
use tracing::debug;
pub use scylla::client::caching_session::*;
pub use scylla::client::session::*;
pub use scylla::client::session_builder::*;
pub use scylla::client::*;
#[derive(Debug, Clone)]
pub struct Client {
session: Arc<CachingSession>,
crud_params: Option<CrudParams>,
}
impl Client {
pub async fn default() -> Result<Self> {
let con_params = ConnectionParams::default();
Self::connect(&con_params).await
}
pub fn from_session(session: &Arc<CachingSession>) -> Result<Self> {
Ok(Self {
session: session.clone(),
crud_params: None,
})
}
pub async fn connect(con_params: &ConnectionParams) -> Result<Self> {
debug!("Connecting to {}", con_params.uri);
let session = con_params.caching().await?;
let client = Self {
session: Arc::new(session),
crud_params: None,
};
if let Some(keyspace) = &con_params.use_keyspace {
if con_params.recreate_keyspace {
client.recreate_keyspace(keyspace).await?;
} else {
client.create_keyspace(keyspace).await?;
}
client.use_keyspace(keyspace).await?;
}
for filename in &con_params.init_files {
client.execute_file(filename).await?;
}
if con_params.migrate {
Self::migrate(client.session.get_session(), &con_params.use_keyspace).await?;
}
Ok(client)
}
}
impl Client {
pub fn with_params(mut self, params: impl Into<CrudParams>) -> Self {
_ = self.crud_params.insert(params.into());
self
}
}
impl Client {
pub fn session(&self) -> Arc<CachingSession> {
self.session.clone()
}
}
impl Client {
pub async fn get<'a, Val, E>(&self, query: CharybdisQuery<'a, Val, E, ModelRow>) -> Result<E>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send,
{
debug!("Get query: {}", query.query_string());
let res = self
.query_apply_params(query)
.execute(&self.session)
.await?;
Ok(res)
}
pub async fn get_optional<'a, Val, E>(
&self,
query: CharybdisQuery<'a, Val, E, OptionalModelRow>,
) -> Result<Option<E>>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send,
{
debug!("Get query: {}", query.query_string());
let res = self
.query_apply_params(query)
.execute(&self.session)
.await?;
Ok(res)
}
pub async fn get_many<'a, Val, E>(
&self,
queries: Vec<CharybdisQuery<'a, Val, E, ModelRow>>,
) -> Result<Vec<E>>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send + Clone,
{
let mut futures = vec![];
for query in queries {
let future = self.get::<Val, E>(query);
futures.push(future);
}
let results = join_all(futures).await;
let result: Vec<E> = results
.iter()
.filter_map(|result| match result {
Ok(entity) => Some(entity.clone()),
_ => None,
})
.collect();
Ok(result)
}
pub async fn count<'a, Val, E>(
&self,
query: CharybdisQuery<'a, Val, E, ModelStream>,
) -> Result<usize>
where
Val: SerializeRow + Sync + Send + Debug,
E: Model + Sync + Send + 'static,
{
Ok(self.stream(query).await?.count().await)
}
pub async fn update<E>(&self, entity: &E) -> Result<()>
where
E: Model + Update + Sync + Send + 'static,
{
self.update_query(entity.update()).await?;
Ok(())
}
async fn update_query<'a, Val, E>(
&self,
query: CharybdisQuery<'a, Val, E, ModelMutation>,
) -> Result<()>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send,
{
debug!("Update query: {}", query.query_string());
_ = self
.query_apply_params(query)
.execute(&self.session)
.await?;
Ok(())
}
pub async fn update_many<'a, E>(&self, iter: &[E], chunk_size: usize) -> Result<()>
where
E: ModelBatch<'a> + Sync + Send + 'a,
{
self.batch_apply_params(E::batch())
.chunked_update(&self.session, iter, chunk_size)
.await?;
Ok(())
}
pub async fn insert<E>(&self, entity: &E) -> Result<()>
where
E: Model + Insert + Sync + Send + 'static,
{
self.insert_query(entity.insert()).await?;
Ok(())
}
async fn insert_query<'a, Val, E>(
&self,
query: CharybdisQuery<'a, Val, E, ModelMutation>,
) -> Result<()>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send,
{
debug!("Insert query: {}", query.query_string());
_ = self
.query_apply_params(query)
.execute(&self.session)
.await?;
Ok(())
}
pub async fn insert_many<'a, E>(&self, iter: &[E], chunk_size: usize) -> Result<()>
where
E: ModelBatch<'a> + Sync + Send + 'a,
{
self.batch_apply_params(E::batch())
.chunked_insert(&self.session, iter, chunk_size)
.await?;
Ok(())
}
pub async fn delete<E>(&self, entity: &E) -> Result<()>
where
E: Model + Delete + Sync + Send + 'static,
{
self.delete_query(entity.delete()).await?;
Ok(())
}
async fn delete_query<'a, Val, E>(
&self,
query: CharybdisQuery<'a, Val, E, ModelMutation>,
) -> Result<()>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send,
{
debug!("Delete query: {}", query.query_string());
_ = self
.query_apply_params(query)
.execute(&self.session)
.await?;
Ok(())
}
pub async fn delete_many<'a, E>(&self, iter: &[E], chunk_size: usize) -> Result<()>
where
E: ModelBatch<'a> + Sync + Send + 'a,
{
self.batch_apply_params(E::batch())
.chunked_delete(&self.session, iter, chunk_size)
.await?;
Ok(())
}
pub async fn stream<'a, Val, E>(
&self,
query: CharybdisQuery<'a, Val, E, ModelStream>,
) -> Result<CharybdisModelStream<E>>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send + 'static,
{
debug!("Stream query: {}", query.query_string());
let res = self
.query_apply_params(query)
.execute(&self.session)
.await?;
Ok(res)
}
}
impl Client {
pub async fn drop_table(&self, name: &str) -> Result<()> {
let query = format!("DROP TABLE IF EXISTS {name};");
self.execute(&query, &[]).await?;
Ok(())
}
}
impl Client {
pub async fn keyspaces(&self) -> Result<Vec<String>> {
let query = "SELECT keyspace_name FROM system_schema.keyspaces;";
let res = self.session.execute_unpaged(query, &[]).await?;
let keyspaces: Vec<String> = res
.into_rows_result()?
.rows::<(String,)>()?
.filter_map(|s| s.ok()) .map(|(keyspace_name,)| keyspace_name) .collect();
Ok(keyspaces)
}
pub fn get_keyspace(&self) -> Option<String> {
let keyspace = self.session.get_session().get_keyspace();
keyspace.map(|k| k.to_string())
}
pub async fn use_keyspace(&self, name: &str) -> Result<()> {
self.session.get_session().use_keyspace(name, true).await?;
Ok(())
}
pub async fn recreate_keyspace(&self, name: &str) -> Result<()> {
self.drop_keyspace(name).await?;
self.create_keyspace(name).await?;
Ok(())
}
pub async fn with_recreate_keyspace(self, name: &str) -> Result<Self> {
self.recreate_keyspace(name).await?;
Ok(self)
}
pub async fn create_keyspace(&self, name: &str) -> Result<()> {
let query = format!("CREATE KEYSPACE IF NOT EXISTS {name} WITH REPLICATION = {{ 'class' : 'SimpleStrategy', 'replication_factor' : 1 }};");
self.execute(&query, &[]).await?;
Ok(())
}
pub async fn drop_keyspace(&self, name: &str) -> Result<()> {
let query = format!("DROP KEYSPACE IF EXISTS {name};");
self.execute(&query, &[]).await?;
Ok(())
}
pub async fn with_keyspace(self, name: &str) -> Result<Self> {
self.create_keyspace(name).await?;
Ok(self)
}
pub async fn with_keyspaces(self, names: &[&str]) -> Result<Self> {
for name in names {
self.create_keyspace(name).await?;
}
Ok(self)
}
pub async fn without_keyspace(self, name: &str) -> Result<Self> {
self.drop_keyspace(name).await?;
Ok(self)
}
pub async fn without_keyspaces(self, names: &[&str]) -> Result<Self> {
for name in names {
self.drop_keyspace(name).await?;
}
Ok(self)
}
}
impl Client {
pub async fn execute(&self, query: &str, values: impl SerializeRow) -> Result<QueryResult> {
debug!("Executing query: {}", query);
let res = self.session.execute_unpaged(query, values).await?;
Ok(res)
}
pub async fn execute_file(&self, filename: &str) -> Result<()> {
debug!("Init file '{}'", filename);
let current_path = std::env::current_dir().unwrap();
let file_path = Path::new(filename);
let full_path = current_path.join(file_path);
let raw_queries = tokio::fs::read_to_string(full_path)
.await
.unwrap_or_else(|_| panic!("Could not read file"));
let queries = raw_queries
.split(";")
.map(|query| query.trim())
.collect::<Vec<&str>>();
for query in queries {
if query.is_empty() {
continue;
}
self.execute(query, &[]).await?;
}
Ok(())
}
pub async fn migrate(session: &Session, use_keyspace: &Option<String>) -> Result<()> {
debug!("Migration started");
let mut builder = MigrationBuilder::new();
if let Some(keyspace) = use_keyspace {
builder = builder.keyspace(keyspace.to_owned());
}
let migration = builder.build(session).await;
migration.run().await;
Ok(())
}
fn batch_apply_params<'a, Val, E>(
&self,
batch: CharybdisModelBatch<'a, Val, E>,
) -> CharybdisModelBatch<'a, Val, E>
where
Val: SerializeRow + Sync + Send,
E: ModelBatch<'a>,
{
if let Some(params) = &self.crud_params {
params.apply_batch(batch)
} else {
batch
}
}
fn query_apply_params<'a, Val, E, Qe>(
&self,
query: CharybdisQuery<'a, Val, E, Qe>,
) -> CharybdisQuery<'a, Val, E, Qe>
where
Val: SerializeRow + Sync + Send,
E: Model + Sync + Send,
Qe: QueryExecutor<E>,
{
if let Some(params) = &self.crud_params {
params.apply_query(query)
} else {
query
}
}
}
#[cfg(test)]
mod tests {
type Result<T> = core::result::Result<T, Box<dyn std::error::Error>>;
use super::*;
use crate::scylla::{
charybdis::{self, macros::charybdis_model, types::Text},
Client, ConnectionParams,
};
#[charybdis_model(
table_name = users,
partition_keys = [id],
clustering_keys = [],
global_secondary_indexes = [name],
local_secondary_indexes = [],
)]
#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct Tst {
id: Text,
name: Option<Text>,
}
impl Tst {
fn with_id(id: &str) -> Self {
Self {
id: id.to_string(),
name: None,
}
}
fn with_name(mut self, name: impl Into<String>) -> Self {
_ = self.name.insert(name.into());
self
}
}
async fn get_client() -> Client {
let params = ConnectionParams {
migrate: false,
use_keyspace: Some("test".into()),
..Default::default()
};
let client = Client::connect(¶ms).await.unwrap();
client
.execute(
"
CREATE TABLE IF NOT EXISTS users (
id text PRIMARY KEY,
name text
);
",
&[],
)
.await
.unwrap();
client
.execute("CREATE INDEX IF NOT EXISTS ON users (name);", &[])
.await
.unwrap();
client
}
#[tokio::test]
async fn test_scylla_get() -> Result<()> {
let client = get_client().await;
let id = "test_scylla_get";
let model = Tst::with_id(id);
client.insert(&model).await?;
assert_eq!(model, client.get(Tst::find_by_id(id.into())).await?);
client.delete(&model).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_get_optional() -> Result<()> {
let client = get_client().await;
let id = "test_scylla_maybe_get";
assert!(client
.get_optional(Tst::maybe_find_first_by_id(id.into()))
.await?
.is_none());
let model = Tst::with_id(id);
client.insert(&model).await?;
assert!(client
.get_optional(Tst::maybe_find_first_by_id(id.into()))
.await?
.is_some());
client.delete(&model).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_get_many() -> Result<()> {
let client = get_client().await;
let fx_name = "test_scylla_get_many";
let ids = vec![
"test_scylla_get_many1",
"test_scylla_get_many2",
"test_scylla_get_many3",
];
let models = ids
.iter()
.map(|id| Tst::with_id(id).with_name(fx_name))
.collect::<Vec<Tst>>();
client.insert_many(&models, 3).await?;
let queries = ids
.iter()
.map(|id| Tst::find_by_id(id.to_string()))
.collect::<Vec<_>>();
let mut got = client.get_many(queries).await?;
got.sort();
assert_eq!(models[0], got[0]);
assert_eq!(models[1], got[1]);
assert_eq!(models[2], got[2]);
client.delete_many(&got, 3).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_stream() -> Result<()> {
let client = get_client().await;
let fx_name = "test_scylla_stream";
let models = [
Tst::with_id("test_scylla_stream1").with_name(fx_name),
Tst::with_id("test_scylla_stream2").with_name(fx_name),
Tst::with_id("test_scylla_stream3").with_name(fx_name),
Tst::with_id("test_scylla_stream4").with_name(fx_name),
Tst::with_id("test_scylla_stream5").with_name(fx_name),
Tst::with_id("test_scylla_stream6").with_name(fx_name),
];
client.insert_many(&models, 6).await?;
let mut stream = client
.stream(Tst::find_by_name(fx_name.to_string()))
.await?;
let mut got = vec![];
while let Some(Ok(model)) = stream.next().await {
got.push(model);
}
assert_eq!(6, got.len());
got.sort();
assert_eq!(models[0], got[0]);
assert_eq!(models[1], got[1]);
assert_eq!(models[2], got[2]);
assert_eq!(models[3], got[3]);
assert_eq!(models[4], got[4]);
assert_eq!(models[5], got[5]);
client.delete_many(&got, 6).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_insert() -> Result<()> {
let client = get_client().await;
let id = "test_scylla_insert";
let model = Tst::with_id(id);
client.insert(&model).await?;
assert_eq!(model, client.get(Tst::find_by_id(id.into())).await?);
client.delete(&model).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_insert_many() -> Result<()> {
let client = get_client().await;
let fx_name = "test_scylla_insert_many";
let models = [
Tst::with_id("test_scylla_insert_many1").with_name(fx_name),
Tst::with_id("test_scylla_insert_many2").with_name(fx_name),
Tst::with_id("test_scylla_insert_many3").with_name(fx_name),
];
client.insert_many(&models, 3).await?;
let mut find = client
.stream(Tst::find_by_name(fx_name.into()))
.await?
.try_collect()
.await?;
assert_eq!(3, find.len());
find.sort();
assert_eq!(models[0], find[0]);
assert_eq!(models[1], find[1]);
assert_eq!(models[2], find[2]);
client.delete_many(&find, 3).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_update() -> Result<()> {
let client = get_client().await;
let id = "test_scylla_update";
let model = Tst::with_id(id);
client.insert(&model).await?;
let updated = model.with_name("new name");
client.update(&updated).await?;
assert_eq!(updated, client.get(Tst::find_by_id(id.into())).await?);
client.delete(&updated).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_update_many() -> Result<()> {
let client = get_client().await;
let fx_name = "name";
let fx_new_name = "new_name";
let models = [
Tst::with_id("test_scylla_update_many1").with_name(fx_name),
Tst::with_id("test_scylla_update_many2").with_name(fx_name),
Tst::with_id("test_scylla_update_many3").with_name(fx_name),
];
client.insert_many(&models, 3).await?;
let mut find = client
.stream(Tst::find_by_name(fx_name.into()))
.await?
.try_collect()
.await?;
assert_eq!(3, find.len());
find.sort();
assert_eq!(models[0], find[0]);
assert_eq!(models[1], find[1]);
assert_eq!(models[2], find[2]);
let mut updated = vec![];
for model in models {
let new_model = model.with_name(fx_new_name);
updated.push(new_model);
}
client.update_many(&updated, 3).await?;
let mut find = client
.stream(Tst::find_by_name(fx_new_name.into()))
.await?
.try_collect()
.await?;
assert_eq!(3, find.len());
find.sort();
assert_eq!(updated[0], find[0]);
assert_eq!(updated[1], find[1]);
assert_eq!(updated[2], find[2]);
client.delete_many(&find, 3).await?;
Ok(())
}
#[tokio::test]
async fn test_scylla_delete() -> Result<()> {
let client = get_client().await;
let id = "test_scylla_delete";
let model = Tst::with_id(id);
client.insert(&model).await?;
client.delete(&model).await?;
assert!(client.get(Tst::find_by_id(id.into())).await.is_err());
Ok(())
}
#[tokio::test]
async fn test_scylla_delete_many() -> Result<()> {
let client = get_client().await;
let fx_name = "test_scylla_delete_many";
let models = [
Tst::with_id("test_scylla_delete_many1").with_name(fx_name),
Tst::with_id("test_scylla_delete_many2").with_name(fx_name),
Tst::with_id("test_scylla_delete_many3").with_name(fx_name),
];
client.insert_many(&models, 3).await?;
let count = client.count(Tst::find_by_name(fx_name.into())).await?;
assert_eq!(3, count);
client.delete_many(&models, 3).await?;
let count = client.count(Tst::find_by_name(fx_name.into())).await?;
assert_eq!(0, count);
Ok(())
}
#[tokio::test]
async fn test_scylla_count() -> Result<()> {
let client = get_client().await;
let fx_name = "test_scylla_count";
let models = [
Tst::with_id("test_scylla_count1").with_name(fx_name),
Tst::with_id("test_scylla_count2").with_name(fx_name),
Tst::with_id("test_scylla_count3").with_name(fx_name),
];
client.insert_many(&models, 3).await?;
let count = client.count(Tst::find_by_name(fx_name.into())).await?;
assert_eq!(3, count);
client.delete_many(&models, 3).await?;
let count = client.count(Tst::find_by_name(fx_name.into())).await?;
assert_eq!(0, count);
Ok(())
}
}