1use connectrpc::client::{ClientConfig, HttpClient};
4use headwaters_proto::connect_gen::headwaters::read::v1::ReadServiceClient;
5
6use crate::types::*;
7use crate::{Error, HeadwatersConfig};
8
9#[derive(Clone)]
33pub struct HeadwatersClient {
34 inner: ReadServiceClient<HttpClient>,
35}
36
37impl HeadwatersClient {
38 pub fn connect(base_url: impl Into<String>) -> Result<Self, Error> {
41 Self::new(HeadwatersConfig::new(base_url))
42 }
43
44 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 pub fn from_env() -> Result<Self, Error> {
74 Self::new(HeadwatersConfig::from_env()?)
75 }
76
77 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}