Add selection support to the Fivetran source connector (#37390)

This adds selection support to the Fivetran source component, meaning that users will be able to opt out of syncing some components, tables and columns.

GitOrigin-RevId: acb1ef09e81ad3aa1467ad35814d208f61e2ba97

Author: Nicolapps

Cargo.lock

@@ -51,6 +51,7 @@ tonic-build = { workspace = true }
 cmd_util = { workspace = true }
 convex = { workspace = true, features = ["testing"] }
 convex_fivetran_common = { workspace = true, features = ["testing"] }
+pretty_assertions = { workspace = true }
 proptest = { workspace = true }
 proptest-derive = { workspace = true }
 rand = { workspace = true }

crates/fivetran_source/Cargo.toml

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
 use convex_fivetran_common::fivetran_sdk::{
     selection::Selection as FivetranSelection,
     SchemaSelection as FivetranSchemaSelection,
+    Selection as FivetranRootSelection,
     TableSelection as FivetranTableSelection,
     TablesWithSchema as FivetranSelectionWithSchema,
 };
@@ -20,8 +21,8 @@ use super::SelectionArg;
 ///
 /// This is the serializable version of `StreamingExportSelection` in the
 /// database crate.
-#[derive(Serialize, Deserialize)]
-#[cfg_attr(test, derive(Eq, PartialEq, Debug, Clone, proptest_derive::Arbitrary))]
+#[derive(Serialize, Deserialize, Clone)]
+#[cfg_attr(test, derive(Eq, PartialEq, Debug, proptest_derive::Arbitrary))]
 pub struct Selection {
     #[serde(flatten)]
     #[cfg_attr(
@@ -58,8 +59,8 @@ impl Default for Selection {
 }
 
 /// Serializable version of `StreamingExportComponentSelection`.
-#[derive(Serialize, Deserialize)]
-#[cfg_attr(test, derive(Clone, Eq, PartialEq, Debug, proptest_derive::Arbitrary))]
+#[derive(Serialize, Deserialize, Clone)]
+#[cfg_attr(test, derive(Eq, PartialEq, Debug, proptest_derive::Arbitrary))]
 pub enum ComponentSelection {
     #[serde(rename = "excl")]
     Excluded,
@@ -79,8 +80,8 @@ pub enum ComponentSelection {
 
 /// Serializable version of
 /// `StreamingExportTableSelection` + `StreamingExportColumnSelection`
-#[derive(Serialize, Deserialize)]
-#[cfg_attr(test, derive(Clone, Eq, PartialEq, Debug, proptest_derive::Arbitrary))]
+#[derive(Serialize, Deserialize, Clone)]
+#[cfg_attr(test, derive(Eq, PartialEq, Debug, proptest_derive::Arbitrary))]
 pub enum TableSelection {
     #[serde(rename = "excl")]
     Excluded,
@@ -157,6 +158,14 @@ impl From<SelectionArg> for Selection {
 /// component)
 pub const DEFAULT_FIVETRAN_SCHEMA_NAME: &str = "convex";
 
+impl TryFrom<Option<FivetranRootSelection>> for Selection {
+    type Error = anyhow::Error;
+
+    fn try_from(value: Option<FivetranRootSelection>) -> Result<Self, Self::Error> {
+        Selection::try_from(value.and_then(|val| val.selection))
+    }
+}
+
 impl TryFrom<Option<FivetranSelection>> for Selection {
     type Error = anyhow::Error;
 
@@ -432,8 +441,9 @@ mod tests_selection_fivetran_conversion {
     use super::*;
 
     #[test]
-    fn test_schema_equals_none_converts_to_everything_included() {
-        let result: Result<Selection, _> = Selection::try_from(None);
+    fn test_none_converts_to_everything_included() {
+        let none: Option<FivetranSelection> = None;
+        let result: Result<Selection, _> = Selection::try_from(none);
         assert_eq!(
             result.unwrap(),
             Selection {

crates/fivetran_source/src/api_types/selection.rs

@@ -31,7 +31,10 @@ use tonic::{
 };
 
 use crate::{
-    api_types::selection::DEFAULT_FIVETRAN_SCHEMA_NAME,
+    api_types::selection::{
+        Selection,
+        DEFAULT_FIVETRAN_SCHEMA_NAME,
+    },
     convex_api::{
         ComponentPath,
         ConvexApi,
@@ -158,7 +161,10 @@ impl SourceConnector for ConvexConnector {
 
         let source = ConvexApi { config };
 
-        let sync = sync(source, state);
+        let selection = Selection::try_from(inner.selection)
+            .map_err(|error| Status::internal(error.to_string()))?;
+
+        let sync = sync(source, state, selection);
         Ok(Response::new(
             sync.map_ok(FivetranUpdateResponse::from)
                 .map_err(|error| Status::internal(error.to_string()))

crates/fivetran_source/src/connector.rs

@@ -23,7 +23,10 @@ use serde::{
 };
 
 use crate::api_types::{
-    selection::DEFAULT_FIVETRAN_SCHEMA_NAME,
+    selection::{
+        Selection,
+        DEFAULT_FIVETRAN_SCHEMA_NAME,
+    },
     DocumentDeltasArgs,
     DocumentDeltasResponse,
     DocumentDeltasValue,
@@ -53,12 +56,14 @@ pub trait Source: Display + Send {
         &self,
         snapshot: Option<i64>,
         cursor: Option<ListSnapshotCursor>,
+        selection: Selection,
     ) -> anyhow::Result<ListSnapshotResponse>;
 
     /// See https://docs.convex.dev/http-api/#get-apidocument_deltas
     async fn document_deltas(
         &self,
         cursor: DocumentDeltasCursor,
+        selection: Selection,
     ) -> anyhow::Result<DocumentDeltasResponse>;
 
     /// Get a list of columns for each table and component on the Convex
@@ -173,13 +178,14 @@ impl Source for ConvexApi {
         &self,
         snapshot: Option<i64>,
         cursor: Option<ListSnapshotCursor>,
+        selection: Selection,
     ) -> anyhow::Result<ListSnapshotResponse> {
         self.post(
             "list_snapshot",
             ListSnapshotArgs {
                 snapshot,
                 cursor: cursor.map(|c| c.into()),
-                selection: SelectionArg::default(),
+                selection: SelectionArg::Exact { selection },
                 format: Some("convex_encoded_json".to_string()),
             },
         )
@@ -189,12 +195,13 @@ impl Source for ConvexApi {
     async fn document_deltas(
         &self,
         cursor: DocumentDeltasCursor,
+        selection: Selection,
     ) -> anyhow::Result<DocumentDeltasResponse> {
         self.post(
             "document_deltas",
             DocumentDeltasArgs {
                 cursor: Some(cursor.into()),
-                selection: SelectionArg::default(),
+                selection: SelectionArg::Exact { selection },
                 format: Some("convex_encoded_json".to_string()),
             },
         )
@@ -242,7 +249,7 @@ pub struct ListSnapshotCursor(pub String);
 #[cfg_attr(test, derive(proptest_derive::Arbitrary))]
 pub struct DocumentDeltasCursor(pub i64);
 
-#[derive(Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash, Display)]
+#[derive(Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash, Display, Debug)]
 pub struct TableName(pub String);
 
 #[cfg(test)]
@@ -273,7 +280,7 @@ impl ComponentPath {
     }
 }
 
-#[derive(Display)]
+#[derive(Display, Debug)]
 pub struct FieldName(pub String);
 
 #[cfg(test)]

crates/fivetran_source/src/convex_api.rs

@@ -25,6 +25,7 @@ use serde::{
 use value_type::Inner as FivetranValue;
 
 use crate::{
+    api_types::selection::Selection,
     convert::to_fivetran_row,
     convex_api::{
         DocumentDeltasCursor,
@@ -90,6 +91,7 @@ pub enum Checkpoint {
 }
 
 /// A simplification of the messages sent to Fivetran in the `update` endpoint.
+#[derive(Debug)]
 pub enum UpdateMessage {
     Update {
         schema_name: Option<String>,
@@ -140,9 +142,10 @@ impl From<UpdateMessage> for FivetranUpdateResponse {
 pub fn sync(
     source: impl Source + 'static,
     state: Option<State>,
+    selection: Selection,
 ) -> BoxStream<'static, anyhow::Result<UpdateMessage>> {
     let Some(state) = state else {
-        return initial_sync(source, None, Some(BTreeSet::new())).boxed();
+        return initial_sync(source, None, Some(BTreeSet::new()), selection).boxed();
     };
 
     let State {
@@ -152,9 +155,11 @@ pub fn sync(
     } = state;
     match checkpoint {
         Checkpoint::InitialSync { snapshot, cursor } => {
-            initial_sync(source, Some((snapshot, cursor)), tables_seen).boxed()
+            initial_sync(source, Some((snapshot, cursor)), tables_seen, selection).boxed()
+        },
+        Checkpoint::DeltaUpdates { cursor } => {
+            delta_sync(source, cursor, tables_seen, selection).boxed()
         },
-        Checkpoint::DeltaUpdates { cursor } => delta_sync(source, cursor, tables_seen).boxed(),
     }
 }
 
@@ -164,6 +169,7 @@ async fn initial_sync(
     source: impl Source,
     mut checkpoint: Option<(i64, ListSnapshotCursor)>,
     mut tables_seen: Option<BTreeSet<String>>,
+    selection: Selection,
 ) {
     let log_msg = if let Some((snapshot, _)) = checkpoint {
         format!("Resuming an initial sync from {source} at {snapshot}")
@@ -175,7 +181,9 @@ async fn initial_sync(
     let snapshot = loop {
         let snapshot = checkpoint.as_ref().map(|c| c.0);
         let cursor = checkpoint.as_ref().map(|c| c.1.clone());
-        let res = source.list_snapshot(snapshot, cursor.clone()).await?;
+        let res = source
+            .list_snapshot(snapshot, cursor.clone(), selection.clone())
+            .await?;
 
         for value in res.values {
             if let Some(ref mut tables_seen) = tables_seen {
@@ -235,13 +243,14 @@ async fn delta_sync(
     source: impl Source,
     cursor: DocumentDeltasCursor,
     mut tables_seen: Option<BTreeSet<String>>,
+    selection: Selection,
 ) {
     log(&format!("Delta sync from {source} starting at {cursor}."));
 
     let mut cursor = cursor;
     let mut has_more = true;
     while has_more {
-        let response = source.document_deltas(cursor).await?;
+        let response = source.document_deltas(cursor, selection.clone()).await?;
 
         for value in response.values {
             if let Some(ref mut tables_seen) = tables_seen {