use std::sync::Arc;
use sqlx::PgPool;
use tonic::{Request, Response, Status};
use crate::metrics::{MetricsRecorder, NoopMetrics};
use crate::proto::udb::core::embedding::services::v1 as embedding_pb;
use crate::proto::udb::core::embedding::services::v1::embedding_service_server::EmbeddingService;
use crate::runtime::DataBrokerRuntime;
use crate::runtime::catalog::CatalogManager;
use crate::runtime::channels::ChannelManager;
pub use crate::proto::udb::core::embedding::services::v1::embedding_service_server::EmbeddingServiceServer;
use super::DataBrokerService;
mod chunking;
mod config;
mod errors;
mod events;
mod handlers;
mod model;
mod store;
#[cfg(test)]
mod tests;
mod workers;
pub(crate) use config::{EMBEDDING_WORK_EMITTER_BATCH, embedding_work_emitter_interval};
pub(crate) use workers::run_embedding_work_emitter_once;
pub struct EmbeddingServiceImpl {
pub(crate) pg_pool: Option<PgPool>,
pub(crate) runtime: Option<Arc<DataBrokerRuntime>>,
pub(crate) catalog: Option<Arc<CatalogManager>>,
pub(crate) outbox_relation: Option<String>,
pub(crate) channels: Option<ChannelManager>,
pub(crate) metrics: Arc<dyn MetricsRecorder>,
}
impl EmbeddingServiceImpl {
pub fn new() -> Self {
Self {
pg_pool: None,
runtime: None,
catalog: None,
outbox_relation: None,
channels: None,
metrics: Arc::new(NoopMetrics),
}
}
pub fn with_postgres(mut self, pool: Option<PgPool>) -> Self {
self.pg_pool = pool;
self
}
pub(crate) fn with_runtime(mut self, runtime: Option<Arc<DataBrokerRuntime>>) -> Self {
self.runtime = runtime;
self
}
pub(crate) fn with_catalog(mut self, catalog: Option<Arc<CatalogManager>>) -> Self {
self.catalog = catalog;
self
}
pub(crate) fn with_outbox(mut self, relation: Option<String>) -> Self {
self.outbox_relation = relation;
self
}
pub(crate) fn with_channels(mut self, channels: Option<ChannelManager>) -> Self {
self.channels = channels;
self
}
pub(crate) fn with_metrics(mut self, metrics: Arc<dyn MetricsRecorder>) -> Self {
self.metrics = metrics;
self
}
pub(crate) fn require_runtime(&self) -> Result<&DataBrokerRuntime, Status> {
self.runtime.as_deref().ok_or_else(|| {
errors::embedding_capability_status(
"native_entity_dispatch",
"runtime_native_entity_dispatch",
"embedding service requires runtime native-entity dispatch (no runtime configured)",
)
})
}
pub(crate) fn require_catalog(&self) -> Result<&CatalogManager, Status> {
self.catalog.as_deref().ok_or_else(|| {
errors::embedding_capability_status(
"catalog_lookup",
"active_catalog",
"embedding service requires the active catalog (no catalog configured)",
)
})
}
}
impl Default for EmbeddingServiceImpl {
fn default() -> Self {
Self::new()
}
}
#[tonic::async_trait]
impl EmbeddingService for EmbeddingServiceImpl {
async fn register_source(
&self,
request: Request<embedding_pb::RegisterSourceRequest>,
) -> Result<Response<embedding_pb::RegisterSourceResponse>, Status> {
handlers::register_source(self, request).await
}
async fn list_sources(
&self,
request: Request<embedding_pb::ListSourcesRequest>,
) -> Result<Response<embedding_pb::ListSourcesResponse>, Status> {
handlers::list_sources(self, request).await
}
async fn delete_source(
&self,
request: Request<embedding_pb::DeleteSourceRequest>,
) -> Result<Response<embedding_pb::DeleteSourceResponse>, Status> {
handlers::delete_source(self, request).await
}
async fn backfill(
&self,
request: Request<embedding_pb::BackfillRequest>,
) -> Result<Response<embedding_pb::BackfillResponse>, Status> {
handlers::backfill(self, request).await
}
async fn report_embedding(
&self,
request: Request<embedding_pb::ReportEmbeddingRequest>,
) -> Result<Response<embedding_pb::ReportEmbeddingResponse>, Status> {
handlers::report_embedding(self, request).await
}
async fn retrieve(
&self,
request: Request<embedding_pb::RetrieveRequest>,
) -> Result<Response<embedding_pb::RetrieveResponse>, Status> {
handlers::retrieve(self, request).await
}
}
impl DataBrokerService {
pub(crate) fn build_embedding_service(&self) -> EmbeddingServiceImpl {
let runtime = self.runtime.load_full();
let pg_pool = runtime
.native_store_pool_for_service("embedding", true, "")
.ok();
let outbox = runtime.config().cdc.outbox_relation();
let channels = Some(runtime.channels().clone());
EmbeddingServiceImpl::new()
.with_postgres(pg_pool)
.with_runtime(Some(runtime))
.with_catalog(Some(self.catalog.clone()))
.with_outbox(Some(outbox))
.with_channels(channels)
.with_metrics(self.metrics.clone())
}
}