Skip to main content

solana_bigtable_connection/
bigtable.rs

1// Primitives for reading/writing BigTable tables
2
3use {
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    /// Establish a connection to the BigTable instance named `instance_name`.  If read-only access
121    /// is required, the `read_only` flag should be used to reduce the requested OAuth2 scope.
122    ///
123    /// The GOOGLE_APPLICATION_CREDENTIALS environment variable will be used to determine the
124    /// program name that contains the BigTable instance in addition to access credentials.
125    ///
126    /// The BIGTABLE_EMULATOR_HOST environment variable is also respected.
127    ///
128    /// The BIGTABLE_PROXY environment variable is used to configure the gRPC connection through a
129    /// forward proxy (see HTTP_PROXY).
130    ///
131    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                        // tonic handles TLS as a separate layer
202                        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    /// Create a new BigTable client.
220    ///
221    /// Clients require `&mut self`, due to `Tonic::transport::Channel` limitations, however
222    /// creating new clients is cheap and thus can be used as a work around for ease of use.
223    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                // The comments for `read_rows_response::CellChunk` provide essential details for
339                // understanding how the below decoding works...
340                trace!("chunk {}: {:?}", i, chunk);
341
342                // Starting a new row?
343                if !chunk.row_key.is_empty() {
344                    row_key = String::from_utf8(chunk.row_key).ok(); // Require UTF-8 for row keys
345                }
346
347                // Starting a new cell?
348                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(); // Require UTF-8 for cell names
354                    cell_timestamp = chunk.timestamp_micros;
355                    cell_version_ok = true;
356                } else {
357                    // Continuing the existing cell.  Check if this is the start of another version of the cell
358                    if chunk.timestamp_micros != 0 {
359                        if chunk.timestamp_micros < cell_timestamp {
360                            cell_version_ok = false; // ignore older versions of the cell
361                        } else {
362                            // newer version of the cell, remove the older cell
363                            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                // End of a row?
374                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    /// Get `table` row keys in lexical order.
404    ///
405    /// If `start_at` is provided, the row key listing will start with key.
406    /// Otherwise the listing will start from the start of the table.
407    ///
408    /// If `end_at` is provided, the row key listing will end at the key. Otherwise it will
409    /// continue until the `rows_limit` is reached or the end of the table, whichever comes first.
410    /// If `rows_limit` is zero, the listing will continue until the end of the table.
411    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                                // Return minimal number of cells
440                                filter: Some(row_filter::Filter::CellsPerRowLimitFilter(1)),
441                            },
442                            RowFilter {
443                                // Only return the latest version of each cell
444                                filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
445                            },
446                            RowFilter {
447                                // Strip the cell values
448                                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    /// Get latest data from `table`.
462    ///
463    /// All column families are accepted, and only the latest version of each column cell will be
464    /// returned.
465    ///
466    /// If `start_at` is provided, the row key listing will start with key, or the next key in the
467    /// table if the explicit key does not exist. Otherwise the listing will start from the start
468    /// of the table.
469    ///
470    /// If `end_at` is provided, the row key listing will end at the key. Otherwise it will
471    /// continue until the `rows_limit` is reached or the end of the table, whichever comes first.
472    /// If `rows_limit` is zero, the listing will continue until the end of the table.
473    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                    // Only return the latest version of each cell
500                    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    /// Get latest data from multiple rows of `table`, if those rows exist.
510    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, // return all keys
523                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                    // Only return the latest version of each cell
532                    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    /// Get latest data from a single row of `table`, if that row exists. Returns an error if that
542    /// row does not exist.
543    ///
544    /// All column families are accepted, and only the latest version of each column cell will be
545    /// returned.
546    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                    // Only return the latest version of each cell
565                    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    /// Delete one or more `table` rows
579    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    /// Store data for one or more `table` rows in the `family_name` Column family
620    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            // overwriting existing cell
632            // https://cloud.google.com/bigtable/docs/gc-latest-value
633            0
634        } else {
635            // server assigned
636            -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}