Move logstreaming handlers to open source/self-hosted backend binary. (#38394)

Add handlers/code to the backend. Still need to update the embedded dashboard to include it too.

GitOrigin-RevId: deacbd2a94b58934b6773b811b68fb2335d0102b

Author: nipunn1313

Cargo.lock

@@ -17,6 +17,7 @@ testing = [
     "function_runner/testing",
     "isolate/testing",
     "keybroker/testing",
+    "log_streaming/testing",
     "metrics/testing",
     "model/testing",
     "node_executor/testing",
@@ -67,6 +68,7 @@ humansize = { workspace = true }
 isolate = { path = "../isolate" }
 itertools = { workspace = true }
 keybroker = { path = "../keybroker" }
+log_streaming = { path = "../log_streaming" }
 lru = { workspace = true }
 maplit = { workspace = true }
 metrics = { path = "../metrics" }
@@ -121,6 +123,7 @@ events = { path = "../events", features = ["testing"] }
 function_runner = { path = "../function_runner", features = ["testing"] }
 isolate = { path = "../isolate", features = ["testing"] }
 keybroker = { path = "../keybroker", features = ["testing"] }
+log_streaming = { path = "../log_streaming", features = ["testing"] }
 metrics = { path = "../metrics", features = ["testing"] }
 model = { path = "../model", features = ["testing"] }
 must-let = { workspace = true }

crates/application/Cargo.toml

@@ -21,6 +21,10 @@ use std::{
     },
 };
 
+use ::log_streaming::{
+    LogManager,
+    LogManagerClient,
+};
 use airbyte_import::{
     AirbyteRecord,
     PrimaryKey,
@@ -74,7 +78,10 @@ use common::{
         report_error,
         JsError,
     },
-    http::RequestDestination,
+    http::{
+        fetch::FetchClient,
+        RequestDestination,
+    },
     knobs::{
         APPLICATION_MAX_CONCURRENT_UPLOADS,
         MAX_JOBS_CANCEL_BATCH,
@@ -191,6 +198,7 @@ use keybroker::{
     Identity,
     KeyBroker,
 };
+use log_streaming::add_local_log_sink_on_startup;
 use maplit::{
     btreemap,
     btreeset,
@@ -201,6 +209,7 @@ use model::{
         AIRBYTE_PRIMARY_KEY_INDEX_DESCRIPTOR,
     },
     auth::AuthInfoModel,
+    backend_info::BackendInfoModel,
     backend_state::BackendStateModel,
     canonical_urls::{
         types::CanonicalUrl,
@@ -389,6 +398,7 @@ pub mod cron_jobs;
 pub mod deploy_config;
 mod exports;
 pub mod function_log;
+mod log_streaming;
 pub mod log_visibility;
 mod metrics;
 mod module_cache;
@@ -548,11 +558,11 @@ pub struct Application<RT: Runtime> {
     export_worker: Arc<Mutex<Box<dyn SpawnHandle>>>,
     system_table_cleanup_worker: Arc<Mutex<Box<dyn SpawnHandle>>>,
     migration_worker: Arc<Mutex<Option<Box<dyn SpawnHandle>>>>,
-    log_sender: Arc<dyn LogSender>,
     log_visibility: Arc<dyn LogVisibility<RT>>,
     module_cache: ModuleCache<RT>,
     system_env_var_names: HashSet<EnvVarName>,
     app_auth: Arc<ApplicationAuth>,
+    log_manager_client: LogManagerClient,
 }
 
 impl<RT: Runtime> Clone for Application<RT> {
@@ -579,11 +589,11 @@ impl<RT: Runtime> Clone for Application<RT> {
             export_worker: self.export_worker.clone(),
             system_table_cleanup_worker: self.system_table_cleanup_worker.clone(),
             migration_worker: self.migration_worker.clone(),
-            log_sender: self.log_sender.clone(),
             log_visibility: self.log_visibility.clone(),
             module_cache: self.module_cache.clone(),
             system_env_var_names: self.system_env_var_names.clone(),
             app_auth: self.app_auth.clone(),
+            log_manager_client: self.log_manager_client.clone(),
         }
     }
 }
@@ -671,10 +681,11 @@ impl<RT: Runtime> Application<RT> {
         segment_term_metadata_fetcher: Arc<dyn SegmentTermMetadataFetcher>,
         persistence: Arc<dyn Persistence>,
         node_actions: Actions<RT>,
-        log_sender: Arc<dyn LogSender>,
         log_visibility: Arc<dyn LogVisibility<RT>>,
         app_auth: Arc<ApplicationAuth>,
         cache: QueryCache,
+        fetch_client: Arc<dyn FetchClient>,
+        local_log_sink: Option<String>,
     ) -> anyhow::Result<Self> {
         let module_cache =
             ModuleCache::new(runtime.clone(), application_storage.modules_storage.clone()).await;
@@ -724,10 +735,32 @@ impl<RT: Runtime> Application<RT> {
             runtime.spawn("system_table_cleanup_worker", system_table_cleanup_worker),
         ));
 
+        // If local_log_sink is passed in, this is a local instance, so we enable log
+        // streaming by default. Otherwise, it's hard to grant the
+        // entitlement in testing and in load generator. If not local, we
+        // read the entitlement from the database.
+        let mut tx = database.begin(Identity::system()).await?;
+        let log_streaming_allowed = if let Some(path) = local_log_sink {
+            add_local_log_sink_on_startup(database.clone(), path).await?;
+            true
+        } else {
+            let mut bi = BackendInfoModel::new(&mut tx);
+            bi.is_log_streaming_allowed().await?
+        };
+
+        let log_manager_client = LogManager::start(
+            runtime.clone(),
+            database.clone(),
+            fetch_client.clone(),
+            instance_name.clone(),
+            log_streaming_allowed,
+        )
+        .await;
+
         let function_log = FunctionExecutionLog::new(
             runtime.clone(),
             database.usage_counter(),
-            log_sender.clone(),
+            Arc::new(log_manager_client.clone()),
         );
         let runner = Arc::new(ApplicationFunctionRunner::new(
             runtime.clone(),
@@ -816,11 +849,11 @@ impl<RT: Runtime> Application<RT> {
             snapshot_import_worker,
             system_table_cleanup_worker,
             migration_worker,
-            log_sender,
             log_visibility,
             module_cache,
             system_env_var_names: default_system_env_vars.into_keys().collect(),
             app_auth,
+            log_manager_client,
         })
     }
 
@@ -848,6 +881,10 @@ impl<RT: Runtime> Application<RT> {
         self.function_log.clone()
     }
 
+    pub fn log_manager_client(&self) -> &LogManagerClient {
+        &self.log_manager_client
+    }
+
     pub fn now_ts_for_reads(&self) -> RepeatableTimestamp {
         self.database.now_ts_for_reads()
     }
@@ -2990,7 +3027,7 @@ impl<RT: Runtime> Application<RT> {
             })
             .try_collect()?;
 
-        self.log_sender.send_logs(logs);
+        self.log_manager_client.send_logs(logs);
         Ok(ts)
     }
 
@@ -3088,7 +3125,7 @@ impl<RT: Runtime> Application<RT> {
             })
             .try_collect()?;
 
-        self.log_sender.send_logs(logs);
+        self.log_manager_client.send_logs(logs);
         Ok((t, stats))
     }
 
