headwaters_client/
client.rs

1//! The [`HeadwatersClient`] and its builder.
2
3use connectrpc::client::{ClientConfig, ClientTransport, HttpClient};
4use headwaters_proto::connect_gen::headwaters::read::v1::ReadServiceClient;
5
6use crate::types::*;
7use crate::{Error, HeadwatersConfig};
8
9/// An async client for the Headwaters read API.
10///
11/// Construct one from a base URL with [`connect`](Self::connect), from the
12/// environment with [`from_env`](Self::from_env), or from a prepared
13/// [`HeadwatersConfig`] with [`new`](Self::new). Each method maps to one read
14/// RPC and returns an owned response message; the open-ended `facets` / `data`
15/// fields on those messages are `google.protobuf.Struct`, which you can turn
16/// into a [`serde_json::Value`] with [`struct_to_json`].
17///
18/// # Examples
19///
20/// ```no_run
21/// # async fn run() -> Result<(), headwaters_client::Error> {
22/// use headwaters_client::HeadwatersClient;
23///
24/// let client = HeadwatersClient::connect("http://localhost:8091")?;
25/// let namespaces = client.list_namespaces().await?;
26/// for ns in &namespaces.namespaces {
27///     println!("{}", ns.name);
28/// }
29/// # Ok(())
30/// # }
31/// ```
32///
33/// The transport `T` defaults to [`HttpClient`] (plaintext/TLS over hyper),
34/// which is what [`connect`](Self::connect), [`new`](Self::new), and
35/// [`from_env`](Self::from_env) build. With the `cloud-auth` feature, the
36/// cloud constructors return a client over olai-http's `CloudTransport`, which
37/// signs each request for a cloud provider (AWS/Azure/GCP/Databricks).
38#[derive(Clone)]
39pub struct HeadwatersClient<T = HttpClient>
40where
41    T: ClientTransport,
42{
43    inner: ReadServiceClient<T>,
44}
45
46impl HeadwatersClient<HttpClient> {
47    /// Connect to a server at `base_url` (e.g. `http://localhost:8091`) with
48    /// default settings. Shorthand for `new(HeadwatersConfig::new(base_url))`.
49    pub fn connect(base_url: impl Into<String>) -> Result<Self, Error> {
50        Self::new(HeadwatersConfig::new(base_url))
51    }
52
53    /// Build a client from a [`HeadwatersConfig`].
54    ///
55    /// Only `http://` base URLs are supported today; an `https://` URL returns
56    /// [`Error::InvalidBaseUrl`] (TLS support is a planned follow-up).
57    pub fn new(config: HeadwatersConfig) -> Result<Self, Error> {
58        let uri: http::Uri = config
59            .base_url
60            .parse()
61            .map_err(|_| Error::InvalidBaseUrl(config.base_url.clone()))?;
62        if uri.scheme_str() == Some("https") {
63            return Err(Error::InvalidBaseUrl(format!(
64                "{}: https is not supported yet (TLS is a planned follow-up)",
65                config.base_url
66            )));
67        }
68
69        let mut client_config = ClientConfig::new(uri).with_default_timeout(config.timeout);
70        if let Some(token) = &config.token {
71            client_config =
72                client_config.with_default_header("authorization", format!("Bearer {token}"));
73        }
74
75        Ok(Self {
76            inner: ReadServiceClient::new(HttpClient::plaintext(), client_config),
77        })
78    }
79
80    /// Build a client from the environment (`HEADWATERS_URL`, `HEADWATERS_TOKEN`,
81    /// `HEADWATERS_TIMEOUT_MS`). See [`HeadwatersConfig::from_env`].
82    pub fn from_env() -> Result<Self, Error> {
83        Self::new(HeadwatersConfig::from_env()?)
84    }
85}
86
87impl<T> HeadwatersClient<T>
88where
89    T: ClientTransport,
90    <T::ResponseBody as http_body::Body>::Error: std::fmt::Display,
91{
92    // --- namespaces ---
93
94    /// List every namespace.
95    pub async fn list_namespaces(&self) -> Result<ListNamespacesResponse, Error> {
96        Ok(self
97            .inner
98            .list_namespaces(ListNamespacesRequest::default())
99            .await?
100            .into_owned())
101    }
102
103    // --- jobs ---
104
105    /// List jobs, optionally scoped to one namespace (`""` = all).
106    pub async fn list_jobs(
107        &self,
108        namespace: &str,
109        limit: i32,
110        offset: i32,
111    ) -> Result<ListJobsResponse, Error> {
112        Ok(self
113            .inner
114            .list_jobs(ListJobsRequest {
115                namespace: namespace.to_string(),
116                limit,
117                offset,
118                ..Default::default()
119            })
120            .await?
121            .into_owned())
122    }
123
124    /// Get one job by namespace and name.
125    pub async fn get_job(&self, namespace: &str, name: &str) -> Result<JobDetail, Error> {
126        Ok(self
127            .inner
128            .get_job(GetJobRequest {
129                namespace: namespace.to_string(),
130                name: name.to_string(),
131                ..Default::default()
132            })
133            .await?
134            .into_owned())
135    }
136
137    /// List the runs of one job.
138    pub async fn get_job_runs(
139        &self,
140        namespace: &str,
141        name: &str,
142    ) -> Result<ListRunsResponse, Error> {
143        Ok(self
144            .inner
145            .get_job_runs(GetJobRunsRequest {
146                namespace: namespace.to_string(),
147                name: name.to_string(),
148                ..Default::default()
149            })
150            .await?
151            .into_owned())
152    }
153
154    // --- datasets ---
155
156    /// List datasets, optionally scoped to one namespace (`""` = all).
157    pub async fn list_datasets(
158        &self,
159        namespace: &str,
160        limit: i32,
161        offset: i32,
162    ) -> Result<ListDatasetsResponse, Error> {
163        Ok(self
164            .inner
165            .list_datasets(ListDatasetsRequest {
166                namespace: namespace.to_string(),
167                limit,
168                offset,
169                ..Default::default()
170            })
171            .await?
172            .into_owned())
173    }
174
175    /// Get one dataset by namespace and name.
176    pub async fn get_dataset(&self, namespace: &str, name: &str) -> Result<Dataset, Error> {
177        Ok(self
178            .inner
179            .get_dataset(GetDatasetRequest {
180                namespace: namespace.to_string(),
181                name: name.to_string(),
182                ..Default::default()
183            })
184            .await?
185            .into_owned())
186    }
187
188    /// List the schema versions of one dataset, newest first.
189    pub async fn list_dataset_versions(
190        &self,
191        namespace: &str,
192        name: &str,
193        limit: i32,
194        offset: i32,
195    ) -> Result<ListDatasetVersionsResponse, Error> {
196        Ok(self
197            .inner
198            .list_dataset_versions(ListDatasetVersionsRequest {
199                namespace: namespace.to_string(),
200                name: name.to_string(),
201                limit,
202                offset,
203                ..Default::default()
204            })
205            .await?
206            .into_owned())
207    }
208
209    // --- search / graph / events / facets ---
210
211    /// Substring search over job and dataset names. `kind` restricts to one
212    /// entity kind (`ENTITY_KIND_UNSPECIFIED` = both); `namespace` restricts to
213    /// one namespace (`""` = all).
214    pub async fn search(
215        &self,
216        q: &str,
217        limit: i32,
218        kind: EntityKind,
219        namespace: &str,
220    ) -> Result<SearchResponse, Error> {
221        Ok(self
222            .inner
223            .search(SearchRequest {
224                q: q.to_string(),
225                limit,
226                r#type: kind.into(),
227                namespace: namespace.to_string(),
228                ..Default::default()
229            })
230            .await?
231            .into_owned())
232    }
233
234    /// Entity-level lineage: the jobs and datasets within `depth` hops of the
235    /// seed `node_id`, with the edges between them.
236    pub async fn get_lineage(&self, node_id: &str, depth: i32) -> Result<LineageGraph, Error> {
237        Ok(self
238            .inner
239            .get_lineage(GetLineageRequest {
240                node_id: node_id.to_string(),
241                depth,
242                ..Default::default()
243            })
244            .await?
245            .into_owned())
246    }
247
248    /// Column-level lineage for the seed `node_id` (a `dataset:` or
249    /// `datasetField:` nodeId).
250    pub async fn get_column_lineage(&self, node_id: &str) -> Result<LineageGraph, Error> {
251        Ok(self
252            .inner
253            .get_column_lineage(GetColumnLineageRequest {
254                node_id: node_id.to_string(),
255                ..Default::default()
256            })
257            .await?
258            .into_owned())
259    }
260
261    /// The raw OpenLineage events feed, newest first.
262    pub async fn list_events(&self, limit: i32, offset: i32) -> Result<ListEventsResponse, Error> {
263        Ok(self
264            .inner
265            .list_events(ListEventsRequest {
266                limit,
267                offset,
268                ..Default::default()
269            })
270            .await?
271            .into_owned())
272    }
273
274    /// The merged facets observed across all of a run's events.
275    pub async fn get_run_facets(&self, run_id: &str) -> Result<RunFacetsResponse, Error> {
276        Ok(self
277            .inner
278            .get_run_facets(GetRunFacetsRequest {
279                run_id: run_id.to_string(),
280                ..Default::default()
281            })
282            .await?
283            .into_owned())
284    }
285
286    // --- tags ---
287
288    /// List every tag in the catalog.
289    pub async fn list_tags(&self) -> Result<ListTagsResponse, Error> {
290        Ok(self
291            .inner
292            .list_tags(ListTagsRequest::default())
293            .await?
294            .into_owned())
295    }
296
297    /// Every field a tag reaches downstream through column lineage.
298    pub async fn get_tag_downstream(&self, tag: &str) -> Result<TagPropagation, Error> {
299        Ok(self
300            .inner
301            .get_tag_downstream(GetTagDownstreamRequest {
302                tag: tag.to_string(),
303                ..Default::default()
304            })
305            .await?
306            .into_owned())
307    }
308
309    // --- stats ---
310
311    /// Time-bucketed counts of lineage events. `period` is `HOUR` | `DAY` |
312    /// `WEEK` | `MONTH`; `limit` caps the number of buckets.
313    pub async fn get_lineage_event_stats(
314        &self,
315        period: &str,
316        limit: i32,
317    ) -> Result<StatsResponse, Error> {
318        Ok(self
319            .inner
320            .get_lineage_event_stats(GetLineageEventStatsRequest {
321                period: period.to_string(),
322                limit,
323                ..Default::default()
324            })
325            .await?
326            .into_owned())
327    }
328
329    /// Time-bucketed first-seen counts of an asset kind (`jobs` | `datasets`).
330    pub async fn get_asset_stats(
331        &self,
332        asset: &str,
333        period: &str,
334        limit: i32,
335    ) -> Result<StatsResponse, Error> {
336        Ok(self
337            .inner
338            .get_asset_stats(GetAssetStatsRequest {
339                asset: asset.to_string(),
340                period: period.to_string(),
341                limit,
342                ..Default::default()
343            })
344            .await?
345            .into_owned())
346    }
347}
348
349#[cfg(feature = "cloud-auth")]
350mod cloud {
351    use connectrpc::client::ClientConfig;
352    use headwaters_proto::connect_gen::headwaters::read::v1::ReadServiceClient;
353    use olai_http::CloudClient;
354    use olai_http::connectrpc::CloudTransport;
355
356    use super::HeadwatersClient;
357    use crate::{Error, HeadwatersConfig};
358
359    /// Build a [`ClientConfig`] for a cloud-fronted server. Unlike the bearer
360    /// path, this sets no `authorization` default header — the [`CloudClient`]
361    /// signer owns request authentication.
362    fn cloud_config(config: &HeadwatersConfig) -> Result<ClientConfig, Error> {
363        let uri: http::Uri = config
364            .base_url
365            .parse()
366            .map_err(|_| Error::InvalidBaseUrl(config.base_url.clone()))?;
367        Ok(ClientConfig::new(uri).with_default_timeout(config.timeout))
368    }
369
370    impl HeadwatersClient<CloudTransport> {
371        /// Build a client that authenticates with a pre-built [`CloudClient`].
372        ///
373        /// Use this to bring your own credential configuration (e.g. an
374        /// [`CloudClient`] constructed with a specific runtime handle or a
375        /// custom credential provider). For the common cases, prefer the
376        /// per-provider constructors
377        /// ([`new_aws`](Self::new_aws), [`new_azure`](Self::new_azure),
378        /// [`new_google`](Self::new_google),
379        /// [`new_databricks`](Self::new_databricks)).
380        ///
381        /// `https://` base URLs are expected here — TLS is handled by the
382        /// underlying olai-http reqwest client.
383        pub fn new_with_cloud(
384            config: HeadwatersConfig,
385            client: CloudClient,
386        ) -> Result<Self, Error> {
387            let client_config = cloud_config(&config)?;
388            Ok(Self {
389                inner: ReadServiceClient::new(CloudTransport::new(client), client_config),
390            })
391        }
392
393        /// Authenticate with AWS credentials (SigV4). `options` are olai-http
394        /// AWS config keys (see [`olai_http::aws`]); pass an empty iterator to
395        /// rely on the ambient environment / instance metadata.
396        pub fn new_aws<I, K, V>(config: HeadwatersConfig, options: I) -> Result<Self, Error>
397        where
398            I: IntoIterator<Item = (K, V)>,
399            K: AsRef<str>,
400            V: Into<String>,
401        {
402            Self::new_with_cloud(config, CloudClient::new_aws(options, None)?)
403        }
404
405        /// Authenticate with Azure AD credentials. `options` are olai-http Azure
406        /// config keys (see [`olai_http::azure`]).
407        pub fn new_azure<I, K, V>(config: HeadwatersConfig, options: I) -> Result<Self, Error>
408        where
409            I: IntoIterator<Item = (K, V)>,
410            K: AsRef<str>,
411            V: Into<String>,
412        {
413            Self::new_with_cloud(config, CloudClient::new_azure(options, None)?)
414        }
415
416        /// Authenticate with Google Cloud credentials. `options` are olai-http
417        /// GCP config keys (see [`olai_http::gcp`]).
418        pub fn new_google<I, K, V>(config: HeadwatersConfig, options: I) -> Result<Self, Error>
419        where
420            I: IntoIterator<Item = (K, V)>,
421            K: AsRef<str>,
422            V: Into<String>,
423        {
424            Self::new_with_cloud(config, CloudClient::new_google(options, None)?)
425        }
426
427        /// Authenticate with Databricks credentials (OAuth M2M / OIDC / PAT).
428        /// `options` are olai-http Databricks config keys (see
429        /// [`olai_http::databricks`]).
430        pub fn new_databricks<I, K, V>(config: HeadwatersConfig, options: I) -> Result<Self, Error>
431        where
432            I: IntoIterator<Item = (K, V)>,
433            K: AsRef<str>,
434            V: Into<String>,
435        {
436            Self::new_with_cloud(config, CloudClient::new_databricks(options, None)?)
437        }
438
439        /// Build a cloud-authenticated client from the environment.
440        ///
441        /// Reads the base URL and timeout from the usual variables
442        /// ([`HeadwatersConfig::from_env`]) and selects the provider from
443        /// [`HEADWATERS_AUTH`](crate::ENV_AUTH) (`aws` | `azure` | `gcp` |
444        /// `databricks`). Each provider discovers its own credentials from the
445        /// ambient environment, so no per-provider options are passed here.
446        ///
447        /// Returns [`Error::InvalidBaseUrl`] if `HEADWATERS_AUTH` is unset or
448        /// names an unknown provider.
449        pub fn from_env() -> Result<Self, Error> {
450            let config = HeadwatersConfig::from_env()?;
451            let provider = std::env::var(crate::ENV_AUTH).unwrap_or_default();
452            let empty: [(&str, &str); 0] = [];
453            match provider.to_ascii_lowercase().as_str() {
454                "aws" => Self::new_aws(config, empty),
455                "azure" => Self::new_azure(config, empty),
456                "gcp" | "google" => Self::new_google(config, empty),
457                "databricks" => Self::new_databricks(config, empty),
458                other => Err(Error::InvalidBaseUrl(format!(
459                    "unknown {}={other:?} (expected aws|azure|gcp|databricks)",
460                    crate::ENV_AUTH
461                ))),
462            }
463        }
464    }
465}
headwaters_client/client.rs

headwaters_client/
client.rs
