database migration for backfilling the webhook log sink HMAC secret (#43305)

GitOrigin-RevId: 932731797091e8ab4221308e91adc71e8dd5df34

Author: reeceyang

crates/migrations_model/src/lib.rs

@@ -24,13 +24,14 @@ use value::{
 
 pub mod migr_119;
 pub mod migr_121;
+pub mod migr_124;
 
 pub type DatabaseVersion = i64;
 // The version for the format of the database. We support all previous
 // migrations unless explicitly dropping support.
 // Add a user name next to the version when you make a change to highlight merge
 // conflicts.
-pub const DATABASE_VERSION: DatabaseVersion = 123; // emma
+pub const DATABASE_VERSION: DatabaseVersion = 124; // reece
 
 pub struct MigrationExecutor<RT: Runtime> {
     pub db: Database<RT>,
@@ -86,6 +87,14 @@ impl<RT: Runtime> MigrationExecutor<RT> {
                 // table for each component, _scheduled_job_args
                 MigrationCompletionCriterion::MigrationComplete(to_version)
             },
+            124 => {
+                let mut tx = self.db.begin_system().await?;
+                migr_124::run_migration(&mut tx).await?;
+                self.db
+                    .commit_with_write_source(tx, "migration_124")
+                    .await?;
+                MigrationCompletionCriterion::MigrationComplete(to_version)
+            },
             // NOTE: Make sure to increase DATABASE_VERSION when adding new migrations.
             _ => anyhow::bail!("Version did not define a migration! {}", to_version),
         };

crates/migrations_model/src/migr_124/log_sinks/mod.rs

@@ -0,0 +1,100 @@
+use std::sync::LazyLock;
+
+use common::{
+    self,
+    document::ParsedDocument,
+    runtime::Runtime,
+};
+use database::{
+    system_tables::{
+        SystemIndex,
+        SystemTable,
+    },
+    Transaction,
+};
+use value::{
+    TableName,
+    TableNamespace,
+};
+
+pub mod types;
+use types::{
+    LogSinksRow,
+    SinkState,
+    SinkType,
+};
+
+pub static LOG_SINKS_TABLE: LazyLock<TableName> = LazyLock::new(|| {
+    "_log_sinks"
+        .parse()
+        .expect("Invalid built-in _log_sinks table")
+});
+
+pub struct LogSinksTable;
+impl SystemTable for LogSinksTable {
+    type Metadata = LogSinksRow;
+
+    fn table_name() -> &'static TableName {
+        &LOG_SINKS_TABLE
+    }
+
+    fn indexes() -> Vec<SystemIndex<Self>> {
+        vec![]
+    }
+}
+
+pub struct LogSinksModel<'a, RT: Runtime> {
+    tx: &'a mut Transaction<RT>,
+}
+
+impl<'a, RT: Runtime> LogSinksModel<'a, RT> {
+    pub fn new(tx: &'a mut Transaction<RT>) -> Self {
+        Self { tx }
+    }
+
+    pub async fn get_by_provider(
+        &mut self,
+        provider: SinkType,
+    ) -> anyhow::Result<Option<ParsedDocument<LogSinksRow>>> {
+        let mut result: Vec<_> = self
+            .get_by_provider_including_tombstoned(provider.clone())
+            .await?
+            .into_iter()
+            .filter(|doc| doc.status != SinkState::Tombstoned)
+            .collect();
+        anyhow::ensure!(
+            result.len() <= 1,
+            "Multiple sinks found of the same type: {:?}",
+            provider
+        );
+        Ok(result.pop())
+    }
+
+    async fn get_by_provider_including_tombstoned(
+        &mut self,
+        provider: SinkType,
+    ) -> anyhow::Result<Vec<ParsedDocument<LogSinksRow>>> {
+        let result: Vec<_> = self
+            .get_all()
+            .await?
+            .into_iter()
+            .filter(|doc| doc.config.sink_type() == provider)
+            .collect();
+        Ok(result)
+    }
+
+    pub async fn get_all(&mut self) -> anyhow::Result<Vec<ParsedDocument<LogSinksRow>>> {
+        let result = self
+            .tx
+            .query_system(
+                TableNamespace::Global,
+                &SystemIndex::<LogSinksTable>::by_id(),
+            )?
+            .all()
+            .await?
+            .into_iter()
+            .map(|arc_row| (*arc_row).clone())
+            .collect();
+        Ok(result)
+    }
+}

crates/migrations_model/src/migr_124/log_sinks/types/axiom.rs

@@ -0,0 +1,105 @@
+use std::{
+    fmt,
+    str::FromStr,
+};
+
+use common::{
+    log_streaming::LogEventFormatVersion,
+    pii::PII,
+};
+use serde::{
+    Deserialize,
+    Serialize,
+};
+
+#[derive(Debug, Clone, PartialEq)]
+#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
+pub struct AxiomConfig {
+    pub api_key: PII<String>,
+    pub dataset_name: String,
+    pub attributes: Vec<AxiomAttribute>,
+    pub version: LogEventFormatVersion,
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SerializedAxiomConfig {
+    pub api_key: String,
+    pub dataset_name: String,
+    pub attributes: Vec<SerializedAxiomAttribute>,
+    pub version: Option<String>,
+}
+
+impl From<AxiomConfig> for SerializedAxiomConfig {
+    fn from(value: AxiomConfig) -> Self {
+        Self {
+            api_key: value.api_key.0,
+            dataset_name: value.dataset_name,
+            attributes: value
+                .attributes
+                .into_iter()
+                .map(SerializedAxiomAttribute::from)
+                .collect(),
+            version: Some(value.version.to_string()),
+        }
+    }
+}
+
+impl TryFrom<SerializedAxiomConfig> for AxiomConfig {
+    type Error = anyhow::Error;
+
+    fn try_from(value: SerializedAxiomConfig) -> Result<Self, Self::Error> {
+        Ok(Self {
+            api_key: PII(value.api_key),
+            dataset_name: value.dataset_name,
+            attributes: value
+                .attributes
+                .into_iter()
+                .map(AxiomAttribute::from)
+                .collect(),
+            version: value
+                .version
+                .map(|v| LogEventFormatVersion::from_str(v.as_str()))
+                .transpose()?
+                .unwrap_or(LogEventFormatVersion::V1),
+        })
+    }
+}
+
+impl fmt::Display for AxiomConfig {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "AxiomConfig {{ version:{:?} ... }}", self.version)
+    }
+}
+
+#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
+pub struct AxiomAttribute {
+    pub key: String,
+    pub value: String,
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SerializedAxiomAttribute {
+    pub key: String,
+    pub value: String,
+}
+
+impl From<AxiomAttribute> for SerializedAxiomAttribute {
+    fn from(attribute: AxiomAttribute) -> Self {
+        Self {
+            key: attribute.key,
+            value: attribute.value,
+        }
+    }
+}
+
+impl From<SerializedAxiomAttribute> for AxiomAttribute {
+    fn from(value: SerializedAxiomAttribute) -> Self {
+        Self {
+            key: value.key,
+            value: value.value,
+        }
+    }
+}