Add bigtable emulator and first tests.

src/bigtable/emulator.rs

@@ -0,0 +1,110 @@
+//! Testing infra to make use of the bigtable emulator.
+//! <https://cloud.google.com/bigtable/docs/emulator>
+//!
+//! Follow installation directions from link above to set up your local development. Once setup,
+//! you should be able to run the pubsub emulator driven tests.
+
+use futures::{future::BoxFuture, FutureExt};
+
+use crate::{
+    bigtable,
+    builder::ClientBuilder,
+    emulator::{self, EmulatorData},
+};
+
+type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+/// Struct to hold a started PubSub emulator process. Process is closed when struct is dropped.
+pub struct EmulatorClient {
+    inner: crate::emulator::EmulatorClient,
+    instance: String,
+}
+
+const DATA: EmulatorData = EmulatorData {
+    gcloud_param: "bigtable",
+    kill_pattern: "bigtable",
+    availability_check: create_bigtable_client,
+    extra_args: Vec::new(),
+};
+
+const INSTANCE_ID: &str = "test-instance";
+
+impl EmulatorClient {
+    /// Create a new emulator instance with a default project name and instance name
+    pub async fn new() -> Result<Self, BoxError> {
+        Ok(EmulatorClient {
+            inner: emulator::EmulatorClient::new(DATA).await?,
+            instance: INSTANCE_ID.into(),
+        })
+    }
+
+    /// Create a new emulator instance with the given project name and instance name
+    pub async fn with_project_and_instance(
+        project_name: impl Into<String>,
+        instance_name: impl Into<String>,
+    ) -> Result<Self, BoxError> {
+        Ok(EmulatorClient {
+            inner: emulator::EmulatorClient::with_project(DATA, project_name).await?,
+            instance: instance_name.into(),
+        })
+    }
+
+    /// Get the endpoint at which the emulator is listening for requests
+    pub fn endpoint(&self) -> String {
+        self.inner.endpoint()
+    }
+
+    /// Get the project name with which the emulator was initialized
+    pub fn project(&self) -> &str {
+        self.inner.project()
+    }
+
+    /// Get the instance name with which the emulator was initialized
+    pub fn instance(&self) -> &str {
+        &self.instance
+    }
+
+    /// Get a client builder which is pre-configured to work with this emulator instance
+    pub fn builder(&self) -> &ClientBuilder {
+        self.inner.builder()
+    }
+
+    /// Create a new table under this emulator's given project name and instance name.
+    ///
+    /// The column families will be created wth effectively no garbage collection.
+    pub async fn create_table(
+        &self,
+        table_name: &str,
+        column_families: impl IntoIterator<Item = impl Into<String>>,
+    ) -> Result<(), BoxError> {
+        let config = bigtable::admin::BigtableTableAdminConfig {
+            endpoint: self.endpoint(),
+            ..bigtable::admin::BigtableTableAdminConfig::default()
+        };
+
+        let mut admin = self
+            .builder()
+            .build_bigtable_admin_client(config, &self.project(), &self.instance)
+            .await?;
+
+        let column_families = column_families
+            .into_iter()
+            .map(|name| (name.into(), bigtable::admin::Rule::MaxNumVersions(i32::MAX)));
+        admin.create_table(table_name, column_families).await?;
+
+        Ok(())
+    }
+}
+
+fn create_bigtable_client(port: &str) -> BoxFuture<Result<(), tonic::transport::Error>> {
+    async move {
+        bigtable::api::bigtable::v2::bigtable_client::BigtableClient::connect(format!(
+            "http://{}:{}",
+            crate::emulator::HOST,
+            port
+        ))
+        .await?;
+        Ok(())
+    }
+    .boxed()
+}

src/bigtable/mod.rs

@@ -23,6 +23,10 @@ pub mod mutation;
 pub use client_builder::BigtableConfig;
 pub use mutation::{MutateRowRequest, MutateRowsError, MutateRowsRequest};
 
