use crate::connection::{PendingConnection, StdbConnection, StdbConnectionConfig};
use bevy_ecs::{
prelude::{Commands, Res, ResMut},
system::SystemParam,
};
use bevy_tasks::IoTaskPool;
use spacetimedb_sdk::{
__codegen::{DbConnection, SpacetimeModule},
DbContext,
};
#[derive(Clone, Debug, Default)]
pub struct StdbConnectOptions {
pub token: Option<String>,
pub uri: Option<String>,
pub database_name: Option<String>,
}
impl StdbConnectOptions {
pub fn from_token(token: impl Into<String>) -> Self {
Self {
token: Some(token.into()),
uri: None,
database_name: None,
}
}
pub fn from_uri(uri: impl Into<String>) -> Self {
Self {
token: None,
uri: Some(uri.into()),
database_name: None,
}
}
pub fn from_database_name(database_name: impl Into<String>) -> Self {
Self {
token: None,
uri: None,
database_name: Some(database_name.into()),
}
}
pub fn from_target(uri: impl Into<String>, database_name: impl Into<String>) -> Self {
Self {
token: None,
uri: Some(uri.into()),
database_name: Some(database_name.into()),
}
}
}
#[derive(SystemParam)]
pub struct StdbCommands<'w, 's, C, M>
where
C: DbConnection<Module = M> + DbContext + Send + Sync + 'static,
M: SpacetimeModule<DbConnection = C> + 'static,
{
config: ResMut<'w, StdbConnectionConfig<C, M>>,
connection: Option<Res<'w, StdbConnection<C>>>,
pending_connection: Option<Res<'w, PendingConnection<C>>>,
commands: Commands<'w, 's>,
}
impl<C, M> StdbCommands<'_, '_, C, M>
where
C: DbConnection<Module = M> + DbContext + Send + Sync + 'static,
M: SpacetimeModule<DbConnection = C> + 'static,
{
pub fn connect(&mut self, options: StdbConnectOptions) {
if self.connection.is_some() || self.pending_connection.is_some() {
return;
}
self.connect_impl(options);
}
pub fn reconnect(&mut self, options: StdbConnectOptions) {
self.disconnect();
self.connect_impl(options);
}
fn connect_impl(&mut self, options: StdbConnectOptions) {
if let Some(uri) = options.uri {
self.config.uri = uri;
}
if let Some(database_name) = options.database_name {
self.config.database_name = database_name;
}
if let Some(token) = options.token {
self.config.token = Some(token);
}
let config = self.config.clone();
let task = IoTaskPool::get().spawn(async move { config.build_connection().await });
self.commands.insert_resource(PendingConnection::<C>(task));
}
pub fn disconnect(&mut self) {
if let Some(conn) = &self.connection {
let _ = conn.disconnect();
}
self.commands.remove_resource::<StdbConnection<C>>();
self.commands.remove_resource::<PendingConnection<C>>();
}
}