rust(features): add new flow after stream initialization + stream init from grpc channel (#229)

Author: solidiquis

rust/crates/sift_connect/src/grpc/config.rs

@@ -8,13 +8,13 @@ pub const SIFT_CONFIG_NAME: &str = "sift.toml";
 /// query the corresponding table from [`SIFT_CONFIG_NAME`] located at
 /// [these locations](https://docs.rs/dirs/6.0.0/dirs/fn.config_local_dir.html)
 /// depending on your operating system. If `None` is provided, then the top-level table is used.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum Credentials {
     Profile(Option<String>),
     Config { uri: String, apikey: String },
 }
 
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub(crate) struct SiftChannelConfig {
     pub uri: String,
     pub apikey: String,

rust/crates/sift_error/src/lib.rs

@@ -106,6 +106,8 @@ pub enum ErrorKind {
     GrpcConnectError,
     /// Indicates that the program was unable to retrieve the run being requested.
     RetrieveRunError,
+    /// Indicates that the program was unable to retrieve the asset being requested.
+    RetrieveAssetError,
     /// Indicates a failure to update a run.
     UpdateRunError,
     /// Indicates that the program was unable to retrieve the ingestion config being requested.
@@ -190,6 +192,7 @@ impl fmt::Display for ErrorKind {
         match self {
             Self::GrpcConnectError => write!(f, "GrpcConnectError"),
             Self::RetriesExhausted => write!(f, "RetriesExhausted"),
+            Self::RetrieveAssetError => write!(f, "RetrieveAssetError"),
             Self::RetrieveRunError => write!(f, "RetrieveRunError"),
             Self::RetrieveIngestionConfigError => write!(f, "RetrieveIngestionConfigError"),
             Self::EmptyResponseError => write!(f, "EmptyResponseError"),

rust/crates/sift_rs/src/wrappers/assets.rs

@@ -0,0 +1,54 @@
+use std::ops::{Deref, DerefMut};
+
+use async_trait::async_trait;
+use sift_connect::SiftChannel;
+use sift_error::prelude::*;
+
+use crate::assets::v1::{Asset, GetAssetRequest, asset_service_client::AssetServiceClient};
+
+/// Return an implementation of [AssetServiceWrapper] which also exposes methods from the
+/// raw [AssetServiceClient].
+pub fn new_asset_service(grpc_channel: SiftChannel) -> impl AssetServiceWrapper {
+    AssetServiceImpl(AssetServiceClient::new(grpc_channel))
+}
+
+/// Convenience methods
+#[async_trait]
+pub trait AssetServiceWrapper: Deref<Target = AssetServiceClient<SiftChannel>> + DerefMut {
+    /// Retrieves an asset by ID
+    async fn try_get_asset_by_id(&mut self, asset_id: &str) -> Result<Asset>;
+}
+
+/// A convenience wrapper around [AssetServiceClient].
+struct AssetServiceImpl(AssetServiceClient<SiftChannel>);
+
+#[async_trait]
+impl AssetServiceWrapper for AssetServiceImpl {
+    async fn try_get_asset_by_id(&mut self, asset_id: &str) -> Result<Asset> {
+        let req = GetAssetRequest {
+            asset_id: asset_id.into(),
+        };
+        let resp = self
+            .get_asset(req)
+            .await
+            .map_err(|e| Error::new(ErrorKind::RetrieveAssetError, e))?;
+
+        resp.into_inner().asset.ok_or_else(|| {
+            Error::new_empty_response("unexpected empty response from AssetService/GetAsset")
+        })
+    }
+}
+
+impl Deref for AssetServiceImpl {
+    type Target = AssetServiceClient<SiftChannel>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for AssetServiceImpl {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

rust/crates/sift_rs/src/wrappers/ingestion_configs.rs

@@ -78,11 +78,6 @@ impl IngestionConfigServiceWrapper for IngestionConfigServiceImpl {
                 "ingestion config client key cannot be blank",
             ));
         }
-        if flows.is_empty() {
-            return Err(Error::new_arg_error(
-                "cannot create an ingestion config with no flows",
-            ));
-        }
 
         self.create_ingestion_config(CreateIngestionConfigRequest {
             asset_name: asset_name.to_string(),

rust/crates/sift_rs/src/wrappers/mod.rs

@@ -1,3 +1,6 @@
+/// Offers a wrapper over Sift's assets API.
+pub mod assets;
+
 /// Offers a wrapper over Sift's ingestion configs API.
 pub mod ingestion_configs;
 

rust/crates/sift_stream/src/backup/disk/mod.rs

@@ -87,7 +87,7 @@ where
         };
         let backups_dir = backups_root.join(new_dir_name);
 
-        match fs::create_dir(&backups_dir) {
+        match fs::create_dir_all(&backups_dir) {
             Err(err) if err.kind() != IoErrorKind::AlreadyExists => {
                 return Err(Error::new(ErrorKind::BackupsError, err))
                     .with_context(|| format!("failed to create directory for backups at {}", backups_dir.display()))

rust/crates/sift_stream/src/stream/builder.rs

@@ -11,6 +11,7 @@ use sift_rs::{
     ping::v1::{PingRequest, ping_service_client::PingServiceClient},
     runs::v2::Run,
     wrappers::{
+        assets::{AssetServiceWrapper, new_asset_service},
         ingestion_configs::{IngestionConfigServiceWrapper, new_ingestion_config_service},
         runs::{RunServiceWrapper, new_run_service},
     },
@@ -39,7 +40,8 @@ pub const DEFAULT_CHECKPOINT_INTERVAL: Duration = Duration::from_secs(60);
 /// [tokio](https://docs.rs/tokio/latest/tokio/) asynchronous runtime is required, otherwise
 /// attempts to call [SiftStreamBuilder::build] will panic.
 pub struct SiftStreamBuilder<C> {
-    credentials: Credentials,
+    credentials: Option<Credentials>,
+    channel: Option<SiftChannel>,
     recovery_strategy: Option<RecoveryStrategy>,
     checkpoint_interval: Duration,
     ingestion_config: Option<IngestionConfigForm>,
@@ -171,7 +173,8 @@ where
         self
     }
 
-    /// Disables TLS. Useful for testing.
+    /// Disables TLS. Useful for testing. This is ignored if [SiftStreamBuilder::from_channel] is
+    /// used to initialize the builder.
     pub fn disable_tls(mut self) -> SiftStreamBuilder<C> {
         self.enable_tls = false;
         self
@@ -286,10 +289,26 @@ where
 
 /// Builds a [SiftStream] specifically for ingestion-config based streaming.
 impl SiftStreamBuilder<IngestionConfigMode> {
-    /// Initializes a new builder for ingestion-config-based streaming.
+    /// Initializes a new builder for ingestion-config-based streaming from [Credentials].
     pub fn new(credentials: Credentials) -> SiftStreamBuilder<IngestionConfigMode> {
         SiftStreamBuilder {
-            credentials,
+            credentials: Some(credentials),
+            channel: None,
+            enable_tls: true,
+            ingestion_config: None,
+            run: None,
+            run_id: None,
+            kind: PhantomData,
+            checkpoint_interval: DEFAULT_CHECKPOINT_INTERVAL,
+            recovery_strategy: None,
+        }
+    }
+
+    /// Initializes a new builder for ingestion-config-based streaming from a [SiftChannel].
+    pub fn from_channel(channel: SiftChannel) -> SiftStreamBuilder<IngestionConfigMode> {
+        SiftStreamBuilder {
+            credentials: None,
+            channel: Some(channel),
             enable_tls: true,
             ingestion_config: None,
             run: None,
@@ -304,26 +323,37 @@ impl SiftStreamBuilder<IngestionConfigMode> {
     /// streaming.
     pub async fn build(self) -> Result<SiftStream<IngestionConfigMode>> {
         let SiftStreamBuilder {
-            credentials,
             checkpoint_interval,
+            channel: grpc_channel,
+            credentials,
+            enable_tls,
             ingestion_config,
+            recovery_strategy,
             run,
             run_id,
-            recovery_strategy,
-            enable_tls,
             ..
         } = self;
 
         let Some(ingestion_config) = ingestion_config else {
             return Err(Error::new_arg_error("ingestion_config is required"));
         };
 
-        let mut sift_channel_builder = SiftChannelBuilder::new(credentials);
+        let channel = match grpc_channel {
+            Some(ch) => ch,
+            None if credentials.is_some() => {
+                let mut sift_channel_builder = SiftChannelBuilder::new(credentials.unwrap());
 
-        if enable_tls {
-            sift_channel_builder = sift_channel_builder.use_tls(true);
-        }
-        let channel = sift_channel_builder.build()?;
+                if enable_tls {
+                    sift_channel_builder = sift_channel_builder.use_tls(true);
+                }
+                sift_channel_builder.build()?
+            }
+            None => {
+                return Err(Error::new_arg_error(
+                    "either credentials or a gRPC channel must be provided",
+                ));
+            }
+        };
 
         // Since the gRPC connection is lazy, we'll connect right away and ensure the connection is
         // valid.
@@ -420,7 +450,8 @@ impl SiftStreamBuilder<IngestionConfigMode> {
         #[cfg(feature = "tracing")]
         tracing::info_span!("load_ingestion_config");
 
-        let mut ingestion_config_service = new_ingestion_config_service(grpc_channel);
+        let mut ingestion_config_service = new_ingestion_config_service(grpc_channel.clone());
+        let mut asset_service = new_asset_service(grpc_channel);
 
         let IngestionConfigForm {
             asset_name,
@@ -437,24 +468,31 @@ impl SiftStreamBuilder<IngestionConfigMode> {
                     .try_create_ingestion_config(&asset_name, &client_key, &flows)
                     .await?;
 
-                let flows = ingestion_config_service
-                    .try_filter_flows(&ingestion_config.ingestion_config_id, "")
-                    .await?;
+                let new_flows = {
+                    if flows.is_empty() {
+                        Vec::new()
+                    } else {
+                        ingestion_config_service
+                            .try_filter_flows(&ingestion_config.ingestion_config_id, "")
+                            .await?
+                    }
+                };
 
                 #[cfg(feature = "tracing")]
                 {
-                    let flow_names = flows
-                        .iter()
-                        .map(|f| f.name.as_str())
-                        .collect::<Vec<&str>>()
-                        .join(",");
-                    tracing::info!(
-                        ingestion_config_id = ingestion_config.ingestion_config_id,
-                        flow_names = flow_names,
-                        "created new ingestion config"
-                    );
+                    if !new_flows.is_empty() {
+                        let flow_names = new_flows
+                            .iter()
+                            .map(|f| f.name.as_str())
+                            .collect::<Vec<&str>>()
+                            .join(",");
+                        tracing::info!(
+                            ingestion_config_id = ingestion_config.ingestion_config_id,
+                            flow_names = flow_names,
+                            "created new ingestion config"
+                        );
+                    }
                 }
-
                 Ok((ingestion_config, flows))
             }
             Err(err) => Err(err),
@@ -466,13 +504,32 @@ impl SiftStreamBuilder<IngestionConfigMode> {
                     "an existing ingestion config was found with the provided client-key"
                 );
 
+                let asset = asset_service
+                    .try_get_asset_by_id(&ingestion_config.asset_id)
+                    .await
+                    .context("failed to retrieve asset specified by ingestion config")?;
+
+                if asset.name != asset_name {
+                    return Err(Error::new_msg(
+                        ErrorKind::IncompatibleIngestionConfigChange,
+                        format!(
+                            "local ingestion config references asset '{asset_name}' but this existing config in Sift refers to asset '{}'",
+                            asset.name
+                        ),
+                    ));
+                }
+
                 let flow_names = flows
                     .iter()
                     .map(|f| format!("'{}'", f.name))
                     .collect::<Vec<String>>()
                     .join(",");
 
-                let filter = format!("flow_name in [{flow_names}]");
+                let filter = flow_names
+                    .is_empty()
+                    .then(String::new)
+                    .unwrap_or_else(|| format!("flow_name in [{flow_names}]"));
+
                 let existing_flows = ingestion_config_service
                     .try_filter_flows(&ingestion_config.ingestion_config_id, &filter)
                     .await?;
@@ -534,3 +591,15 @@ impl SiftStreamBuilder<IngestionConfigMode> {
         }
     }
 }
+
+impl From<Credentials> for SiftStreamBuilder<IngestionConfigMode> {
+    fn from(value: Credentials) -> Self {
+        Self::new(value)
+    }
+}
+
+impl From<SiftChannel> for SiftStreamBuilder<IngestionConfigMode> {
+    fn from(value: SiftChannel) -> Self {
+        Self::from_channel(value)
+    }
+}

rust/crates/sift_stream/src/stream/mode/ingestion_config.rs

@@ -13,6 +13,7 @@ use sift_rs::{
     },
     ingestion_configs::v2::{FlowConfig, IngestionConfig},
     runs::v2::Run,
+    wrappers::ingestion_configs::{IngestionConfigServiceWrapper, new_ingestion_config_service},
 };
 use std::{
     collections::HashMap,
@@ -298,6 +299,29 @@ impl SiftStream<IngestionConfigMode> {
         }
     }
 
+    /// Modify the existing ingestion config by adding new flows that weren't accounted for during
+    /// initialization.
+    pub async fn add_new_flows(&mut self, flow_configs: &[FlowConfig]) -> Result<()> {
+        new_ingestion_config_service(self.grpc_channel.clone())
+            .try_create_flows(
+                &self.mode.ingestion_config.ingestion_config_id,
+                flow_configs,
+            )
+            .await
+            .context("SiftStream::add_new_flows")?;
+
+        for flow_config in flow_configs {
+            self.mode
+                .flows_by_name
+                .entry(flow_config.name.clone())
+                .and_modify(|flows| flows.push(flow_config.clone()))
+                .or_insert_with(|| vec![flow_config.clone()]);
+
+            tracing::info!(flow = flow_config.name, "successfully registered new flow");
+        }
+        Ok(())
+    }
+
     async fn backup_data(&mut self, req: &IngestWithConfigDataStreamRequest) -> Result<()> {
         if let Some(backups_manager) = self.mode.backups_manager.as_mut() {
             match backups_manager {
@@ -709,7 +733,8 @@ impl SiftStream<IngestionConfigMode> {
         })
     }
 
-    /// Flows passed into this function should have names match `flow_name`.
+    /// Flows passed into this function should have names match `flow_name`. The only case
+    /// in which this returns `None` is if there is no [FlowConfig] for the given `message`.
     pub(crate) fn message_to_ingest_req(
         message: &Flow,
         ingestion_config_id: &str,