feat(profiles): Add dataset configurations and processor for profile chunks

rust_snuba/Cargo.toml

@@ -33,7 +33,7 @@ pyo3 = { version = "0.18.1", features = ["chrono"] }
 reqwest = { version = "0.11.11", features = ["stream"] }
 rust_arroyo = { version = "*", git = "https://github.com/getsentry/arroyo" }
 sentry = { version = "0.32.0", features = ["anyhow", "tracing"] }
-sentry-kafka-schemas = "0.1.81"
+sentry-kafka-schemas = "0.1.82"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0" }
 thiserror = "1.0"

rust_snuba/src/processors/mod.rs

@@ -3,6 +3,7 @@ mod functions;
 mod generic_metrics;
 mod metrics_summaries;
 mod outcomes;
+mod profile_chunks;
 mod profiles;
 mod querylog;
 mod release_health_metrics;
@@ -62,6 +63,7 @@ define_processing_functions! {
     ("GenericGaugesMetricsProcessor", "snuba-generic-metrics", ProcessingFunctionType::ProcessingFunction(generic_metrics::process_gauge_message)),
     ("PolymorphicMetricsProcessor", "snuba-metrics", ProcessingFunctionType::ProcessingFunction(release_health_metrics::process_metrics_message)),
     ("ErrorsProcessor", "events", ProcessingFunctionType::ProcessingFunctionWithReplacements(errors::process_message_with_replacement)),
+    ("ProfileChunksProcessor", "snuba-profile-chunks", ProcessingFunctionType::ProcessingFunction(profile_chunks::process_message)),
 }
 
 // COGS is recorded for these processors

rust_snuba/src/processors/profile_chunks.rs

@@ -0,0 +1,91 @@
+use crate::config::ProcessorConfig;
+use anyhow::Context;
+use chrono::DateTime;
+use rust_arroyo::backends::kafka::types::KafkaPayload;
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+use crate::processors::utils::enforce_retention;
+use crate::types::{InsertBatch, KafkaMessageMetadata};
+
+pub fn process_message(
+    payload: KafkaPayload,
+    metadata: KafkaMessageMetadata,
+    config: &ProcessorConfig,
+) -> anyhow::Result<InsertBatch> {
+    let payload_bytes = payload.payload().context("Expected payload")?;
+    let msg: InputMessage = serde_json::from_slice(payload_bytes)?;
+
+    let mut row = Chunk {
+        chunk: msg,
+        offset: metadata.offset,
+        partition: metadata.partition,
+    };
+
+    row.chunk.retention_days = Some(enforce_retention(
+        row.chunk.retention_days,
+        &config.env_config,
+    ));
+
+    let origin_timestamp = DateTime::from_timestamp(row.chunk.received, 0);
+
+    InsertBatch::from_rows([row], origin_timestamp)
+}
+
+#[derive(Debug, Deserialize, Serialize, JsonSchema)]
+struct InputMessage {
+    project_id: u64,
+    profiler_id: Uuid,
+    chunk_id: Uuid,
+    start_timestamp: f64,
+    end_timestamp: f64,
+    received: i64,
+    retention_days: Option<u16>,
+}
+
+#[derive(Debug, Deserialize, Serialize, JsonSchema)]
+struct Chunk {
+    #[serde(flatten)]
+    chunk: InputMessage,
+
+    #[serde(default)]
+    offset: u64,
+    #[serde(default)]
+    partition: u16,
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::SystemTime;
+
+    use crate::processors::tests::run_schema_type_test;
+
+    use super::*;
+
+    #[test]
+    fn test_chunk() {
+        let data = r#"{
+            "chunk_id": "0432a0a4c25f4697bf9f0a2fcbe6a814",
+            "end_timestamp": 1710805689.1234567,
+            "profiler_id": "4d229f1d3807421ba62a5f8bc295d836",
+            "project_id": 1,
+            "received": 1694357860,
+            "retention_days": 30,
+            "start_timestamp": 1710805688.1234567
+        }"#;
+        let payload = KafkaPayload::new(None, None, Some(data.as_bytes().to_vec()));
+        let meta = KafkaMessageMetadata {
+            partition: 0,
+            offset: 1,
+            timestamp: DateTime::from(SystemTime::now()),
+        };
+        process_message(payload, meta, &ProcessorConfig::default())
+            .expect("The message should be processed");
+    }
+
+    #[test]
+    fn schema() {
+        run_schema_type_test::<InputMessage>("snuba-profile-chunks", None);
+    }
+}

