From e8bc60ba711a80414abc6ad5c5a71ebb9aeee827 Mon Sep 17 00:00:00 2001
From: Micaiah Reid <4290054+vabanaerytk@users.noreply.github.com>
Date: Tue, 19 Mar 2024 11:08:37 -0400
Subject: [PATCH] feat: use stacks.rocksdb for predicate scan (#514)

### Description

Previously, when using the `chainhook predicates scan` command,
Chainhook would download a very large archive file to use as a data
source for the scan **every time the command was run**.

The Chainhook service (`chainhook service start`), however, implements
rocksdb to allow storing this dataset locally.

This PR uses the rocksdb that a user may have created locally to improve
performance of the `chainhook predicates scan` command. If a rocksdb
doesn't exist or can't be created, Chainhook will fallback to the
original method of downloading the archive.

Combined with #513, this PR reduces a simple scan of a few blocks from
~20 minutes to .02 seconds.

Fixes #485
---
 components/chainhook-cli/src/cli/mod.rs | 38 +++++++++++++++++++------
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/components/chainhook-cli/src/cli/mod.rs b/components/chainhook-cli/src/cli/mod.rs
index 3fa1b1a..ed51d0b 100644
--- a/components/chainhook-cli/src/cli/mod.rs
+++ b/components/chainhook-cli/src/cli/mod.rs
@@ -3,6 +3,7 @@ use crate::config::Config;
 use crate::scan::bitcoin::scan_bitcoin_chainstate_via_rpc_using_predicate;
 use crate::scan::stacks::{
     consolidate_local_stacks_chainstate_using_csv, scan_stacks_chainstate_via_csv_using_predicate,
+    scan_stacks_chainstate_via_rocksdb_using_predicate,
 };
 use crate::service::http_api::document_predicate_api_server;
 use crate::service::Service;
@@ -485,14 +486,35 @@ async fn handle_command(opts: Opts, ctx: Context) -> Result<(), String> {
                                 ));
                             }
                         };
-                        // TODO: if a stacks.rocksdb is present, use it.
-                        // TODO: update Stacks archive file if required.
-                        scan_stacks_chainstate_via_csv_using_predicate(
-                            &predicate_spec,
-                            &mut config,
-                            &ctx,
-                        )
-                        .await?;
+                        match open_readonly_stacks_db_conn(&config.expected_cache_path(), &ctx) {
+                            Ok(db_conn) => {
+                                let _ = consolidate_local_stacks_chainstate_using_csv(
+                                    &mut config,
+                                    &ctx,
+                                )
+                                .await;
+                                scan_stacks_chainstate_via_rocksdb_using_predicate(
+                                    &predicate_spec,
+                                    None,
+                                    &db_conn,
+                                    &config,
+                                    &ctx,
+                                )
+                                .await?;
+                            }
+                            Err(e) => {
+                                info!(
+                                    ctx.expect_logger(),
+                                    "Could not open db. This will greatly increase scan times. Error: {}", e
+                                );
+                                scan_stacks_chainstate_via_csv_using_predicate(
+                                    &predicate_spec,
+                                    &mut config,
+                                    &ctx,
+                                )
+                                .await?;
+                            }
+                        };
                     }
                 }
             }