use std::{fmt, future::Future, sync::Arc};
#[cfg(feature = "surrealdb")]
use crate::databases::surrealdb::backend::SurrealBackend;
use crate::{
CapabilitySet, CatalogRequest, CatalogResponse, DataRequest, DataResponse, IndexRequest,
IndexResponse, QueryRequest, QueryResponse, StorageError,
databases::{config::DatabaseConfig, lancedb::backend::LanceBackend},
};
pub trait DatabaseBackend: Send + Sync + 'static {
fn capabilities(&self) -> CapabilitySet;
fn catalog(
&self,
request: CatalogRequest,
) -> impl Future<Output = Result<CatalogResponse, StorageError>> + Send + '_;
fn data(
&self,
request: DataRequest,
) -> impl Future<Output = Result<DataResponse, StorageError>> + Send + '_;
fn query(
&self,
request: QueryRequest,
) -> impl Future<Output = Result<QueryResponse, StorageError>> + Send + '_;
fn index(
&self,
request: IndexRequest,
) -> impl Future<Output = Result<IndexResponse, StorageError>> + Send + '_;
}
#[derive(Clone)]
pub struct Database {
inner: Arc<DatabaseInner>,
}
enum DatabaseInner {
Lance(LanceBackend),
#[cfg(feature = "surrealdb")]
Surreal(SurrealBackend),
}
impl Database {
pub async fn connect(config: DatabaseConfig) -> Result<Self, StorageError> {
config.validate()?;
match config {
DatabaseConfig::Lance { config, common } => LanceBackend::connect(common, config)
.await
.map(|backend| Self {
inner: Arc::new(DatabaseInner::Lance(backend)),
}),
DatabaseConfig::Surreal { config, common } => {
if cfg!(feature = "surrealdb") {
#[cfg(feature = "surrealdb")]
return SurrealBackend::connect(common, config)
.await
.map(|backend| Self {
inner: Arc::new(DatabaseInner::Surreal(backend)),
});
}
drop((config, common));
Err(StorageError::InvalidConfig(
"SurrealDB requires the `surrealdb` feature".into(),
))
}
}
}
}
impl fmt::Debug for Database {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Database { .. }")
}
}
impl DatabaseBackend for Database {
fn capabilities(&self) -> CapabilitySet {
self.inner.capabilities()
}
fn catalog(
&self,
request: CatalogRequest,
) -> impl Future<Output = Result<CatalogResponse, StorageError>> + Send + '_ {
self.inner.catalog(request)
}
fn data(
&self,
request: DataRequest,
) -> impl Future<Output = Result<DataResponse, StorageError>> + Send + '_ {
self.inner.data(request)
}
fn query(
&self,
request: QueryRequest,
) -> impl Future<Output = Result<QueryResponse, StorageError>> + Send + '_ {
self.inner.query(request)
}
fn index(
&self,
request: IndexRequest,
) -> impl Future<Output = Result<IndexResponse, StorageError>> + Send + '_ {
self.inner.index(request)
}
}
impl DatabaseBackend for DatabaseInner {
fn capabilities(&self) -> CapabilitySet {
match self {
Self::Lance(backend) => backend.capabilities(),
#[cfg(feature = "surrealdb")]
Self::Surreal(backend) => backend.capabilities(),
}
}
async fn catalog(&self, request: CatalogRequest) -> Result<CatalogResponse, StorageError> {
match self {
Self::Lance(backend) => backend.catalog(request).await,
#[cfg(feature = "surrealdb")]
Self::Surreal(backend) => backend.catalog(request).await,
}
}
async fn data(&self, request: DataRequest) -> Result<DataResponse, StorageError> {
request.validate()?;
match self {
Self::Lance(backend) => backend.data(request).await,
#[cfg(feature = "surrealdb")]
Self::Surreal(backend) => backend.data(request).await,
}
}
async fn query(&self, request: QueryRequest) -> Result<QueryResponse, StorageError> {
match self {
Self::Lance(backend) => backend.query(request).await,
#[cfg(feature = "surrealdb")]
Self::Surreal(backend) => backend.query(request).await,
}
}
async fn index(&self, request: IndexRequest) -> Result<IndexResponse, StorageError> {
request.validate()?;
match self {
Self::Lance(backend) => backend.index(request).await,
#[cfg(feature = "surrealdb")]
Self::Surreal(backend) => backend.index(request).await,
}
}
}