micromegas/client/
flightsql_client_factory.rs

1use super::flightsql_client::Client;
2use anyhow::{Context, Result};
3use async_trait::async_trait;
4use http::Uri;
5use tonic::transport::{Channel, ClientTlsConfig};
6
7/// A trait for creating FlightSQL clients.
8#[async_trait]
9pub trait FlightSQLClientFactory: Send + Sync {
10    async fn make_client(&self) -> Result<Client>;
11}
12
13/// A FlightSQL client factory that uses a bearer token for authentication.
14pub struct BearerFlightSQLClientFactory {
15    url: String,
16    token: String,
17    client_type: Option<String>,
18    extra_metadata: Vec<(String, String)>,
19}
20
21impl BearerFlightSQLClientFactory {
22    /// Creates a new `BearerFlightSQLClientFactory`.
23    ///
24    /// # Arguments
25    ///
26    /// * `url` - The FlightSQL server URL.
27    /// * `token` - The bearer token to use for authentication.
28    pub fn new(url: String, token: String) -> Self {
29        Self {
30            url,
31            token,
32            client_type: None,
33            extra_metadata: Vec::new(),
34        }
35    }
36
37    /// Creates a new `BearerFlightSQLClientFactory` with a specific client type identifier.
38    ///
39    /// # Arguments
40    ///
41    /// * `url` - The FlightSQL server URL.
42    /// * `token` - The bearer token to use for authentication.
43    /// * `client_type` - The client type identifier (e.g., "web", "cli", "python").
44    pub fn new_with_client_type(url: String, token: String, client_type: String) -> Self {
45        Self {
46            url,
47            token,
48            client_type: Some(client_type),
49            extra_metadata: Vec::new(),
50        }
51    }
52
53    /// Creates a new `BearerFlightSQLClientFactory` that reads the URL from the
54    /// `MICROMEGAS_FLIGHTSQL_URL` environment variable.
55    pub fn from_env(token: String) -> Result<Self> {
56        let url = std::env::var("MICROMEGAS_FLIGHTSQL_URL")
57            .with_context(|| "error reading MICROMEGAS_FLIGHTSQL_URL environment variable")?;
58        Ok(Self {
59            url,
60            token,
61            client_type: None,
62            extra_metadata: Vec::new(),
63        })
64    }
65
66    /// Creates a new `BearerFlightSQLClientFactory` that reads the URL from the
67    /// `MICROMEGAS_FLIGHTSQL_URL` environment variable, with a client type.
68    pub fn from_env_with_client_type(token: String, client_type: String) -> Result<Self> {
69        let url = std::env::var("MICROMEGAS_FLIGHTSQL_URL")
70            .with_context(|| "error reading MICROMEGAS_FLIGHTSQL_URL environment variable")?;
71        Ok(Self {
72            url,
73            token,
74            client_type: Some(client_type),
75            extra_metadata: Vec::new(),
76        })
77    }
78
79    /// Attaches an additional gRPC metadata header sent with every request made by
80    /// clients from this factory (e.g., the web app's notebook/cell origin labels).
81    /// `key` must already be a valid lowercase gRPC metadata key.
82    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
83        self.extra_metadata.push((key.into(), value.into()));
84        self
85    }
86}
87
88/// Rewrites the `grpc://`/`grpc+tls://` scheme convention used by data source configs into the
89/// `http://`/`https://` scheme tonic's `Channel` expects for its TLS decision. `http://`/`https://`
90/// URLs pass through unchanged.
91pub fn normalize_channel_scheme(url: &str) -> String {
92    let lower = url.to_ascii_lowercase();
93    if lower.starts_with("grpc+tls://") {
94        format!("https://{}", &url["grpc+tls://".len()..])
95    } else if lower.starts_with("grpc://") {
96        format!("http://{}", &url["grpc://".len()..])
97    } else {
98        url.to_string()
99    }
100}
101
102#[async_trait]
103impl FlightSQLClientFactory for BearerFlightSQLClientFactory {
104    async fn make_client(&self) -> Result<Client> {
105        let normalized_url = normalize_channel_scheme(&self.url);
106        let flight_url = normalized_url
107            .parse::<Uri>()
108            .with_context(|| "parsing flightsql url")?;
109        let mut endpoint = Channel::builder(flight_url.clone());
110        if flight_url.scheme_str() == Some("https") {
111            let tls_config = ClientTlsConfig::new().with_native_roots();
112            endpoint = endpoint
113                .tls_config(tls_config)
114                .with_context(|| "tls_config")?;
115        }
116        let channel = endpoint
117            .connect()
118            .await
119            .with_context(|| "connecting grpc channel")?;
120        let mut client = Client::new(channel);
121        let auth_value = if self.token.starts_with("Bearer ") {
122            self.token.clone()
123        } else {
124            format!("Bearer {}", self.token)
125        };
126
127        client
128            .inner_mut()
129            .set_header(http::header::AUTHORIZATION.as_str(), auth_value);
130
131        // Set client type header if provided
132        if let Some(client_type) = &self.client_type {
133            client
134                .inner_mut()
135                .set_header("x-client-type", client_type.clone());
136        }
137
138        // Set any additional per-factory metadata (e.g., notebook/cell origin labels)
139        for (key, value) in &self.extra_metadata {
140            client.inner_mut().set_header(key, value.clone());
141        }
142
143        // Preserve dictionary encoding for bandwidth efficiency
144        client.inner_mut().set_header("preserve_dictionary", "true");
145
146        Ok(client)
147    }
148}
micromegas/client/flightsql_client_factory.rs

micromegas/client/
flightsql_client_factory.rs
