use std::fmt::Display;
use std::future::Future;
use std::hash::{BuildHasher as _, Hasher as _, RandomState};
use std::ops::Range;
use std::sync::Arc;
use std::time::Duration;
use arrow_array::RecordBatch;
use arrow_schema::{Schema, SchemaRef};
use async_trait::async_trait;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use reqwest::Url;
use reqwest_middleware::ClientWithMiddleware;
use spiral_contract::auth::AuthHeader;
use spiral_contract::call::CallError;
use spiral_contract::safelong::SafeUnsignedLong;
use spiral_contract::stream::ByteStream as ContractStream;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::api::scandal::ScandalClient;
use crate::auth::{Token, TokenSource, exchange_token, token_endpoint};
use crate::http::default_http_client;
use crate::ipc::{ArrowIpcReader, BodyError};
use crate::scandal::api::{
AbortBlobUploadRequest, AlignedFragment, AppendRowsRequest, ArrowSchema, AttachColumnsRequest,
AttachedColumns, BindProjectionRequest, ByteRange, ChangePlan, CompleteBlobUploadRequest,
DeleteRowsRequest, EvaluateRequest, FinalizeWriteRequest, PackBlobsRequest, PatchRowsRequest,
PlanChangesRequest, PlanScanRequest, ReadBlobRequest, ReadLegacyTableRequest, ReadScanRequest,
ScandalApi, SearchVectorsRequest, StartBlobUploadRequest, StatBlobRequest, TakeRowsRequest,
UploadBlobPartRequest, WriteAlignedRequest, WriteBlobRequest,
};
use crate::scandal::blob::{Blob, BlobPart, BlobStat, BlobUpload};
use crate::scandal::error::ScandalError;
use crate::scandal::expression::NamedExpression;
use crate::scandal::functions::FunctionInfo;
use crate::scandal::types::{
FragmentReceipt, PreparedWriteReceipt, Scan, ScanPlan, Split, SplitTake,
};
use crate::scandal::wire::{self, CAPABILITY_AUDIENCE};
use crate::scandal::{
ByteStream, Evaluation, RecordBatchStream, Scandal, SplitRead, StorageBinding, TableBinding,
WriteBase,
};
const RENEW_CAPABILITY_BEFORE: Duration = Duration::from_secs(60);
const UNAVAILABLE_RETRIES: u32 = 6;
const UNAVAILABLE_WAIT: Duration = Duration::from_millis(250);
#[derive(Clone)]
pub struct RemoteScandal {
client: ScandalClient,
exchange: ClientWithMiddleware,
capabilities: Option<Arc<Capabilities>>,
}
impl RemoteScandal {
pub fn new(base_url: Url) -> Result<Self, ScandalError> {
let http = default_http_client().map_err(local_error)?;
let client = ScandalClient::new(base_url.as_str()).map_err(call_error)?;
Ok(Self {
client,
exchange: http.clone(),
capabilities: None,
}
.with_http(http))
}
pub fn with_http(mut self, http: impl Into<ClientWithMiddleware>) -> Self {
let http = http.into();
self.exchange = http.clone();
self.client = self.client.with_http(http);
self
}
pub fn with_caller(
mut self,
caller: TokenSource,
auth_url: &Url,
api_url: &Url,
) -> Result<Self, ScandalError> {
self.capabilities = Some(Arc::new(Capabilities {
caller,
token_endpoint: token_endpoint(auth_url).map_err(local_error)?,
api: api_url.as_str().trim_end_matches('/').to_owned(),
current: Mutex::new(None),
}));
Ok(self)
}
async fn auth(&self) -> Result<AuthHeader, ScandalError> {
match &self.capabilities {
Some(capabilities) => {
let capability = capabilities.current(&self.exchange).await?;
Ok(AuthHeader::new(format!(
"Bearer {}",
capability.expose_secret()
)))
}
None => Ok(AuthHeader::new("")),
}
}
async fn admitted<T, Request>(
&self,
mut request: impl FnMut() -> Request,
) -> Result<T, ScandalError>
where
Request: Future<Output = Result<T, ScandalError>>,
{
let mut bound = UNAVAILABLE_WAIT;
for _ in 0..UNAVAILABLE_RETRIES {
match request().await {
Err(ScandalError::Unavailable { .. }) => {
tokio::time::sleep(bound.mul_f64(jitter())).await;
bound *= 2;
}
result => return result,
}
}
request().await
}
}
#[async_trait]
impl Scandal for RemoteScandal {
async fn read_legacy_table(
&self,
request: &ReadLegacyTableRequest,
) -> Result<ByteStream, ScandalError> {
let body = self
.admitted(|| async {
self.client
.read_legacy_table(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await?;
Ok(body.map_err(|error| local_error(&*error)).boxed())
}
async fn list_functions(&self) -> Result<Vec<FunctionInfo>, ScandalError> {
let catalog = self
.admitted(|| async {
self.client
.list_functions(self.auth().await?)
.await
.map_err(call_error)
})
.await?;
Ok(catalog.functions)
}
async fn plan_scan(
&self,
scan: &Scan,
storage: &StorageBinding,
) -> Result<ScanPlan, ScandalError> {
let request = PlanScanRequest {
scan: scan.clone(),
storage: Some(storage.clone()),
};
self.admitted(|| async {
self.client
.plan_scan(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn read_scan(
&self,
scan: &Scan,
split: &Split,
storage: &StorageBinding,
) -> Result<SplitRead, ScandalError> {
let request = ReadScanRequest {
scan: scan.clone(),
split: split.clone(),
storage: Some(storage.clone()),
};
let body = self
.admitted(|| async {
self.client
.read_scan(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await?;
split_read(body).await
}
async fn take_rows(
&self,
scan: &Scan,
takes: &[SplitTake],
storage: &StorageBinding,
) -> Result<SplitRead, ScandalError> {
let request = TakeRowsRequest {
scan: scan.clone(),
takes: takes.to_vec(),
storage: Some(storage.clone()),
};
let body = self
.admitted(|| async {
self.client
.take_rows(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await?;
split_read(body).await
}
async fn read_blob(
&self,
blob: &Blob,
range: Option<Range<u64>>,
storage: &StorageBinding,
) -> Result<ByteStream, ScandalError> {
let request = ReadBlobRequest {
blob: blob.clone(),
range: range.map(|range| ByteRange {
start: range.start.into(),
end: range.end.into(),
}),
storage: Some(storage.clone()),
};
let body = self
.admitted(|| async {
self.client
.read_blob(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await?;
Ok(body.map_err(|error| local_error(&*error)).boxed())
}
async fn stat_blob(
&self,
blob: &Blob,
storage: &StorageBinding,
) -> Result<BlobStat, ScandalError> {
let request = StatBlobRequest {
blob: blob.clone(),
storage: Some(storage.clone()),
};
self.admitted(|| async {
self.client
.stat_blob(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn write_blob(
&self,
table: &TableBinding,
media_type: &str,
max_bytes: u64,
input: ByteStream,
) -> Result<Blob, ScandalError> {
let request = WriteBlobRequest {
table: table.clone(),
media_type: media_type.to_owned(),
max_bytes: SafeUnsignedLong(max_bytes),
};
self.client
.write_blob(self.auth().await?, request, ContractStream::new(input))
.await
.map_err(call_error)
}
async fn pack_blobs(
&self,
table: &TableBinding,
attempt: Option<Uuid>,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<Vec<RecordBatch>, ScandalError> {
let request = PackBlobsRequest {
table: table.clone(),
attempt,
};
let rows = ContractStream::new(wire::encode_input(schema, input));
let body = self
.client
.pack_blobs(self.auth().await?, request, rows)
.await
.map_err(call_error)?;
let (_, batches) = arrow(body).await?;
batches.try_collect().await
}
async fn start_blob_upload(
&self,
table: &TableBinding,
attempt: Uuid,
media_type: &str,
part_bytes: Option<u64>,
) -> Result<BlobUpload, ScandalError> {
let request = StartBlobUploadRequest {
table: table.clone(),
attempt,
media_type: media_type.to_owned(),
part_bytes: part_bytes.map(SafeUnsignedLong),
};
self.admitted(|| async {
self.client
.start_blob_upload(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn upload_blob_part(
&self,
upload: &BlobUpload,
part: u32,
input: ByteStream,
storage: &StorageBinding,
) -> Result<BlobPart, ScandalError> {
let request = UploadBlobPartRequest {
upload: upload.clone(),
storage: Some(storage.clone()),
part,
};
let chunks = input.try_collect::<Vec<_>>().await?;
self.admitted(|| async {
let body = stream::iter(chunks.clone().into_iter().map(Ok::<_, ScandalError>));
self.client
.upload_blob_part(
self.auth().await?,
request.clone(),
ContractStream::new(body),
)
.await
.map_err(call_error)
})
.await
}
async fn complete_blob_upload(
&self,
upload: &BlobUpload,
parts: &[BlobPart],
storage: &StorageBinding,
) -> Result<Blob, ScandalError> {
let request = CompleteBlobUploadRequest {
upload: upload.clone(),
storage: Some(storage.clone()),
parts: parts.to_vec(),
};
self.admitted(|| async {
self.client
.complete_blob_upload(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn abort_blob_upload(
&self,
upload: &BlobUpload,
storage: &StorageBinding,
) -> Result<(), ScandalError> {
let request = AbortBlobUploadRequest {
upload: upload.clone(),
storage: Some(storage.clone()),
};
self.admitted(|| async {
self.client
.abort_blob_upload(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn append_rows(
&self,
table: &TableBinding,
base: Option<Uuid>,
attempt: Option<Uuid>,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<Vec<FragmentReceipt>, ScandalError> {
self.append_partition(
&AppendRowsRequest {
table: table.clone(),
base,
attempt,
properties: None,
schema_changes: None,
},
schema,
input,
)
.await
}
async fn append_partition(
&self,
request: &AppendRowsRequest,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<Vec<FragmentReceipt>, ScandalError> {
let rows = ContractStream::new(wire::encode_input(schema, input));
let appended = self
.client
.append_rows(self.auth().await?, request.clone(), rows)
.await
.map_err(call_error)?;
Ok(appended.receipts)
}
async fn patch_rows(
&self,
request: &PatchRowsRequest,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<Vec<FragmentReceipt>, ScandalError> {
let rows = ContractStream::new(wire::encode_input(schema, input));
let written = self
.client
.patch_rows(self.auth().await?, request.clone(), rows)
.await
.map_err(call_error)?;
Ok(written.receipts)
}
async fn delete_rows(
&self,
request: &DeleteRowsRequest,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<Vec<FragmentReceipt>, ScandalError> {
let rows = ContractStream::new(wire::encode_input(schema, input));
let written = self
.client
.delete_rows(self.auth().await?, request.clone(), rows)
.await
.map_err(call_error)?;
Ok(written.receipts)
}
async fn attach_columns(
&self,
request: &AttachColumnsRequest,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<AttachedColumns, ScandalError> {
let rows = ContractStream::new(wire::encode_input(schema, input));
self.client
.attach_columns(self.auth().await?, request.clone(), rows)
.await
.map_err(call_error)
}
async fn write_aligned(
&self,
request: &WriteAlignedRequest,
schema: SchemaRef,
input: RecordBatchStream,
) -> Result<AlignedFragment, ScandalError> {
let rows = ContractStream::new(wire::encode_input(schema, input));
self.client
.write_aligned(self.auth().await?, request.clone(), rows)
.await
.map_err(call_error)
}
async fn plan_changes(&self, request: &PlanChangesRequest) -> Result<ChangePlan, ScandalError> {
self.admitted(|| async {
self.client
.plan_changes(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn finalize_write(
&self,
table: &TableBinding,
base: &WriteBase,
receipts: &[FragmentReceipt],
) -> Result<PreparedWriteReceipt, ScandalError> {
let request = FinalizeWriteRequest {
table: table.clone(),
base: base.clone(),
receipts: receipts.to_vec(),
};
self.admitted(|| async {
self.client
.finalize_write(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await
}
async fn bind_projection(
&self,
input: &Schema,
projection: &[NamedExpression],
) -> Result<SchemaRef, ScandalError> {
let request = BindProjectionRequest {
input: ArrowSchema::from_schema(input)?,
projection: projection.to_vec(),
};
let schema = self
.admitted(|| async {
self.client
.bind_projection(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await?;
Ok(SchemaRef::new(schema.to_schema()?))
}
async fn search_vectors(
&self,
request: &SearchVectorsRequest,
) -> Result<Evaluation, ScandalError> {
request.validate()?;
let body = self
.admitted(|| async {
self.client
.search_vectors(self.auth().await?, request.clone())
.await
.map_err(call_error)
})
.await?;
let (schema, batches) = arrow(body).await?;
Ok(Evaluation { schema, batches })
}
async fn evaluate(
&self,
projection: &[NamedExpression],
input_schema: SchemaRef,
input: RecordBatchStream,
tables: &[TableBinding],
) -> Result<Evaluation, ScandalError> {
let request = EvaluateRequest {
projection: projection.to_vec(),
tables: tables.to_vec(),
};
let rows = ContractStream::new(wire::encode_input(input_schema, input));
let body = self
.client
.evaluate(self.auth().await?, request, rows)
.await
.map_err(call_error)?;
let (schema, batches) = arrow(body).await?;
Ok(Evaluation { schema, batches })
}
}
async fn split_read(body: ContractStream) -> Result<SplitRead, ScandalError> {
let (schema, batches) = arrow(body).await?;
Ok(SplitRead {
schema: wire::read_scan_rows(&schema)?,
batches: batches
.map(|batch| batch.and_then(wire::detach_progress))
.boxed(),
})
}
async fn arrow(body: ContractStream) -> Result<(SchemaRef, RecordBatchStream), ScandalError> {
let source = body.map_err(|error| error as BodyError).boxed();
let reader = ArrowIpcReader::new(source).await.map_err(local_error)?;
let schema = wire::execution_rows(&reader.schema())?;
let batches = stream::unfold(reader, |mut reader| async move {
reader.next_batch().await.map(|batch| {
(
batch.map_err(local_error).and_then(wire::detach_execution),
reader,
)
})
});
Ok((schema, wire::execution_batches(batches.boxed())))
}
struct Capabilities {
caller: TokenSource,
token_endpoint: Url,
api: String,
current: Mutex<Option<Narrowed>>,
}
struct Narrowed {
parent: Token,
capability: Token,
}
impl Capabilities {
async fn current(&self, http: &ClientWithMiddleware) -> Result<Token, ScandalError> {
let unauthorized = |reason: String| ScandalError::Unauthenticated {
message: format!("no scandal capability: {reason}"),
};
let token = self
.caller
.token()
.await
.map_err(|error| unauthorized(error.to_string()))?
.ok_or_else(|| unauthorized("the caller has no credentials".to_owned()))?;
let mut current = self.current.lock().await;
if let Some(narrowed) = current.as_ref()
&& narrowed.parent.expose_secret() == token.expose_secret()
&& narrowed
.capability
.remaining_lifetime()
.is_ok_and(|remaining| remaining > RENEW_CAPABILITY_BEFORE)
{
return Ok(narrowed.capability.clone());
}
let capability = exchange_token(
http,
&self.token_endpoint,
&token,
&[CAPABILITY_AUDIENCE, &self.api],
)
.await
.map_err(|error| unauthorized(error.to_string()))?;
*current = Some(Narrowed {
parent: token,
capability: capability.clone(),
});
Ok(capability)
}
}
fn jitter() -> f64 {
let random = RandomState::new().build_hasher().finish();
(random >> 11) as f64 / (1u64 << 53) as f64
}
fn local_error(error: impl Display) -> ScandalError {
ScandalError::Internal {
message: error.to_string(),
}
}
fn call_error(error: CallError<ScandalError>) -> ScandalError {
match error {
CallError::Service(failure) => match failure.into_declared() {
Ok(error) => error,
Err(failure) => {
let envelope = failure.envelope();
let message = envelope.to_string();
match envelope.error_code.as_str() {
"UNAUTHENTICATED" => ScandalError::Unauthenticated { message },
"INVALID_ARGUMENT" => ScandalError::InvalidArgument { message },
_ => ScandalError::from_wire(envelope),
}
}
},
error => local_error(error),
}
}
#[cfg(test)]
pub(crate) mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde_json::json;
use wiremock::matchers::{body_string_contains, header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::*;
use crate::auth::TokenSource;
pub(crate) fn unsigned_jwt(sub: &str, ttl: Duration) -> String {
let exp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap() + ttl;
let encode = |value: serde_json::Value| URL_SAFE_NO_PAD.encode(value.to_string());
format!(
"{}.{}.signature",
encode(json!({ "alg": "RS256", "typ": "JWT" })),
encode(json!({ "sub": sub, "exp": exp.as_secs() })),
)
}
pub(crate) async fn issue_capability(
server: &MockServer,
caller: &str,
capability: &str,
exchanges: u64,
requests: u64,
) {
Mock::given(method("POST"))
.and(path("/oauth/token"))
.and(body_string_contains(format!("subject_token={caller}&")))
.and(body_string_contains(
"resource=https%3A%2F%2Fscandal.spiraldb.dev&resource=",
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token": capability,
"token_type": "Bearer",
"expires_in": 900,
})))
.expect(exchanges)
.mount(server)
.await;
Mock::given(method("GET"))
.and(path("/functions"))
.and(header("authorization", format!("Bearer {capability}")))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({ "functions": [] })))
.expect(requests)
.mount(server)
.await;
}
async fn scandal_for(server: &MockServer, caller: &str) -> RemoteScandal {
let url: Url = server.uri().parse().unwrap();
RemoteScandal::new(url.clone())
.unwrap()
.with_caller(TokenSource::fixed(caller.into()), &url, &url)
.unwrap()
}
#[tokio::test]
async fn requests_carry_a_capability_exchanged_from_the_callers_token() {
let server = MockServer::start().await;
let capability = unsigned_jwt("user_test", Duration::from_secs(900));
issue_capability(&server, "stack-token", &capability, 1, 2).await;
let scandal = scandal_for(&server, "stack-token").await;
for _ in 0..2 {
scandal.list_functions().await.unwrap();
}
}
#[tokio::test]
async fn renews_a_capability_before_it_expires() {
let server = MockServer::start().await;
let capability = unsigned_jwt("user_test", RENEW_CAPABILITY_BEFORE / 2);
issue_capability(&server, "stack-token", &capability, 2, 2).await;
let scandal = scandal_for(&server, "stack-token").await;
for _ in 0..2 {
scandal.list_functions().await.unwrap();
}
}
#[tokio::test]
async fn narrows_the_callers_new_token_when_it_changes() {
let server = MockServer::start().await;
let first = unsigned_jwt("user_test", Duration::from_secs(900));
let second = unsigned_jwt("user_other", Duration::from_secs(900));
issue_capability(&server, "first-token", &first, 1, 1).await;
issue_capability(&server, "second-token", &second, 1, 1).await;
let calls = Arc::new(AtomicUsize::new(0));
let caller = TokenSource::new(move || {
let call = calls.fetch_add(1, Ordering::SeqCst);
async move {
Ok(Some(Token::from(if call == 0 {
"first-token"
} else {
"second-token"
})))
}
});
let url: Url = server.uri().parse().unwrap();
let scandal = RemoteScandal::new(url.clone())
.unwrap()
.with_caller(caller, &url, &url)
.unwrap();
for _ in 0..2 {
scandal.list_functions().await.unwrap();
}
}
#[tokio::test]
async fn a_refused_exchange_fails_before_scandal_is_called() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(400).set_body_json(json!({
"error": "invalid_request",
"error_description": "the subject token is not for https://scandal.spiraldb.dev",
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/functions"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({ "functions": [] })))
.expect(0)
.mount(&server)
.await;
let error = scandal_for(&server, "spfs-capability")
.await
.list_functions()
.await
.unwrap_err();
assert!(
matches!(error, ScandalError::Unauthenticated { .. }),
"{error:?}"
);
}
}