1use {
4 crate::{
5 access_token::{AccessToken, Scope},
6 compression::{compress_best, decompress},
7 root_ca_certificate, CredentialType,
8 },
9 backoff::{future::retry, ExponentialBackoff},
10 log::*,
11 std::{
12 str::FromStr,
13 time::{Duration, Instant},
14 },
15 thiserror::Error,
16 tonic::{
17 codegen::InterceptedService, transport::ClientTlsConfig, Request,
18 Status,
19 },
20};
21
22mod google {
23 mod rpc {
24 include!(concat!(
25 env!("CARGO_MANIFEST_DIR"),
26 concat!("/proto/google.rpc.rs")
27 ));
28 }
29 pub mod bigtable {
30 pub mod v2 {
31 include!(concat!(
32 env!("CARGO_MANIFEST_DIR"),
33 concat!("/proto/google.bigtable.v2.rs")
34 ));
35 }
36 }
37}
38use google::bigtable::v2::*;
39
40pub type RowKey = String;
41pub type RowData = Vec<(CellName, CellValue)>;
42pub type RowDataSlice<'a> = &'a [(CellName, CellValue)];
43pub type CellName = String;
44pub type CellValue = Vec<u8>;
45pub enum CellData<B, P> {
46 Bincode(B),
47 Protobuf(P),
48}
49
50#[derive(Debug, Error)]
51pub enum Error {
52 #[error("AccessToken: {0}")]
53 AccessToken(String),
54
55 #[error("Certificate: {0}")]
56 Certificate(String),
57
58 #[error("I/O: {0}")]
59 Io(std::io::Error),
60
61 #[error("Transport: {0}")]
62 Transport(tonic::transport::Error),
63
64 #[error("Invalid URI {0}: {1}")]
65 InvalidUri(String, String),
66
67 #[error("Row not found")]
68 RowNotFound,
69
70 #[error("Row write failed")]
71 RowWriteFailed,
72
73 #[error("Row delete failed")]
74 RowDeleteFailed,
75
76 #[error("Object not found: {0}")]
77 ObjectNotFound(String),
78
79 #[error("Object is corrupt: {0}")]
80 ObjectCorrupt(String),
81
82 #[error("RPC: {0}")]
83 Rpc(tonic::Status),
84
85 #[error("Timeout")]
86 Timeout,
87}
88
89impl std::convert::From<std::io::Error> for Error {
90 fn from(err: std::io::Error) -> Self {
91 Self::Io(err)
92 }
93}
94
95impl std::convert::From<tonic::transport::Error> for Error {
96 fn from(err: tonic::transport::Error) -> Self {
97 Self::Transport(err)
98 }
99}
100
101impl std::convert::From<tonic::Status> for Error {
102 fn from(err: tonic::Status) -> Self {
103 Self::Rpc(err)
104 }
105}
106
107pub type Result<T> = std::result::Result<T, Error>;
108type InterceptedRequestResult = std::result::Result<Request<()>, Status>;
109
110#[derive(Clone)]
111pub struct BigTableConnection {
112 access_token: Option<AccessToken>,
113 channel: tonic::transport::Channel,
114 table_prefix: String,
115 app_profile_id: String,
116 timeout: Option<Duration>,
117}
118
119impl BigTableConnection {
120 pub async fn new(
132 instance_name: &str,
133 app_profile_id: &str,
134 read_only: bool,
135 timeout: Option<Duration>,
136 credential_type: CredentialType,
137 ) -> Result<Self> {
138 match std::env::var("BIGTABLE_EMULATOR_HOST") {
139 Ok(endpoint) => {
140 info!("Connecting to bigtable emulator at {}", endpoint);
141
142 Ok(Self {
143 access_token: None,
144 channel: tonic::transport::Channel::from_shared(format!("http://{}", endpoint))
145 .map_err(|err| Error::InvalidUri(endpoint, err.to_string()))?
146 .connect_lazy(),
147 table_prefix: format!("projects/emulator/instances/{}/tables/", instance_name),
148 app_profile_id: app_profile_id.to_string(),
149 timeout,
150 })
151 }
152
153 Err(_) => {
154 let access_token = AccessToken::new(
155 if read_only {
156 Scope::BigTableDataReadOnly
157 } else {
158 Scope::BigTableData
159 },
160 credential_type,
161 )
162 .await
163 .map_err(Error::AccessToken)?;
164
165 let table_prefix = format!(
166 "projects/{}/instances/{}/tables/",
167 access_token.project(),
168 instance_name
169 );
170
171 let endpoint = {
172 let endpoint =
173 tonic::transport::Channel::from_static("https://bigtable.googleapis.com")
174 .tls_config(
175 ClientTlsConfig::new()
176 .ca_certificate(
177 root_ca_certificate::load().map_err(Error::Certificate)?,
178 )
179 .domain_name("bigtable.googleapis.com"),
180 )?;
181
182 if let Some(timeout) = timeout {
183 endpoint.timeout(timeout)
184 } else {
185 endpoint
186 }
187 };
188
189 let mut http = hyper::client::HttpConnector::new();
190 http.enforce_http(false);
191 let channel = match std::env::var("BIGTABLE_PROXY") {
192 Ok(proxy_uri) => {
193 let proxy = hyper_proxy::Proxy::new(
194 hyper_proxy::Intercept::All,
195 proxy_uri
196 .parse::<http::Uri>()
197 .map_err(|err| Error::InvalidUri(proxy_uri, err.to_string()))?,
198 );
199 let mut proxy_connector =
200 hyper_proxy::ProxyConnector::from_proxy(http, proxy)?;
201 proxy_connector.set_tls(None);
203 endpoint.connect_with_connector_lazy(proxy_connector)
204 }
205 _ => endpoint.connect_with_connector_lazy(http),
206 };
207
208 Ok(Self {
209 access_token: Some(access_token),
210 channel,
211 table_prefix,
212 app_profile_id: app_profile_id.to_string(),
213 timeout,
214 })
215 }
216 }
217 }
218
219 pub fn client(&self) -> BigTable<impl FnMut(Request<()>) -> InterceptedRequestResult> {
224 let access_token = self.access_token.clone();
225 let client = bigtable_client::BigtableClient::with_interceptor(
226 self.channel.clone(),
227 move |mut req: Request<()>| {
228 if let Some(access_token) = &access_token {
229 match FromStr::from_str(&access_token.get()) {
230 Ok(authorization_header) => {
231 req.metadata_mut()
232 .insert("authorization", authorization_header);
233 }
234 Err(err) => {
235 warn!("Failed to set authorization header: {}", err);
236 }
237 }
238 }
239 Ok(req)
240 },
241 );
242 BigTable {
243 access_token: self.access_token.clone(),
244 client,
245 table_prefix: self.table_prefix.clone(),
246 app_profile_id: self.app_profile_id.clone(),
247 timeout: self.timeout,
248 }
249 }
250
251 pub async fn put_bincode_cells_with_retry<T>(
252 &self,
253 table: &str,
254 cells: &[(RowKey, T)],
255 overwrite: bool,
256 ) -> Result<usize>
257 where
258 T: serde::ser::Serialize,
259 {
260 retry(ExponentialBackoff::default(), || async {
261 let mut client = self.client();
262 Ok(client.put_bincode_cells(table, cells, overwrite).await?)
263 })
264 .await
265 }
266
267 pub async fn delete_rows_with_retry(&self, table: &str, row_keys: &[RowKey]) -> Result<()> {
268 retry(ExponentialBackoff::default(), || async {
269 let mut client = self.client();
270 Ok(client.delete_rows(table, row_keys).await?)
271 })
272 .await
273 }
274
275 pub async fn get_bincode_cells_with_retry<T>(
276 &self,
277 table: &str,
278 row_keys: &[RowKey],
279 ) -> Result<Vec<(RowKey, Result<T>)>>
280 where
281 T: serde::de::DeserializeOwned,
282 {
283 retry(ExponentialBackoff::default(), || async {
284 let mut client = self.client();
285 Ok(client.get_bincode_cells(table, row_keys).await?)
286 })
287 .await
288 }
289
290 pub async fn put_protobuf_cells_with_retry<T>(
291 &self,
292 table: &str,
293 cells: &[(RowKey, T)],
294 overwrite: bool,
295 ) -> Result<usize>
296 where
297 T: prost::Message,
298 {
299 retry(ExponentialBackoff::default(), || async {
300 let mut client = self.client();
301 Ok(client.put_protobuf_cells(table, cells, overwrite).await?)
302 })
303 .await
304 }
305}
306
307pub struct BigTable<F: FnMut(Request<()>) -> InterceptedRequestResult> {
308 access_token: Option<AccessToken>,
309 client: bigtable_client::BigtableClient<InterceptedService<tonic::transport::Channel, F>>,
310 table_prefix: String,
311 app_profile_id: String,
312 timeout: Option<Duration>,
313}
314
315impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
316 async fn decode_read_rows_response(
317 &self,
318 mut rrr: tonic::codec::Streaming<ReadRowsResponse>,
319 ) -> Result<Vec<(RowKey, RowData)>> {
320 let mut rows: Vec<(RowKey, RowData)> = vec![];
321
322 let mut row_key = None;
323 let mut row_data = vec![];
324
325 let mut cell_name = None;
326 let mut cell_timestamp = 0;
327 let mut cell_value = vec![];
328 let mut cell_version_ok = true;
329 let started = Instant::now();
330
331 while let Some(res) = rrr.message().await? {
332 if let Some(timeout) = self.timeout {
333 if Instant::now().duration_since(started) > timeout {
334 return Err(Error::Timeout);
335 }
336 }
337 for (i, mut chunk) in res.chunks.into_iter().enumerate() {
338 trace!("chunk {}: {:?}", i, chunk);
341
342 if !chunk.row_key.is_empty() {
344 row_key = String::from_utf8(chunk.row_key).ok(); }
346
347 if let Some(qualifier) = chunk.qualifier {
349 if let Some(cell_name) = cell_name {
350 row_data.push((cell_name, cell_value));
351 cell_value = vec![];
352 }
353 cell_name = String::from_utf8(qualifier).ok(); cell_timestamp = chunk.timestamp_micros;
355 cell_version_ok = true;
356 } else {
357 if chunk.timestamp_micros != 0 {
359 if chunk.timestamp_micros < cell_timestamp {
360 cell_version_ok = false; } else {
362 cell_version_ok = true;
364 cell_value = vec![];
365 cell_timestamp = chunk.timestamp_micros;
366 }
367 }
368 }
369 if cell_version_ok {
370 cell_value.append(&mut chunk.value);
371 }
372
373 if chunk.row_status.is_some() {
375 if let Some(read_rows_response::cell_chunk::RowStatus::CommitRow(_)) =
376 chunk.row_status
377 {
378 if let Some(cell_name) = cell_name {
379 row_data.push((cell_name, cell_value));
380 }
381
382 if let Some(row_key) = row_key {
383 rows.push((row_key, row_data))
384 }
385 }
386
387 row_key = None;
388 row_data = vec![];
389 cell_value = vec![];
390 cell_name = None;
391 }
392 }
393 }
394 Ok(rows)
395 }
396
397 async fn refresh_access_token(&self) {
398 if let Some(ref access_token) = self.access_token {
399 access_token.refresh().await;
400 }
401 }
402
403 pub async fn get_row_keys(
412 &mut self,
413 table_name: &str,
414 start_at: Option<RowKey>,
415 end_at: Option<RowKey>,
416 rows_limit: i64,
417 ) -> Result<Vec<RowKey>> {
418 self.refresh_access_token().await;
419 let response = self
420 .client
421 .read_rows(ReadRowsRequest {
422 table_name: format!("{}{}", self.table_prefix, table_name),
423 app_profile_id: self.app_profile_id.clone(),
424 rows_limit,
425 rows: Some(RowSet {
426 row_keys: vec![],
427 row_ranges: vec![RowRange {
428 start_key: start_at.map(|row_key| {
429 row_range::StartKey::StartKeyClosed(row_key.into_bytes())
430 }),
431 end_key: end_at
432 .map(|row_key| row_range::EndKey::EndKeyClosed(row_key.into_bytes())),
433 }],
434 }),
435 filter: Some(RowFilter {
436 filter: Some(row_filter::Filter::Chain(row_filter::Chain {
437 filters: vec![
438 RowFilter {
439 filter: Some(row_filter::Filter::CellsPerRowLimitFilter(1)),
441 },
442 RowFilter {
443 filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
445 },
446 RowFilter {
447 filter: Some(row_filter::Filter::StripValueTransformer(true)),
449 },
450 ],
451 })),
452 }),
453 })
454 .await?
455 .into_inner();
456
457 let rows = self.decode_read_rows_response(response).await?;
458 Ok(rows.into_iter().map(|r| r.0).collect())
459 }
460
461 pub async fn get_row_data(
474 &mut self,
475 table_name: &str,
476 start_at: Option<RowKey>,
477 end_at: Option<RowKey>,
478 rows_limit: i64,
479 ) -> Result<Vec<(RowKey, RowData)>> {
480 self.refresh_access_token().await;
481
482 let response = self
483 .client
484 .read_rows(ReadRowsRequest {
485 table_name: format!("{}{}", self.table_prefix, table_name),
486 app_profile_id: self.app_profile_id.clone(),
487 rows_limit,
488 rows: Some(RowSet {
489 row_keys: vec![],
490 row_ranges: vec![RowRange {
491 start_key: start_at.map(|row_key| {
492 row_range::StartKey::StartKeyClosed(row_key.into_bytes())
493 }),
494 end_key: end_at
495 .map(|row_key| row_range::EndKey::EndKeyClosed(row_key.into_bytes())),
496 }],
497 }),
498 filter: Some(RowFilter {
499 filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
501 }),
502 })
503 .await?
504 .into_inner();
505
506 self.decode_read_rows_response(response).await
507 }
508
509 pub async fn get_multi_row_data(
511 &mut self,
512 table_name: &str,
513 row_keys: &[RowKey],
514 ) -> Result<Vec<(RowKey, RowData)>> {
515 self.refresh_access_token().await;
516
517 let response = self
518 .client
519 .read_rows(ReadRowsRequest {
520 table_name: format!("{}{}", self.table_prefix, table_name),
521 app_profile_id: self.app_profile_id.clone(),
522 rows_limit: 0, rows: Some(RowSet {
524 row_keys: row_keys
525 .iter()
526 .map(|k| k.as_bytes().to_vec())
527 .collect::<Vec<_>>(),
528 row_ranges: vec![],
529 }),
530 filter: Some(RowFilter {
531 filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
533 }),
534 })
535 .await?
536 .into_inner();
537
538 self.decode_read_rows_response(response).await
539 }
540
541 pub async fn get_single_row_data(
547 &mut self,
548 table_name: &str,
549 row_key: RowKey,
550 ) -> Result<RowData> {
551 self.refresh_access_token().await;
552
553 let response = self
554 .client
555 .read_rows(ReadRowsRequest {
556 table_name: format!("{}{}", self.table_prefix, table_name),
557 app_profile_id: self.app_profile_id.clone(),
558 rows_limit: 1,
559 rows: Some(RowSet {
560 row_keys: vec![row_key.into_bytes()],
561 row_ranges: vec![],
562 }),
563 filter: Some(RowFilter {
564 filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
566 }),
567 })
568 .await?
569 .into_inner();
570
571 let rows = self.decode_read_rows_response(response).await?;
572 rows.into_iter()
573 .next()
574 .map(|r| r.1)
575 .ok_or(Error::RowNotFound)
576 }
577
578 async fn delete_rows(&mut self, table_name: &str, row_keys: &[RowKey]) -> Result<()> {
580 self.refresh_access_token().await;
581
582 let mut entries = vec![];
583 for row_key in row_keys {
584 entries.push(mutate_rows_request::Entry {
585 row_key: row_key.as_bytes().to_vec(),
586 mutations: vec![Mutation {
587 mutation: Some(mutation::Mutation::DeleteFromRow(
588 mutation::DeleteFromRow {},
589 )),
590 }],
591 });
592 }
593
594 let mut response = self
595 .client
596 .mutate_rows(MutateRowsRequest {
597 table_name: format!("{}{}", self.table_prefix, table_name),
598 app_profile_id: self.app_profile_id.clone(),
599 entries,
600 })
601 .await?
602 .into_inner();
603
604 while let Some(res) = response.message().await? {
605 for entry in res.entries {
606 if let Some(status) = entry.status {
607 if status.code != 0 {
608 eprintln!("delete_rows error {}: {}", status.code, status.message);
609 warn!("delete_rows error {}: {}", status.code, status.message);
610 return Err(Error::RowDeleteFailed);
611 }
612 }
613 }
614 }
615
616 Ok(())
617 }
618
619 async fn put_row_data(
621 &mut self,
622 table_name: &str,
623 family_name: &str,
624 row_data: &[(&RowKey, RowData)],
625 overwrite: bool,
626 ) -> Result<()> {
627 self.refresh_access_token().await;
628
629 let mut entries = vec![];
630 let timestamp = if overwrite {
631 0
634 } else {
635 -1
637 };
638 for (row_key, row_data) in row_data {
639 let mutations = row_data
640 .iter()
641 .map(|(column_key, column_value)| Mutation {
642 mutation: Some(mutation::Mutation::SetCell(mutation::SetCell {
643 family_name: family_name.to_string(),
644 column_qualifier: column_key.clone().into_bytes(),
645 timestamp_micros: timestamp,
646 value: column_value.to_vec(),
647 })),
648 })
649 .collect();
650
651 entries.push(mutate_rows_request::Entry {
652 row_key: (*row_key).clone().into_bytes(),
653 mutations,
654 });
655 }
656
657 let mut response = self
658 .client
659 .mutate_rows(MutateRowsRequest {
660 table_name: format!("{}{}", self.table_prefix, table_name),
661 app_profile_id: self.app_profile_id.clone(),
662 entries,
663 })
664 .await?
665 .into_inner();
666
667 while let Some(res) = response.message().await? {
668 for entry in res.entries {
669 if let Some(status) = entry.status {
670 if status.code != 0 {
671 eprintln!("put_row_data error {}: {}", status.code, status.message);
672 warn!("put_row_data error {}: {}", status.code, status.message);
673 return Err(Error::RowWriteFailed);
674 }
675 }
676 }
677 }
678
679 Ok(())
680 }
681
682 pub async fn get_bincode_cell<T>(&mut self, table: &str, key: RowKey) -> Result<T>
683 where
684 T: serde::de::DeserializeOwned,
685 {
686 let row_data = self.get_single_row_data(table, key.clone()).await?;
687 deserialize_bincode_cell_data(&row_data, table, key.to_string())
688 }
689
690 pub async fn get_bincode_cells<T>(
691 &mut self,
692 table: &str,
693 keys: &[RowKey],
694 ) -> Result<Vec<(RowKey, Result<T>)>>
695 where
696 T: serde::de::DeserializeOwned,
697 {
698 Ok(self
699 .get_multi_row_data(table, keys)
700 .await?
701 .into_iter()
702 .map(|(key, row_data)| {
703 let key_str = key.to_string();
704 (
705 key,
706 deserialize_bincode_cell_data(&row_data, table, key_str),
707 )
708 })
709 .collect())
710 }
711
712 pub async fn get_protobuf_or_bincode_cell<B, P>(
713 &mut self,
714 table: &str,
715 key: RowKey,
716 ) -> Result<CellData<B, P>>
717 where
718 B: serde::de::DeserializeOwned,
719 P: prost::Message + Default,
720 {
721 let row_data = self.get_single_row_data(table, key.clone()).await?;
722 deserialize_protobuf_or_bincode_cell_data(&row_data, table, key)
723 }
724
725 pub async fn get_protobuf_or_bincode_cells<'a, B, P>(
726 &mut self,
727 table: &'a str,
728 row_keys: impl IntoIterator<Item = RowKey>,
729 ) -> Result<impl Iterator<Item = (RowKey, CellData<B, P>)> + 'a>
730 where
731 B: serde::de::DeserializeOwned,
732 P: prost::Message + Default,
733 {
734 Ok(self
735 .get_multi_row_data(
736 table,
737 row_keys.into_iter().collect::<Vec<RowKey>>().as_slice(),
738 )
739 .await?
740 .into_iter()
741 .map(|(key, row_data)| {
742 let key_str = key.to_string();
743 (
744 key,
745 deserialize_protobuf_or_bincode_cell_data(&row_data, table, key_str).unwrap(),
746 )
747 }))
748 }
749
750 pub async fn put_bincode_cells<T>(
751 &mut self,
752 table: &str,
753 cells: &[(RowKey, T)],
754 overwrite: bool,
755 ) -> Result<usize>
756 where
757 T: serde::ser::Serialize,
758 {
759 let mut bytes_written = 0;
760 let mut new_row_data = vec![];
761 for (row_key, data) in cells {
762 let data = compress_best(&bincode::serialize(&data).unwrap())?;
763 bytes_written += data.len();
764 new_row_data.push((row_key, vec![("bin".to_string(), data)]));
765 }
766
767 self.put_row_data(table, "x", &new_row_data, overwrite)
768 .await?;
769 Ok(bytes_written)
770 }
771
772 pub async fn put_protobuf_cells<T>(
773 &mut self,
774 table: &str,
775 cells: &[(RowKey, T)],
776 overwrite: bool,
777 ) -> Result<usize>
778 where
779 T: prost::Message,
780 {
781 let mut bytes_written = 0;
782 let mut new_row_data = vec![];
783 for (row_key, data) in cells {
784 let mut buf = Vec::with_capacity(data.encoded_len());
785 data.encode(&mut buf).unwrap();
786 let data = compress_best(&buf)?;
787 bytes_written += data.len();
788 new_row_data.push((row_key, vec![("proto".to_string(), data)]));
789 }
790
791 let result = self
792 .put_row_data(table, "x", &new_row_data, overwrite)
793 .await;
794 match result {
795 Err(err) => {
796 for (key, data) in new_row_data.iter() {
797 for (col, cell) in data.iter() {
798 error!(
799 "Error writing account key {} col {} len: {}",
800 *key,
801 col,
802 cell.len()
803 );
804 }
805 }
806
807 return Err(err);
808 }
809 Ok(_) => {}
810 }
811 Ok(bytes_written)
812 }
813}
814
815pub(crate) fn deserialize_protobuf_or_bincode_cell_data<B, P>(
816 row_data: RowDataSlice,
817 table: &str,
818 key: RowKey,
819) -> Result<CellData<B, P>>
820where
821 B: serde::de::DeserializeOwned,
822 P: prost::Message + Default,
823{
824 match deserialize_protobuf_cell_data(row_data, table, key.to_string()) {
825 Ok(result) => return Ok(CellData::Protobuf(result)),
826 Err(err) => match err {
827 Error::ObjectNotFound(_) => {}
828 _ => return Err(err),
829 },
830 }
831 deserialize_bincode_cell_data(row_data, table, key).map(CellData::Bincode)
832}
833
834pub(crate) fn deserialize_protobuf_cell_data<T>(
835 row_data: RowDataSlice,
836 table: &str,
837 key: RowKey,
838) -> Result<T>
839where
840 T: prost::Message + Default,
841{
842 let value = &row_data
843 .iter()
844 .find(|(name, _)| name == "proto")
845 .ok_or_else(|| Error::ObjectNotFound(format!("{}/{}", table, key)))?
846 .1;
847
848 let data = decompress(value)?;
849 T::decode(&data[..]).map_err(|err| {
850 warn!("Failed to deserialize {}/{}: {}", table, key, err);
851 Error::ObjectCorrupt(format!("{}/{}", table, key))
852 })
853}
854
855pub(crate) fn deserialize_bincode_cell_data<T>(
856 row_data: RowDataSlice,
857 table: &str,
858 key: RowKey,
859) -> Result<T>
860where
861 T: serde::de::DeserializeOwned,
862{
863 let value = &row_data
864 .iter()
865 .find(|(name, _)| name == "bin")
866 .ok_or_else(|| Error::ObjectNotFound(format!("{}/{}", table, key)))?
867 .1;
868
869 let data = decompress(value)?;
870 bincode::deserialize(&data).map_err(|err| {
871 warn!("Failed to deserialize {}/{}: {}", table, key, err);
872 Error::ObjectCorrupt(format!("{}/{}", table, key))
873 })
874}