[data ingestion] add support for bigtable in main workflow binary (#2…

…0018)

## Description 

adding BigTable KV as a variant to main binary that manages internal
workflows

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] Indexer: 
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [ ] CLI: 
- [ ] Rust SDK:
- [ ] REST API:

Cargo.toml

@@ -646,6 +646,7 @@ sui-json-rpc = { path = "crates/sui-json-rpc" }
 sui-json-rpc-api = { path = "crates/sui-json-rpc-api" }
 sui-json-rpc-types = { path = "crates/sui-json-rpc-types" }
 sui-keys = { path = "crates/sui-keys" }
+sui-kvstore = {path = "crates/sui-kvstore"}
 sui-macros = { path = "crates/sui-macros" }
 sui-metric-checker = { path = "crates/sui-metric-checker" }
 sui-move = { path = "crates/sui-move" }

crates/sui-data-ingestion/Cargo.toml

@@ -31,6 +31,7 @@ tracing.workspace = true
 sui-archival.workspace = true
 sui-storage.workspace = true
 sui-data-ingestion-core.workspace = true
+sui-kvstore.workspace = true
 sui-types.workspace = true
 tempfile.workspace = true
 url.workspace = true

crates/sui-data-ingestion/src/main.rs

@@ -12,6 +12,7 @@ use sui_data_ingestion::{
 };
 use sui_data_ingestion_core::{DataIngestionMetrics, ReaderOptions};
 use sui_data_ingestion_core::{IndexerExecutor, WorkerPool};
+use sui_kvstore::{BigTableClient, KvWorker};
 use tokio::signal;
 use tokio::sync::oneshot;
 
@@ -21,6 +22,7 @@ enum Task {
     Archival(ArchivalConfig),
     Blob(BlobTaskConfig),
     KV(KVStoreTaskConfig),
+    BigTableKV(BigTableTaskConfig),
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug)]
@@ -40,6 +42,11 @@ struct ProgressStoreConfig {
     pub table_name: String,
 }
 
+#[derive(Serialize, Deserialize, Clone, Debug)]
+struct BigTableTaskConfig {
+    instance_id: String,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct IndexerConfig {
     path: PathBuf,
@@ -146,6 +153,15 @@ async fn main() -> Result<()> {
                 );
                 executor.register(worker_pool).await?;
             }
+            Task::BigTableKV(kv_config) => {
+                let client = BigTableClient::new_remote(kv_config.instance_id, false, None).await?;
+                let worker_pool = WorkerPool::new(
+                    KvWorker { client },
+                    task_config.name,
+                    task_config.concurrency,
+                );
+                executor.register(worker_pool).await?;
+            }
         };
     }
     let reader_options = ReaderOptions {

crates/sui-kvstore/src/bigtable/mod.rs

@@ -3,3 +3,4 @@
 
 pub(crate) mod client;
 mod proto;
+pub(crate) mod worker;

crates/sui-kvstore/src/bigtable/worker.rs

@@ -0,0 +1,39 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{BigTableClient, KeyValueStoreWriter, TransactionData};
+use async_trait::async_trait;
+use sui_data_ingestion_core::Worker;
+use sui_types::full_checkpoint_content::CheckpointData;
+
+pub struct KvWorker {
+    pub client: BigTableClient,
+}
+
+#[async_trait]
+impl Worker for KvWorker {
+    type Result = ();
+
+    async fn process_checkpoint(&self, checkpoint: &CheckpointData) -> anyhow::Result<()> {
+        let mut client = self.client.clone();
+        let mut objects = vec![];
+        let mut transactions = vec![];
+        for transaction in &checkpoint.transactions {
+            let full_transaction = TransactionData {
+                transaction: transaction.transaction.clone(),
+                effects: transaction.effects.clone(),
+                events: transaction.events.clone(),
+                checkpoint_number: checkpoint.checkpoint_summary.sequence_number,
+                timestamp: checkpoint.checkpoint_summary.timestamp_ms,
+            };
+            for object in &transaction.output_objects {
+                objects.push(object);
+            }
+            transactions.push(full_transaction);
+        }
+        client.save_objects(&objects).await?;
+        client.save_transactions(&transactions).await?;
+        client.save_checkpoint(checkpoint).await?;
+        Ok(())
+    }
+}

crates/sui-kvstore/src/lib.rs

@@ -4,6 +4,7 @@ mod bigtable;
 use anyhow::Result;
 use async_trait::async_trait;
 pub use bigtable::client::BigTableClient;
+pub use bigtable::worker::KvWorker;
 use sui_types::crypto::AuthorityStrongQuorumSignInfo;
 use sui_types::digests::TransactionDigest;
 use sui_types::effects::{TransactionEffects, TransactionEvents};

crates/sui-kvstore/src/main.rs

@@ -1,44 +1,11 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 use anyhow::Result;
-use async_trait::async_trait;
-use sui_data_ingestion_core::{setup_single_workflow, Worker};
-use sui_kvstore::{BigTableClient, KeyValueStoreWriter, TransactionData};
-use sui_types::full_checkpoint_content::CheckpointData;
+use sui_data_ingestion_core::setup_single_workflow;
+use sui_kvstore::BigTableClient;
+use sui_kvstore::KvWorker;
 use telemetry_subscribers::TelemetryConfig;
 
-struct KvWorker {
-    client: BigTableClient,
-}
-
-#[async_trait]
-impl Worker for KvWorker {
-    type Result = ();
-
-    async fn process_checkpoint(&self, checkpoint: &CheckpointData) -> Result<()> {
-        let mut client = self.client.clone();
-        let mut objects = vec![];
-        let mut transactions = vec![];
-        for transaction in &checkpoint.transactions {
-            let full_transaction = TransactionData {
-                transaction: transaction.transaction.clone(),
-                effects: transaction.effects.clone(),
-                events: transaction.events.clone(),
-                checkpoint_number: checkpoint.checkpoint_summary.sequence_number,
-                timestamp: checkpoint.checkpoint_summary.timestamp_ms,
-            };
-            for object in &transaction.output_objects {
-                objects.push(object);
-            }
-            transactions.push(full_transaction);
-        }
-        client.save_objects(&objects).await?;
-        client.save_transactions(&transactions).await?;
-        client.save_checkpoint(checkpoint).await?;
-        Ok(())
-    }
-}
-
 #[tokio::main]
 async fn main() -> Result<()> {
     let _guard = TelemetryConfig::new().with_env().init();