rust_snuba/src/processors/snapshots/rust_snuba__processors__tests__schemas@snuba-profile-chunks-ProfileChunksProcessor-snuba-profile-chunks__1__valid.json.snap

@@ -0,0 +1,18 @@
+---
+source: src/processors/mod.rs
+description: "{\n  \"project_id\": 1,\n  \"profiler_id\": \"4d229f1d3807421ba62a5f8bc295d836\",\n  \"chunk_id\": \"0432a0a4c25f4697bf9f0a2fcbe6a814\",\n  \"start_timestamp\": 1710805688.1234567,\n  \"end_timestamp\": 1710805689.1234567,\n  \"received\": 1694357860,\n  \"retention_days\": 30\n}\n"
+expression: snapshot_payload
+---
+[
+  {
+    "chunk_id": "0432a0a4-c25f-4697-bf9f-0a2fcbe6a814",
+    "end_timestamp": 1710805689.1234567,
+    "offset": 1,
+    "partition": 0,
+    "profiler_id": "4d229f1d-3807-421b-a62a-5f8bc295d836",
+    "project_id": 1,
+    "received": 1694357860,
+    "retention_days": 30,
+    "start_timestamp": 1710805688.1234567
+  }
+]

snuba/datasets/configuration/profiles/storages/chunks.yaml

@@ -0,0 +1,55 @@
+version: v1
+kind: writable_storage
+name: profile_chunks
+storage:
+  key: profile_chunks
+  set_key: profile_chunks
+readiness_state: partial
+schema:
+  columns:
+    [
+      { name: project_id, type: UInt, args: { size: 64 } },
+      { name: profiler_id, type: UUID },
+      { name: chunk_id, type: UUID },
+      { name: start_timestamp, type: DateTime64, args: { precision: 6 } },
+      { name: end_timestamp, type: DateTime64, args: { precision: 6 } },
+      { name: retention_days, type: UInt, args: { size: 16 } },
+      { name: partition, type: UInt, args: { size: 16 } },
+      { name: offset, type: UInt, args: { size: 64 } },
+    ]
+  local_table_name: profile_chunks_local
+  dist_table_name: profile_chunks_dist
+query_processors:
+  - processor: UUIDColumnProcessor
+    args:
+      columns: !!set
+        profiler_id: null
+        chunk_id: null
+  - processor: TableRateLimit
+allocation_policies:
+  - name: ConcurrentRateLimitAllocationPolicy
+    args:
+      required_tenant_types:
+        - referrer
+        - project_id
+      default_config_overrides:
+        is_enforced: 1
+  - name: BytesScannedWindowAllocationPolicy
+    args:
+      required_tenant_types:
+        - referrer
+      default_config_overrides:
+        is_enforced: 1
+        throttled_thread_number: 1
+        org_limit_bytes_scanned: 100000
+  - name: ReferrerGuardRailPolicy
+    args:
+      required_tenant_types:
+        - referrer
+      default_config_overrides:
+        is_enforced: 1
+mandatory_condition_checkers:
+  - condition: ProjectIdEnforcer
+stream_loader:
+  processor: ProfileChunksProcessor
+  default_topic: snuba-profile-chunks

snuba/datasets/processors/profile_chunks_processor.py

@@ -0,0 +1,6 @@
+from snuba.datasets.processors.rust_compat_processor import RustCompatProcessor
+
+
+class ProfileChunksProcessor(RustCompatProcessor):
+    def __init__(self) -> None:
+        super().__init__("ProfileChunksProcessor")

snuba/utils/streams/topics.py

@@ -39,8 +39,11 @@ class Topic(Enum):
     SUBSCRIPTION_RESULTS_GENERIC_METRICS = "generic-metrics-subscription-results"
 
     QUERYLOG = "snuba-queries"
+
     PROFILES = "processed-profiles"
     PROFILES_FUNCTIONS = "profiles-call-tree"
+    PROFILE_CHUNKS = "snuba-profile-chunks"
+
     REPLAYEVENTS = "ingest-replay-events"
     GENERIC_METRICS = "snuba-generic-metrics"
     GENERIC_METRICS_SETS_COMMIT_LOG = "snuba-generic-metrics-sets-commit-log"