use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use crate::WyrdClient;
#[cfg(feature = "internal")]
use crate::config::ClientConfig;
#[cfg(feature = "internal")]
use crate::environment::Environment;
#[cfg(feature = "internal")]
use crate::error::WyrdClientError;
#[cfg(feature = "internal")]
use crate::global_config::GlobalConfig;
use arrow::datatypes::SchemaRef;
use arrow::json::WriterBuilder;
use arrow::json::writer::JsonArray;
use arrow::record_batch::RecordBatch;
use serde::de::DeserializeOwned;
use tokio::sync::Mutex as AsyncMutex;
use wyrd_queue::QueueConfig;
use wyrd_queue::variant::WyrdJsonEncoderFactory;
#[cfg(feature = "internal")]
use wyrd_queue::{BatchSink, ClientByteGuard, DurableBatchAck, SealedBatch, SinkError};
use wyrd_spec::error::WyrdError;
use wyrd_spec::request_id::RequestId;
use wyrd_spec::vala::api::{
BifrostQueryRequest, BifrostTableDescription, BifrostTableEntry, CancelRunningQueryResponse,
QueryParam, QueryTerminalFrame, RegisterOutcome, RegisterTableResponse, RunningQuerySummary,
};
use wyrd_spec::vala::ids::TraceId;
use wyrd_spec::vala::trace::SpanRecord;
#[cfg(feature = "internal")]
use crate::bifrost::BifrostMetrics;
use crate::bifrost::grpc::BifrostGrpcTransport;
use crate::bifrost::handle::WriterPool;
use crate::bifrost::query::{BifrostClientError, QueryClient, QueryResultStream};
#[cfg(feature = "internal")]
use crate::bifrost::query::{CollectedQueryLimits, CollectedQueryResult};
use crate::bifrost::rows::{SPAN_COLUMNS, span_record};
use crate::bifrost::scope::ClientScope;
use crate::bifrost::sink::BifrostIngestSink;
use crate::bifrost::table::{Correlation, TableConfig, WriterTable};
const SPANS_NOT_CREATED: &str = "WYRD_VALA_404_BIFROST_TABLE_NOT_FOUND";
#[derive(Debug, Default)]
pub struct BifrostOptions {
pub table: Option<TableConfig>,
pub queue: QueueConfig,
}
pub struct Bifrost {
query: QueryClient,
writer: Arc<WriterPool>,
active: Mutex<Option<TableConfig>>,
described: Mutex<HashMap<Arc<str>, WriterTable>>,
describe_gate: AsyncMutex<()>,
}
impl Bifrost {
pub async fn connect(options: BifrostOptions) -> Result<Self, BifrostClientError> {
let client = WyrdClient::from_global().map_err(BifrostClientError::from)?;
Self::connect_with(&client, options).await
}
pub async fn connect_with(
client: &WyrdClient,
options: BifrostOptions,
) -> Result<Self, BifrostClientError> {
let BifrostOptions { table, queue } = options;
queue.validate().map_err(BifrostClientError::Queue)?;
let client = client.connected().await?;
Self::assemble(&client, table, queue).await
}
#[must_use]
#[cfg(feature = "internal")]
pub fn with_sink(
client: &WyrdClient,
table: Option<TableConfig>,
sink: Arc<dyn BatchSink<ClientByteGuard>>,
config: QueueConfig,
) -> Self {
Self {
query: QueryClient::new(client),
writer: Arc::new(WriterPool::new(
ClientScope::from_client(client),
sink,
config,
)),
active: Mutex::new(table),
described: Mutex::new(HashMap::new()),
describe_gate: AsyncMutex::new(()),
}
}
#[must_use]
#[cfg(feature = "internal")]
pub fn query_only(client: &WyrdClient) -> Self {
Self::with_sink(
client,
None,
Arc::new(QueryOnlySink),
QueueConfig::default(),
)
}
pub(crate) async fn assemble(
client: &WyrdClient,
table: Option<TableConfig>,
mut config: QueueConfig,
) -> Result<Self, BifrostClientError> {
let transport = BifrostGrpcTransport::connect(client)
.await
.map_err(BifrostClientError::from)?;
let deadline_ms = u64::try_from(transport.send_deadline().as_millis()).unwrap_or(u64::MAX);
config.flush_timeout_ms = config.flush_timeout_ms.max(deadline_ms);
Ok(Self {
query: QueryClient::new(client),
writer: Arc::new(WriterPool::new(
ClientScope::from_client(client),
Arc::new(BifrostIngestSink {
transport: Arc::new(transport),
}),
config,
)),
active: Mutex::new(table),
described: Mutex::new(HashMap::new()),
describe_gate: AsyncMutex::new(()),
})
}
pub async fn register(&self) -> Result<RegisterOutcome, BifrostClientError> {
let request = {
let active = self.active.lock().expect("active table lock poisoned");
active
.as_ref()
.ok_or(BifrostClientError::NoActiveTable)?
.register_request()
};
let response: RegisterTableResponse = self
.query
.client()
.http
.request_json(reqwest::Method::POST, "/v1/bifrost/tables", Some(&request))
.await?;
let mut active = self.active.lock().expect("active table lock poisoned");
if let Some(table) = active.as_mut()
&& table.register_request() == request
{
table.resolve(&response);
}
Ok(response.outcome)
}
pub fn use_table(&self, table: TableConfig) -> Option<TableConfig> {
self.active
.lock()
.expect("active table lock poisoned")
.replace(table)
}
pub async fn use_table_by_name(&self, fqn: &str) -> Result<(), BifrostClientError> {
let table = TableConfig::describe_with(self.query.client(), fqn).await?;
self.use_table(table);
Ok(())
}
#[cfg(feature = "internal")]
pub async fn writer_table(&self, fqn: &str) -> Result<WriterTable, BifrostClientError> {
self.resolve_writer_table(fqn).await
}
pub(crate) async fn resolve_writer_table(
&self,
fqn: &str,
) -> Result<WriterTable, BifrostClientError> {
if let Some(table) = self.described_table(fqn) {
return Ok(table);
}
let _gate = self.describe_gate.lock().await;
if let Some(table) = self.described_table(fqn) {
return Ok(table);
}
let config = TableConfig::describe_with(self.query.client(), fqn).await?;
let described = WriterTable::new(fqn, config.user_schema().clone());
self.described
.lock()
.expect("described table lock poisoned")
.insert(Arc::from(fqn), described.clone());
Ok(described)
}
#[must_use]
#[cfg(feature = "internal")]
pub fn cached_writer_table(&self, fqn: &str) -> Option<WriterTable> {
self.described_table(fqn)
}
fn described_table(&self, fqn: &str) -> Option<WriterTable> {
self.described
.lock()
.expect("described table lock poisoned")
.get(fqn)
.cloned()
}
#[cfg(feature = "internal")]
pub fn insert_into(
&self,
table: &WriterTable,
row: Vec<u8>,
correlation: Correlation,
) -> Result<(), BifrostClientError> {
self.writer
.insert(
table.fqn(),
table.user_schema(),
row,
correlation.card_uid,
correlation.run_id,
)
.map_err(Into::into)
}
#[cfg(feature = "internal")]
pub fn insert_rows_into(
&self,
table: &WriterTable,
rows: Vec<Vec<u8>>,
correlation: Correlation,
) -> Result<(), BifrostClientError> {
self.insert_observation(table, rows, correlation)
}
pub(crate) fn insert_observation(
&self,
table: &WriterTable,
rows: Vec<Vec<u8>>,
correlation: Correlation,
) -> Result<(), BifrostClientError> {
self.writer
.insert_rows(
table.fqn(),
table.user_schema(),
rows,
correlation.card_uid,
correlation.run_id,
)
.map_err(Into::into)
}
#[must_use]
pub fn table(&self) -> Option<TableConfig> {
self.active
.lock()
.expect("active table lock poisoned")
.clone()
}
pub fn insert(&self, row: Vec<u8>, correlation: Correlation) -> Result<(), BifrostClientError> {
let (table, schema) = {
let active = self.active.lock().expect("active table lock poisoned");
let table = active.as_ref().ok_or(BifrostClientError::NoActiveTable)?;
(table.fqn(), table.user_schema().clone())
};
self.writer
.insert(
&table,
&schema,
row,
correlation.card_uid,
correlation.run_id,
)
.map_err(Into::into)
}
pub async fn write_batch(
&self,
table: &str,
batch: &RecordBatch,
) -> Result<(), BifrostClientError> {
self.writer
.write_batch(table, batch)
.await
.map_err(Into::into)
}
#[cfg(feature = "internal")]
pub fn enqueue_batch(
&self,
table: &str,
batch: RecordBatch,
request_id: Option<RequestId>,
) -> Result<(), BifrostClientError> {
self.writer
.enqueue_batch(table, batch, request_id)
.map_err(Into::into)
}
pub async fn flush(&self) -> Result<(), BifrostClientError> {
self.drain(Drain::Flush).await
}
pub async fn shutdown(&self) -> Result<(), BifrostClientError> {
self.drain(Drain::Shutdown).await
}
async fn drain(&self, kind: Drain) -> Result<(), BifrostClientError> {
let writer = Arc::clone(&self.writer);
tokio::task::spawn_blocking(move || match kind {
Drain::Flush => writer.flush(),
Drain::Shutdown => writer.shutdown(),
})
.await
.expect("bifrost drain task joins")
.map_err(Into::into)
}
pub async fn sql(
&self,
query: &str,
params: &[QueryParam],
) -> Result<QueryResult, BifrostClientError> {
let mut stream = self.stream(query, params, None).await?;
let mut batches = Vec::new();
while let Some(batch) = stream.next_batch().await? {
batches.push(batch);
}
let schema = stream.schema().cloned().ok_or_else(|| {
BifrostClientError::Protocol("query stream omitted its schema".to_owned())
})?;
let terminal = stream
.terminal()
.cloned()
.ok_or(BifrostClientError::IncompleteQueryStream)?;
Ok(QueryResult {
batches,
schema,
terminal,
})
}
pub async fn sql_as<T: DeserializeOwned>(
&self,
query: &str,
params: &[QueryParam],
) -> Result<Vec<T>, BifrostClientError> {
self.sql(query, params).await?.deserialize()
}
pub async fn trace_timeline(
&self,
trace_id: TraceId,
) -> Result<Vec<SpanRecord>, BifrostClientError> {
let query = format!(
"SELECT {SPAN_COLUMNS} FROM vala.traces.spans \
WHERE trace_id = decode($1, 'hex') OR trace_id IN ( \
SELECT trace_id FROM ( \
SELECT trace_id, unnest(links) AS link FROM vala.traces.spans) \
WHERE link['trace_id'] = decode($1, 'hex'))"
);
let result = match self
.sql(&query, &[QueryParam::String(trace_id.to_hex())])
.await
{
Ok(result) => result,
Err(error) if WyrdError::from(&error).code() == SPANS_NOT_CREATED => {
return Ok(Vec::new());
}
Err(error) => return Err(error),
};
let mut spans = result
.batches()
.iter()
.flat_map(|batch| (0..batch.num_rows()).map(move |row| span_record(batch, row)))
.collect::<Result<Vec<_>, _>>()
.map_err(BifrostClientError::RowDeserialization)?;
spans.sort_by(|left, right| {
(
left.start_time,
left.trace_id.as_bytes(),
left.span_id.as_bytes(),
)
.cmp(&(
right.start_time,
right.trace_id.as_bytes(),
right.span_id.as_bytes(),
))
});
let mut seen = HashSet::new();
spans.retain(|span| seen.insert((span.trace_id, span.span_id)));
Ok(spans)
}
pub async fn stream(
&self,
query: &str,
params: &[QueryParam],
deadline: Option<Duration>,
) -> Result<QueryResultStream, BifrostClientError> {
self.query
.query(&BifrostQueryRequest {
params: params.to_vec(),
sql: query.to_owned(),
deadline_ms: deadline
.map(|deadline| i64::try_from(deadline.as_millis()).unwrap_or(i64::MAX)),
})
.await
}
#[cfg(feature = "internal")]
pub async fn query(
&self,
request: &BifrostQueryRequest,
) -> Result<QueryResultStream, BifrostClientError> {
self.query.query(request).await
}
#[cfg(feature = "internal")]
pub async fn collect_bounded(
&self,
request: &BifrostQueryRequest,
limits: CollectedQueryLimits,
) -> Result<CollectedQueryResult, BifrostClientError> {
self.query.collect_bounded(request, limits).await
}
pub async fn running(&self) -> Result<Vec<RunningQuerySummary>, BifrostClientError> {
self.query.running().await
}
pub async fn status(
&self,
request_id: &RequestId,
) -> Result<RunningQuerySummary, BifrostClientError> {
self.query.status(request_id).await
}
pub async fn cancel(
&self,
request_id: &RequestId,
) -> Result<CancelRunningQueryResponse, BifrostClientError> {
self.query.cancel(request_id).await
}
pub async fn list_tables(&self) -> Result<Vec<BifrostTableEntry>, BifrostClientError> {
self.query.list_tables().await
}
pub async fn describe_table(
&self,
namespace: &str,
name: &str,
) -> Result<BifrostTableDescription, BifrostClientError> {
self.query.describe_table(namespace, name).await
}
#[must_use]
#[cfg(feature = "internal")]
pub fn scope(&self) -> &ClientScope {
self.writer.scope()
}
#[must_use]
pub fn producer_count(&self) -> usize {
self.writer.producer_count()
}
#[must_use]
pub fn dropped(&self) -> u64 {
self.writer.dropped()
}
#[must_use]
#[cfg(feature = "internal")]
pub fn metrics(&self) -> BifrostMetrics {
self.writer.metrics()
}
#[cfg(feature = "internal")]
pub fn observe_losses(&self, observer: impl Fn(u64) + Send + Sync + 'static) {
self.writer.observe_losses(observer);
}
pub(crate) fn writer(&self) -> &WriterPool {
&self.writer
}
}
#[derive(Debug, Clone, Copy)]
enum Drain {
Flush,
Shutdown,
}
#[derive(Debug)]
pub struct QueryResult {
batches: Vec<RecordBatch>,
schema: SchemaRef,
terminal: QueryTerminalFrame,
}
impl QueryResult {
#[must_use]
pub fn batches(&self) -> &[RecordBatch] {
&self.batches
}
#[must_use]
pub fn schema(&self) -> &SchemaRef {
&self.schema
}
#[must_use]
pub fn terminal(&self) -> &QueryTerminalFrame {
&self.terminal
}
#[must_use]
pub fn num_rows(&self) -> usize {
self.batches.iter().map(RecordBatch::num_rows).sum()
}
fn deserialize<T: DeserializeOwned>(&self) -> Result<Vec<T>, BifrostClientError> {
let mut bytes = Vec::new();
{
let mut writer = WriterBuilder::new()
.with_encoder_factory(Arc::new(WyrdJsonEncoderFactory))
.build::<_, JsonArray>(&mut bytes);
for batch in &self.batches {
writer
.write(batch)
.map_err(|error| BifrostClientError::Arrow(error.to_string()))?;
}
writer
.finish()
.map_err(|error| BifrostClientError::Arrow(error.to_string()))?;
}
if bytes.is_empty() {
return Ok(Vec::new());
}
serde_json::from_slice(&bytes)
.map_err(|error| BifrostClientError::RowDeserialization(error.to_string()))
}
pub fn to_bytes(&self) -> Result<Vec<u8>, BifrostClientError> {
let mut buffer = Vec::new();
{
let mut writer = arrow::ipc::writer::StreamWriter::try_new(&mut buffer, &self.schema)
.map_err(|error| BifrostClientError::Arrow(error.to_string()))?;
for batch in &self.batches {
writer
.write(batch)
.map_err(|error| BifrostClientError::Arrow(error.to_string()))?;
}
writer
.finish()
.map_err(|error| BifrostClientError::Arrow(error.to_string()))?;
}
Ok(buffer)
}
}
#[must_use]
#[cfg(feature = "internal")]
pub fn register_outcome_name(outcome: RegisterOutcome) -> &'static str {
match outcome {
RegisterOutcome::Created => "created",
RegisterOutcome::AlreadyExists => "already_exists",
}
}
#[cfg(feature = "internal")]
pub fn client_from_options(
server_url: Option<&str>,
credential: Option<&str>,
grpc_url: Option<&str>,
) -> Result<WyrdClient, BifrostClientError> {
client_from_options_in(Environment::Process, server_url, credential, grpc_url)
}
#[cfg(feature = "internal")]
fn client_from_options_in(
environment: Environment,
server_url: Option<&str>,
credential: Option<&str>,
grpc_url: Option<&str>,
) -> Result<WyrdClient, BifrostClientError> {
for (field, value) in [
("server_url", server_url),
("grpc_url", grpc_url),
("credential", credential),
] {
if value.is_some_and(|value| value.trim().is_empty()) {
return Err(BifrostClientError::Client(WyrdClientError::Config {
field: field.to_owned(),
reason: "must not be empty".to_owned(),
}));
}
}
let global = GlobalConfig::load_in(&environment).map_err(BifrostClientError::from)?;
let mut config = ClientConfig::from_environment(environment, &global, server_url, grpc_url);
if let Some(credential) = credential {
config.credential = Some(secrecy::SecretString::from(credential.to_owned()));
}
WyrdClient::with_config(config).map_err(BifrostClientError::from)
}
#[cfg(feature = "internal")]
struct QueryOnlySink;
#[cfg(feature = "internal")]
#[async_trait::async_trait]
impl BatchSink<ClientByteGuard> for QueryOnlySink {
async fn send(
&self,
batch: &SealedBatch<ClientByteGuard>,
) -> Result<DurableBatchAck, SinkError> {
Err(wyrd_queue::SinkError::Terminal(
wyrd_spec::error::WyrdError::Validation {
message: "this Bifrost client was built with Bifrost::query_only and cannot write"
.to_owned(),
details: serde_json::json!({ "table": batch.table }),
},
))
}
}
#[cfg(all(test, feature = "internal"))]
mod tests {
use super::client_from_options_in;
use crate::environment::Environment;
#[test]
fn omitted_server_url_comes_from_the_global_config_file() {
let config_dir = tempfile::tempdir().expect("temporary config directory");
std::fs::write(
config_dir.path().join("config.toml"),
"[client]\nhttp_url = \"http://127.0.0.1:9711\"\n",
)
.expect("write config.toml");
let config_home = config_dir.path().to_str().expect("utf-8 temp path");
let environment =
Environment::from([("WYRD_CONFIG_HOME", config_home), ("HOME", config_home)]);
let client = client_from_options_in(environment, None, Some("wyrd_sk_test"), None)
.expect("client from the configured server");
assert_eq!(client.server_url(), "http://127.0.0.1:9711");
}
#[test]
fn blank_explicit_options_are_refused_by_their_public_name() {
let cases = [
("server_url", Some(""), Some("wyrd_sk_test"), None),
("server_url", Some(" "), Some("wyrd_sk_test"), None),
("grpc_url", None, Some("wyrd_sk_test"), Some(" ")),
("credential", None, Some(""), None),
];
for (field, server_url, credential, grpc_url) in cases {
let error = client_from_options_in(
Environment::from([("HOME", "/nonexistent")]),
server_url,
credential,
grpc_url,
)
.expect_err("a blank explicit option is refused");
let error = wyrd_spec::error::WyrdError::from(error);
assert_eq!(error.code(), "WYRD_CLIENT_400_CONFIG_INVALID", "{field}");
assert_eq!(
error.problem().details,
serde_json::json!({ "field": field, "reason": "must not be empty" }),
);
}
}
}