Open source log sinks model (#34658)

On the path to open sourcing all the models and unifying the model code once again.

GitOrigin-RevId: 65bbb96d6a1ad49649f37816c1cbf1a84647cdf6

Author: nipunn1313

Cargo.lock

@@ -34,11 +34,14 @@ openidconnect = { workspace = true }
 pb = { path = "../pb" }
 proptest = { workspace = true, optional = true }
 proptest-derive = { workspace = true, optional = true }
+proptest-http = { workspace = true, optional = true }
 rand = { workspace = true }
+reqwest = { workspace = true }
 runtime = { path = "../runtime" }
 saffron = { workspace = true }
 search = { path = "../search" }
 semver = { workspace = true }
+sentry = { workspace = true }
 serde = { workspace = true }
 serde_bytes = { workspace = true }
 serde_json = { workspace = true }
@@ -59,6 +62,7 @@ keybroker = { path = "../keybroker", features = ["testing"] }
 metrics = { path = "../metrics", features = ["testing"] }
 proptest = { workspace = true }
 proptest-derive = { workspace = true }
+proptest-http = { workspace = true }
 runtime = { path = "../runtime", features = ["testing"] }
 search = { path = "../search", features = ["testing"] }
 storage = { path = "../storage", features = ["testing"] }
@@ -77,6 +81,7 @@ testing = [
     "metrics/testing",
     "proptest",
     "proptest-derive",
+    "proptest-http",
     "runtime/testing",
     "search/testing",
     "storage/testing",

crates/model/Cargo.toml

@@ -96,6 +96,7 @@ use environment_variables::{
 use exports::EXPORTS_BY_STATE_AND_TS_INDEX;
 use file_storage::FILE_STORAGE_ID_INDEX;
 use keybroker::Identity;
+use log_sinks::LogSinksTable;
 use maplit::btreeset;
 use modules::{
     MODULES_TABLE,
@@ -152,6 +153,7 @@ pub mod exports;
 pub mod external_packages;
 pub mod file_storage;
 pub mod fivetran_import;
+pub mod log_sinks;
 mod metrics;
 pub mod migrations;
 pub mod modules;
@@ -184,6 +186,7 @@ enum DefaultTableNumber {
     CronJobs = 19,
     Schemas = 20,
     CronJobLogs = 21,
+    LogSinks = 23,
     BackendState = 24,
     ExternalPackages = 25,
     ScheduledJobs = 27,
@@ -224,6 +227,7 @@ impl From<DefaultTableNumber> for &'static dyn SystemTable {
             DefaultTableNumber::CronJobs => &CronJobsTable,
             DefaultTableNumber::Schemas => &SchemasTable,
             DefaultTableNumber::CronJobLogs => &CronJobLogsTable,
+            DefaultTableNumber::LogSinks => &LogSinksTable,
             DefaultTableNumber::BackendState => &BackendStateTable,
             DefaultTableNumber::ExternalPackages => &ExternalPackagesTable,
             DefaultTableNumber::ScheduledJobs => &ScheduledJobsTable,
@@ -435,6 +439,7 @@ pub fn app_system_tables() -> Vec<&'static dyn SystemTable> {
         &SnapshotImportsTable,
         &FunctionHandlesTable,
         &CanonicalUrlsTable,
+        &LogSinksTable,
     ];
     system_tables.extend(component_system_tables());
     system_tables

crates/model/src/lib.rs

@@ -0,0 +1,190 @@
+use std::sync::LazyLock;
+
+use common::{
+    document::{
+        ParsedDocument,
+        ResolvedDocument,
+    },
+    query::{
+        Order,
+        Query,
+    },
+    runtime::Runtime,
+};
+use database::{
+    patch_value,
+    ResolvedQuery,
+    SystemMetadataModel,
+    Transaction,
+};
+use errors::ErrorMetadata;
+use value::{
+    ConvexValue,
+    ResolvedDocumentId,
+    TableName,
+    TableNamespace,
+};
+
+use crate::{
+    SystemIndex,
+    SystemTable,
+};
+
+pub mod types;
+use types::{
+    LogSinksRow,
+    SinkConfig,
+    SinkState,
+    SinkType,
+    LOG_SINKS_LIMIT,
+};
+
+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 {
+    fn table_name(&self) -> &'static TableName {
+        &LOG_SINKS_TABLE
+    }
+
+    fn indexes(&self) -> Vec<SystemIndex> {
+        vec![]
+    }
+
+    fn validate_document(&self, document: ResolvedDocument) -> anyhow::Result<()> {
+        ParsedDocument::<LogSinksRow>::try_from(document).map(|_| ())
+    }
+}
+
+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 mut result: Vec<_> = vec![];
+
+        let value_query = Query::full_table_scan(LOG_SINKS_TABLE.clone(), Order::Asc);
+        let mut query_stream = ResolvedQuery::new(self.tx, TableNamespace::Global, value_query)?;
+        while let Some(doc) = query_stream.next(self.tx, None).await? {
+            let row: ParsedDocument<LogSinksRow> = doc.try_into()?;
+            result.push(row);
+        }
+
+        Ok(result)
+    }
+
+    pub async fn patch_status(
+        &mut self,
+        id: ResolvedDocumentId,
+        status: SinkState,
+    ) -> anyhow::Result<()> {
+        SystemMetadataModel::new_global(self.tx)
+            .patch(
+                id,
+                patch_value!("status" => Some(ConvexValue::Object(status.try_into()?)))?,
+            )
+            .await?;
+        Ok(())
+    }
+
+    pub async fn mark_for_removal(&mut self, id: ResolvedDocumentId) -> anyhow::Result<()> {
+        self.patch_status(id, SinkState::Tombstoned).await?;
+        Ok(())
+    }
+
+    pub async fn add_or_update(&mut self, config: SinkConfig) -> anyhow::Result<()> {
+        let sink_type = config.sink_type();
+        let row = LogSinksRow {
+            status: SinkState::Pending,
+            config,
+        };
+
+        // Filter to non-tombstoned log sinks
+        let sinks = self
+            .get_all()
+            .await?
+            .into_iter()
+            .filter(|row| row.status != SinkState::Tombstoned)
+            .collect::<Vec<_>>();
+        if sinks.len() >= LOG_SINKS_LIMIT {
+            return Err(ErrorMetadata::bad_request(
+                "LogSinkQuotaExceeded",
+                "Cannot add more LogSinks, the quota for this project has been reached.",
+            )
+            .into());
+        }
+
+        if let Some(row) = self.get_by_provider(sink_type.clone()).await? {
+            self.mark_for_removal(row.id()).await?;
+        }
+
+        SystemMetadataModel::new_global(self.tx)
+            .insert(&LOG_SINKS_TABLE, row.try_into()?)
+            .await?;
+        Ok(())
+    }
+
+    // It's generally not safe to delete an existing sink without marking it
+    // Tombstoned first since the LogManager will not know to remove the sink.
+    // However, we can do this during startup before the LogManager has started
+    // (like when adding a local log sink)
+    pub async fn add_on_startup(&mut self, config: SinkConfig) -> anyhow::Result<()> {
+        // Search for matching provider
+        if let Some(sink) = self.get_by_provider(config.sink_type()).await? {
+            SystemMetadataModel::new_global(self.tx)
+                .delete(sink.id())
+                .await?;
+        };
+        self.add_or_update(config).await?;
+        Ok(())
+    }
+
+    pub async fn clear(&mut self) -> anyhow::Result<()> {
+        let providers = self.get_all().await?;
+
+        for sink in providers {
+            self.patch_status(sink.id(), SinkState::Tombstoned).await?;
+        }
+        Ok(())
+    }
+}

crates/model/src/log_sinks/mod.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,
+        }
+    }
+}