Skip to main content

headwaters_client/
client.rs

1//! The [`HeadwatersClient`] and its builder.
2
3use connectrpc::client::{ClientConfig, 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#[derive(Clone)]
33pub struct HeadwatersClient {
34    inner: ReadServiceClient<HttpClient>,
35}
36
37impl HeadwatersClient {
38    /// Connect to a server at `base_url` (e.g. `http://localhost:8091`) with
39    /// default settings. Shorthand for `new(HeadwatersConfig::new(base_url))`.
40    pub fn connect(base_url: impl Into<String>) -> Result<Self, Error> {
41        Self::new(HeadwatersConfig::new(base_url))
42    }
43
44    /// Build a client from a [`HeadwatersConfig`].
45    ///
46    /// Only `http://` base URLs are supported today; an `https://` URL returns
47    /// [`Error::InvalidBaseUrl`] (TLS support is a planned follow-up).
48    pub fn new(config: HeadwatersConfig) -> Result<Self, Error> {
49        let uri: http::Uri = config
50            .base_url
51            .parse()
52            .map_err(|_| Error::InvalidBaseUrl(config.base_url.clone()))?;
53        if uri.scheme_str() == Some("https") {
54            return Err(Error::InvalidBaseUrl(format!(
55                "{}: https is not supported yet (TLS is a planned follow-up)",
56                config.base_url
57            )));
58        }
59
60        let mut client_config = ClientConfig::new(uri).with_default_timeout(config.timeout);
61        if let Some(token) = &config.token {
62            client_config =
63                client_config.with_default_header("authorization", format!("Bearer {token}"));
64        }
65
66        Ok(Self {
67            inner: ReadServiceClient::new(HttpClient::plaintext(), client_config),
68        })
69    }
70
71    /// Build a client from the environment (`HEADWATERS_URL`, `HEADWATERS_TOKEN`,
72    /// `HEADWATERS_TIMEOUT_MS`). See [`HeadwatersConfig::from_env`].
73    pub fn from_env() -> Result<Self, Error> {
74        Self::new(HeadwatersConfig::from_env()?)
75    }
76
77    // --- namespaces ---
78
79    /// List every namespace.
80    pub async fn list_namespaces(&self) -> Result<ListNamespacesResponse, Error> {
81        Ok(self
82            .inner
83            .list_namespaces(ListNamespacesRequest::default())
84            .await?
85            .into_owned())
86    }
87
88    // --- jobs ---
89
90    /// List jobs, optionally scoped to one namespace (`""` = all).
91    pub async fn list_jobs(
92        &self,
93        namespace: &str,
94        limit: i32,
95        offset: i32,
96    ) -> Result<ListJobsResponse, Error> {
97        Ok(self
98            .inner
99            .list_jobs(ListJobsRequest {
100                namespace: namespace.to_string(),
101                limit,
102                offset,
103                ..Default::default()
104            })
105            .await?
106            .into_owned())
107    }
108
109    /// Get one job by namespace and name.
110    pub async fn get_job(&self, namespace: &str, name: &str) -> Result<JobDetail, Error> {
111        Ok(self
112            .inner
113            .get_job(GetJobRequest {
114                namespace: namespace.to_string(),
115                name: name.to_string(),
116                ..Default::default()
117            })
118            .await?
119            .into_owned())
120    }
121
122    /// List the runs of one job.
123    pub async fn get_job_runs(
124        &self,
125        namespace: &str,
126        name: &str,
127    ) -> Result<ListRunsResponse, Error> {
128        Ok(self
129            .inner
130            .get_job_runs(GetJobRunsRequest {
131                namespace: namespace.to_string(),
132                name: name.to_string(),
133                ..Default::default()
134            })
135            .await?
136            .into_owned())
137    }
138
139    // --- datasets ---
140
141    /// List datasets, optionally scoped to one namespace (`""` = all).
142    pub async fn list_datasets(
143        &self,
144        namespace: &str,
145        limit: i32,
146        offset: i32,
147    ) -> Result<ListDatasetsResponse, Error> {
148        Ok(self
149            .inner
150            .list_datasets(ListDatasetsRequest {
151                namespace: namespace.to_string(),
152                limit,
153                offset,
154                ..Default::default()
155            })
156            .await?
157            .into_owned())
158    }
159
160    /// Get one dataset by namespace and name.
161    pub async fn get_dataset(&self, namespace: &str, name: &str) -> Result<Dataset, Error> {
162        Ok(self
163            .inner
164            .get_dataset(GetDatasetRequest {
165                namespace: namespace.to_string(),
166                name: name.to_string(),
167                ..Default::default()
168            })
169            .await?
170            .into_owned())
171    }
172
173    /// List the schema versions of one dataset, newest first.
174    pub async fn list_dataset_versions(
175        &self,
176        namespace: &str,
177        name: &str,
178        limit: i32,
179        offset: i32,
180    ) -> Result<ListDatasetVersionsResponse, Error> {
181        Ok(self
182            .inner
183            .list_dataset_versions(ListDatasetVersionsRequest {
184                namespace: namespace.to_string(),
185                name: name.to_string(),
186                limit,
187                offset,
188                ..Default::default()
189            })
190            .await?
191            .into_owned())
192    }
193
194    // --- search / graph / events / facets ---
195
196    /// Substring search over job and dataset names. `kind` restricts to one
197    /// entity kind (`ENTITY_KIND_UNSPECIFIED` = both); `namespace` restricts to
198    /// one namespace (`""` = all).
199    pub async fn search(
200        &self,
201        q: &str,
202        limit: i32,
203        kind: EntityKind,
204        namespace: &str,
205    ) -> Result<SearchResponse, Error> {
206        Ok(self
207            .inner
208            .search(SearchRequest {
209                q: q.to_string(),
210                limit,
211                r#type: kind.into(),
212                namespace: namespace.to_string(),
213                ..Default::default()
214            })
215            .await?
216            .into_owned())
217    }
218
219    /// Entity-level lineage: the jobs and datasets within `depth` hops of the
220    /// seed `node_id`, with the edges between them.
221    pub async fn get_lineage(&self, node_id: &str, depth: i32) -> Result<LineageGraph, Error> {
222        Ok(self
223            .inner
224            .get_lineage(GetLineageRequest {
225                node_id: node_id.to_string(),
226                depth,
227                ..Default::default()
228            })
229            .await?
230            .into_owned())
231    }
232
233    /// Column-level lineage for the seed `node_id` (a `dataset:` or
234    /// `datasetField:` nodeId).
235    pub async fn get_column_lineage(&self, node_id: &str) -> Result<LineageGraph, Error> {
236        Ok(self
237            .inner
238            .get_column_lineage(GetColumnLineageRequest {
239                node_id: node_id.to_string(),
240                ..Default::default()
241            })
242            .await?
243            .into_owned())
244    }
245
246    /// The raw OpenLineage events feed, newest first.
247    pub async fn list_events(&self, limit: i32, offset: i32) -> Result<ListEventsResponse, Error> {
248        Ok(self
249            .inner
250            .list_events(ListEventsRequest {
251                limit,
252                offset,
253                ..Default::default()
254            })
255            .await?
256            .into_owned())
257    }
258
259    /// The merged facets observed across all of a run's events.
260    pub async fn get_run_facets(&self, run_id: &str) -> Result<RunFacetsResponse, Error> {
261        Ok(self
262            .inner
263            .get_run_facets(GetRunFacetsRequest {
264                run_id: run_id.to_string(),
265                ..Default::default()
266            })
267            .await?
268            .into_owned())
269    }
270
271    // --- tags ---
272
273    /// List every tag in the catalog.
274    pub async fn list_tags(&self) -> Result<ListTagsResponse, Error> {
275        Ok(self
276            .inner
277            .list_tags(ListTagsRequest::default())
278            .await?
279            .into_owned())
280    }
281
282    /// Every field a tag reaches downstream through column lineage.
283    pub async fn get_tag_downstream(&self, tag: &str) -> Result<TagPropagation, Error> {
284        Ok(self
285            .inner
286            .get_tag_downstream(GetTagDownstreamRequest {
287                tag: tag.to_string(),
288                ..Default::default()
289            })
290            .await?
291            .into_owned())
292    }
293
294    // --- stats ---
295
296    /// Time-bucketed counts of lineage events. `period` is `HOUR` | `DAY` |
297    /// `WEEK` | `MONTH`; `limit` caps the number of buckets.
298    pub async fn get_lineage_event_stats(
299        &self,
300        period: &str,
301        limit: i32,
302    ) -> Result<StatsResponse, Error> {
303        Ok(self
304            .inner
305            .get_lineage_event_stats(GetLineageEventStatsRequest {
306                period: period.to_string(),
307                limit,
308                ..Default::default()
309            })
310            .await?
311            .into_owned())
312    }
313
314    /// Time-bucketed first-seen counts of an asset kind (`jobs` | `datasets`).
315    pub async fn get_asset_stats(
316        &self,
317        asset: &str,
318        period: &str,
319        limit: i32,
320    ) -> Result<StatsResponse, Error> {
321        Ok(self
322            .inner
323            .get_asset_stats(GetAssetStatsRequest {
324                asset: asset.to_string(),
325                period: period.to_string(),
326                limit,
327                ..Default::default()
328            })
329            .await?
330            .into_owned())
331    }
332}