@@ -3439,7 +3476,7 @@ impl<RT: Runtime> Application<RT> {
     }
 
     pub async fn shutdown(&self) -> anyhow::Result<()> {
-        self.log_sender.shutdown()?;
+        self.log_manager_client.shutdown()?;
         self.table_summary_worker.shutdown().await?;
         self.system_table_cleanup_worker.lock().shutdown();
         self.schema_worker.lock().shutdown();

crates/application/src/lib.rs

@@ -0,0 +1,66 @@
+use common::runtime::Runtime;
+use database::Database;
+use errors::ErrorMetadata;
+use keybroker::Identity;
+use model::{
+    backend_info::BackendInfoModel,
+    log_sinks::{
+        types::{
+            SinkConfig,
+            SinkType,
+        },
+        LogSinksModel,
+    },
+};
+
+use crate::Application;
+
+pub async fn add_local_log_sink_on_startup<RT: Runtime>(
+    db: Database<RT>,
+    path: String,
+) -> anyhow::Result<()> {
+    let mut tx = db.begin(Identity::system()).await?;
+    let mut log_sink_model = LogSinksModel::new(&mut tx);
+    log_sink_model
+        .add_on_startup(SinkConfig::Local(path.clone()))
+        .await?;
+    db.commit_with_write_source(tx, "add_local_log_sink_startup")
+        .await?;
+    tracing::info!("Local log sink configured at {path}.");
+    Ok(())
+}
+
+impl<RT: Runtime> Application<RT> {
+    pub async fn add_log_sink(&self, config: SinkConfig) -> anyhow::Result<()> {
+        let mut tx = self.begin(Identity::system()).await?;
+        let mut model = LogSinksModel::new(&mut tx);
+        model.add_or_update(config).await?;
+        self.commit(tx, "add_log_sink").await?;
+        Ok(())
+    }
+
+    pub async fn remove_log_sink(&self, sink_type: SinkType) -> anyhow::Result<()> {
+        let mut tx = self.begin(Identity::system()).await?;
+        let mut model = LogSinksModel::new(&mut tx);
+
+        let Some(row) = model.get_by_provider(sink_type.clone()).await? else {
+            return Err(ErrorMetadata::bad_request(
+                "SinkDoesntExist",
+                "Cannot remove a sink that is not configured for this project.",
+            )
+            .into());
+        };
+
+        model.mark_for_removal(row.id()).await?;
+        self.commit(tx, "remove_log_sink").await?;
+
+        Ok(())
+    }
+
+    pub async fn ensure_log_streaming_allowed(&self, identity: Identity) -> anyhow::Result<()> {
+        let mut tx = self.begin(identity).await?;
+        BackendInfoModel::new(&mut tx)
+            .ensure_log_streaming_allowed()
+            .await
+    }
+}

crates/application/src/log_streaming.rs

@@ -29,7 +29,6 @@ use common::{
         ACTION_USER_TIMEOUT,
         UDF_CACHE_MAX_SIZE,
     },
-    log_streaming::NoopLogSender,
     persistence::Persistence,
     runtime::Runtime,
     shutdown::ShutdownSignal,
@@ -227,7 +226,7 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
                     modules_storage: application_storage.modules_storage.clone(),
                 },
                 database.clone(),
-                fetch_client,
+                fetch_client.clone(),
             )
             .await?,
         );
@@ -266,13 +265,14 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
             segment_term_metadata_fetcher,
             Arc::new(persistence.clone()),
             actions,
-            Arc::new(NoopLogSender),
             Arc::new(RedactLogsToClient::new(false)),
             Arc::new(ApplicationAuth::new(
                 kb.clone(),
                 Arc::new(NullAccessTokenAuth),
             )),
             QueryCache::new(*UDF_CACHE_MAX_SIZE),
+            fetch_client,
+            None, // local_log_sink
         )
         .await?;