Add metadata

Cargo.toml

@@ -43,6 +43,7 @@ chrono = "0.4.37"
 filesize = "0.2.0"
 serde = "1.0.203"
 serde_json = "1.0.117"
+secrecy = "=0.10.3"
 backoff = "0.4.0"
 gcp-bigquery-client = "0.20.0"
 rayon = "1.10.0"

lib/flight/Cargo.toml

@@ -16,4 +16,5 @@ futures.workspace = true
 bincode.workspace = true
 polars = {workspace = true, features = ["polars-io", "ipc_streaming"]}
 thiserror.workspace = true
-log = "0.4.22"
+secrecy.workspace = true
+log.workspace = true

lib/flight/src/client.rs

@@ -12,6 +12,7 @@ use representation::RDFNodeType;
 use std::collections::HashMap;
 use std::str::FromStr;
 use std::time::Duration;
+use secrecy::{ExposeSecret, SecretString};
 use thiserror::*;
 use tonic::transport::{Channel, Endpoint};
 use tonic::{Request, Status};
@@ -49,6 +50,7 @@ impl ChrontextFlightClient {
     pub async fn query(
         &mut self,
         query: &str,
+        metadata: &HashMap<String, SecretString>,
     ) -> Result<SolutionMappings, ChrontextFlightClientError> {
         let endpoint = Endpoint::from_str(&self.uri)
             .map_err(|e| ChrontextFlightClientError::ConnectError(e))?;
@@ -60,7 +62,10 @@ impl ChrontextFlightClient {
                 .map_err(|x| ChrontextFlightClientError::ConnectError(x))?
         };
         info!("Building request");
-        let request = Request::new(Ticket::new(serialize(query).unwrap()));
+        let mut request = Request::new(Ticket::new(serialize(query).unwrap()));
+        for (k,v) in metadata {
+            request.metadata_mut().insert(k, v.expose_secret().parse().unwrap());
+        }
         info!("Sending request");
         let mut flight_data = client
             .do_get(request)

py_chrontext/Cargo.toml

@@ -22,6 +22,7 @@ log.workspace = true
 env_logger.workspace = true
 spargebra.workspace = true
 oxrdfio.workspace = true
+secrecy.workspace = true
 
 [lib]
 name = "chrontext"

py_chrontext/chrontext/chrontext.pyi

@@ -401,10 +401,11 @@ class DataProduct:
 
 
 class FlightClient:
-    def __init__(self, uri:str):
+    def __init__(self, uri:str, metadata:Dict[str, str]=None):
         """
 
         :param uri: The URI of the Flight server (see engine.serve_flight())
+        :param metadata: gRPC metadata to add to each request
         """
 
     def query(self,

py_chrontext/src/lib.rs

@@ -52,6 +52,7 @@ use representation::solution_mapping::EagerSolutionMappings;
 use representation::BaseRDFNodeType;
 use std::collections::HashMap;
 use std::sync::Arc;
+use secrecy::SecretString;
 use templates::python::{a, py_triple, PyArgument, PyInstance, PyParameter, PyTemplate, PyXSD};
 use tokio::runtime::Builder;
 use virtualization::bigquery::VirtualizedBigQueryDatabase;
@@ -292,14 +293,21 @@ impl PyEngine {
 #[derive(Clone)]
 #[pyclass(name = "FlightClient")]
 pub struct PyFlightClient {
-    uri:String
+    uri:String,
+    metadata: HashMap<String, SecretString>
 }
 
 #[pymethods]
 impl PyFlightClient {
     #[new]
-    pub fn new(uri:String) -> PyResult<Self> {
-        Ok(Self {uri})
+    pub fn new(uri:String, metadata: Option<HashMap<String, String>>) -> PyResult<Self> {
+        let mut metadata_s = HashMap::new();
+        if let Some(metadata) = metadata {
+            for (k,v) in metadata {
+                metadata_s.insert(k, SecretString::from(v));
+            }
+        }
+        Ok(Self {uri, metadata:metadata_s})
     }
 
     pub fn query(
@@ -311,6 +319,7 @@ impl PyFlightClient {
     ) -> PyResult<PyObject> {
         let sparql = sparql.to_string();
         let res = py.allow_threads(move || {
+
             let sparql = sparql;
             let mut builder = Builder::new_multi_thread();
             builder.enable_all();
@@ -320,7 +329,7 @@ impl PyFlightClient {
             let sm = builder
                 .build()
                 .unwrap()
-                .block_on(client.query(&sparql))
+                .block_on(client.query(&sparql, &self.metadata))
                 .map_err(|x|PyChrontextError::FlightClientError(x))?;
             Ok(sm)
         });