+#[cfg(feature = "emulators")]
+#[cfg_attr(docsrs, doc(cfg(feature = "emulators")))]
+pub mod emulator;
+
 #[allow(rustdoc::broken_intra_doc_links, rustdoc::bare_urls, missing_docs)]
 pub mod api {
     pub mod rpc {

tests/bigtable_client.rs

@@ -0,0 +1,153 @@
+#[cfg(all(feature = "bigtable", feature = "emulators"))]
+mod bigtable_client_tests {
+    use futures::TryStreamExt;
+    use hyper::body::Bytes;
+    use ya_gcp::bigtable::{self, admin::Rule, api, emulator::EmulatorClient, ReadRowsRequest};
+
+    #[tokio::test]
+    async fn create_table() {
+        let table_name = "test-table";
+        let emulator = EmulatorClient::new().await.unwrap();
+        let config = bigtable::admin::BigtableTableAdminConfig::new().endpoint(emulator.endpoint());
+        let mut admin = emulator
+            .builder()
+            .build_bigtable_admin_client(config, emulator.project(), emulator.instance())
+            .await
+            .unwrap();
+
+        admin
+            .create_table(table_name, [("column".into(), Rule::MaxNumVersions(1))])
+            .await
+            .unwrap();
+
+        let tables: Vec<_> = admin
+            .list_tables()
+            .await
+            .unwrap()
+            .try_collect()
+            .await
+            .unwrap();
+
+        assert_eq!(tables.len(), 1);
+        assert_eq!(
+            tables[0].name,
+            "projects/test-project/instances/test-instance/tables/test-table"
+        );
+    }
+
+    async fn default_client(table_name: &str) -> (EmulatorClient, bigtable::BigtableClient) {
+        let emulator = EmulatorClient::new().await.unwrap();
+        emulator
+            .create_table(table_name, ["fam1", "fam2"])
+            .await
+            .unwrap();
+
+        let config = bigtable::BigtableConfig::new().endpoint(emulator.endpoint());
+        let client = emulator
+            .builder()
+            .build_bigtable_client(config, emulator.project(), emulator.instance())
+            .await
+            .unwrap();
+        (emulator, client)
+    }
+
+    #[tokio::test]
+    async fn set_and_read_row() {
+        let table_name = "test-table";
+        let (_emulator, mut client) = default_client(table_name).await;
+
+        client
+            .set_row_data(
+                table_name,
+                "fam1".into(),
+                "row1-key",
+                [("col1", "data1"), ("col2", "data2")],
+            )
+            .await
+            .unwrap();
+
+        let row = client
+            .read_one_row(table_name, "row1-key")
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!("row1-key", row.key);
+        assert_eq!(1, row.families.len());
+        assert_eq!("fam1", row.families[0].name);
+        dbg!(&row.families);
+        assert_eq!(2, row.families[0].columns.len());
+        assert_eq!("col1", row.families[0].columns[0].qualifier);
+        assert_eq!("data1", row.families[0].columns[0].cells[0].value);
+        assert_eq!("col2", row.families[0].columns[1].qualifier);
+        assert_eq!("data2", row.families[0].columns[1].cells[0].value);
+
+        let row_range: Vec<_> = client
+            .read_row_range(
+                table_name,
+                Bytes::from("row1-key")..=Bytes::from("row1-key"),
+                None,
+            )
+            .try_collect()
+            .await
+            .unwrap();
+        assert_eq!(row_range, vec![row]);
+    }
+
+    #[tokio::test]
+    async fn cell_versions() {
+        let table_name = "test-table";
+        let (emulator, mut client) = default_client(table_name).await;
+        client
+            .set_row_data_with_timestamp(
+                table_name,
+                "fam1".into(),
+                // The bigtable emulator enforces millisecond granularity
+                6000,
+                "row1-key",
+                [("col1", "data1")],
+            )
+            .await
+            .unwrap();
+        client
+            .set_row_data_with_timestamp(
+                table_name,
+                "fam1".into(),
+                7000,
+                "row1-key",
+                [("col1", "data2")],
+            )
+            .await
+            .unwrap();
+        let row = client
+            .read_one_row(table_name, "row1-key")
+            .await
+            .unwrap()
+            .unwrap();
+        // read_one_row only returns the latest version
+        assert_eq!(1, row.families[0].columns[0].cells.len());
+        assert_eq!("data2", row.families[0].columns[0].cells[0].value);
+        assert_eq!(
+            vec!["data2"],
+            row.most_recent_cells().map(|c| c.value).collect::<Vec<_>>()
+        );
+
+        let req = ReadRowsRequest {
+            table_name: format!(
+                "projects/{}/instances/{}/tables/{table_name}",
+                emulator.project(),
+                emulator.instance()
+            ),
+            rows: Some(api::bigtable::v2::RowSet::default().with_key("row1-key")),
+            ..Default::default()
+        };
+        let rows: Vec<_> = client.read_rows(req).try_collect().await.unwrap();
+        assert_eq!(1, rows.len());
+        let row = &rows[0];
+        dbg!(row);
+        assert_eq!(2, row.families[0].columns[0].cells.len());
+        assert_eq!(
+            vec!["data2"],
+            row.most_recent_cells().map(|c| c.value).collect::<Vec<_>>()
+        );
+    }
+}