linting + add main CLI entry arg

getsentry · Jun 19, 2024 · e5216d0 · e5216d0
1 parent fd20804
commit e5216d0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/rust_snuba/src/consumer.rs b/rust_snuba/src/consumer.rs
@@ -22,7 +22,6 @@ use crate::logging::{setup_logging, setup_sentry};
 use crate::metrics::global_tags::set_global_tag;
 use crate::metrics::statsd::StatsDBackend;
 use crate::processors;
-use crate::strategies::clickhouse::batch;
 use crate::types::{InsertOrReplacement, KafkaMessageMetadata};
 
 #[pyfunction]
@@ -87,8 +86,7 @@ pub fn consumer_impl(
     let max_batch_size = consumer_config.max_batch_size;
     let max_batch_time = Duration::from_millis(consumer_config.max_batch_time_ms);
 
-    let batch_write_timeout =
-        batch_write_timeout_ms.map_or(None, |v| Some(Duration::from_millis(v)));
+    let batch_write_timeout = batch_write_timeout_ms.map(|v| Some(Duration::from_millis(v)));
 
     for storage in &consumer_config.storages {
         tracing::info!(

diff --git a/snuba/cli/rust_consumer.py b/snuba/cli/rust_consumer.py
@@ -186,6 +186,7 @@ def rust_consumer(
     health_check_file: Optional[str],
     enforce_schema: bool,
     stop_at_timestamp: Optional[int],
+    batch_write_timeout_ms: Optional[int]
 ) -> None:
     """
     Experimental alternative to `snuba consumer`
@@ -235,6 +236,7 @@ def rust_consumer(
         python_max_queue_depth,
         health_check_file,
         stop_at_timestamp,
+        batch_write_timeout_ms,
     )
 
     sys.exit(exitcode)