feat(source): introduce gcs source and new s3 source via OpenDAL

ci/scripts/s3-source-test.sh

@@ -29,7 +29,7 @@ echo "--- starting risingwave cluster with connector node"
 cargo make ci-start ci-1cn-1fe
 
 echo "--- Run test"
-python3 -m pip install minio psycopg2-binary
+python3 -m pip install minio psycopg2-binary opendal
 python3 e2e_test/s3/$script
 
 echo "--- Kill cluster"

ci/workflows/main-cron.yml

@@ -468,6 +468,23 @@ steps:
     timeout_in_minutes: 25
     retry: *auto-retry
 
+  - label: "GCS source check on AWS (json parser)"
+    command: "ci/scripts/s3-source-test.sh -p ci-release -s 'gcs_source.py json'"
+    depends_on: build
+    plugins:
+      - seek-oss/aws-sm#v2.3.1:
+          env:
+            GCS_SOURCE_TEST_CONF: ci_gcs_source_test_gcp
+      - docker-compose#v4.9.0:
+          run: rw-build-env
+          config: ci/docker-compose.yml
+          mount-buildkite-agent: true
+          environment:
+            - GCS_SOURCE_TEST_CONF
+      - ./ci/plugins/upload-failure-logs
+    timeout_in_minutes: 25
+    retry: *auto-retry
+
   - label: "S3_v2 source check on AWS (csv parser)"
     command: "ci/scripts/s3-source-test.sh -p ci-release -s 'fs_source_v2.py csv_without_header'"
     if: |

e2e_test/s3/gcs_source.py

@@ -0,0 +1,130 @@
+import os
+import sys
+import csv
+import json
+import random
+import psycopg2
+import opendal
+
+from time import sleep
+from io import StringIO
+from functools import partial
+
+def gen_data(file_num, item_num_per_file):
+    assert item_num_per_file % 2 == 0, \
+        f'item_num_per_file should be even to ensure sum(mark) == 0: {item_num_per_file}'
+    return [
+        [{
+            'id': file_id * item_num_per_file + item_id,
+            'name': f'{file_id}_{item_id}',
+            'sex': item_id % 2,
+            'mark': (-1) ** (item_id % 2),
+        } for item_id in range(item_num_per_file)]
+        for file_id in range(file_num)
+    ]
+
+def format_json(data):
+    return [
+        '\n'.join([json.dumps(item) for item in file])
+        for file in data
+    ]
+
+
+def do_test(config, file_num, item_num_per_file, prefix, fmt, credential):
+    conn = psycopg2.connect(
+        host="localhost",
+        port="4566",
+        user="root",
+        database="dev"
+    )
+
+    # Open a cursor to execute SQL statements
+    cur = conn.cursor()
+
+    def _table():
+        return f'gcs_test_{fmt}'
+
+    def _encode():
+        return 'JSON'
+
+    # Execute a SELECT statement
+    cur.execute(f'''CREATE TABLE {_table()}(
+        id int,
+        name TEXT,
+        sex int,
+        mark int,
+    ) WITH (
+        connector = 'gcs',
+        match_pattern = '{prefix}*.{fmt}',
+        gcs.bucket_name = '{config['GCS_BUCKET']}',
+        gcs.credentials = '{credential}',
+    ) FORMAT PLAIN ENCODE {_encode()};''')
+
+    total_rows = file_num * item_num_per_file
+    MAX_RETRIES = 40
+    for retry_no in range(MAX_RETRIES):
+        cur.execute(f'select count(*) from {_table()}')
+        result = cur.fetchone()
+        if result[0] == total_rows:
+            break
+        print(f"[retry {retry_no}] Now got {result[0]} rows in table, {total_rows} expected, wait 30s")
+        sleep(30)
+
+    stmt = f'select count(*), sum(id), sum(sex), sum(mark) from {_table()}'
+    print(f'Execute {stmt}')
+    cur.execute(stmt)
+    result = cur.fetchone()
+
+    print('Got:', result)
+
+    def _assert_eq(field, got, expect):
+        assert got == expect, f'{field} assertion failed: got {got}, expect {expect}.'
+
+    _assert_eq('count(*)', result[0], total_rows)
+    _assert_eq('sum(id)', result[1], (total_rows - 1) * total_rows / 2)
+    _assert_eq('sum(sex)', result[2], total_rows / 2)
+    _assert_eq('sum(mark)', result[3], 0)
+
+    print('Test pass')
+
+    cur.execute(f'drop table {_table()}')
+    cur.close()
+    conn.close()
+
+
+if __name__ == "__main__":
+    FILE_NUM = 4001
+    ITEM_NUM_PER_FILE = 2
+    data = gen_data(FILE_NUM, ITEM_NUM_PER_FILE)
+
+    fmt = sys.argv[1]
+    FORMATTER = {
+        'json': format_json,
+    }
+    assert fmt in FORMATTER, f"Unsupported format: {fmt}"
+    formatted_files = FORMATTER[fmt](data)
+
+    config = json.loads(os.environ["GCS_SOURCE_TEST_CONF"])
+    run_id = str(random.randint(1000, 9999))
+    _local = lambda idx: f'data_{idx}.{fmt}'
+    _gcs = lambda idx: f"{run_id}_data_{idx}.{fmt}"
+    credential_str = json.dumps(config["GOOGLE_APPLICATION_CREDENTIALS"])
+    # put gcs files
+    op = opendal.Operator("gcs", root="/", bucket=config["GCS_BUCKET"], credential=credential_str)
+
+    print("upload file to gcs")
+    for idx, file_str in enumerate(formatted_files):
+        with open(_local(idx), "w") as f:
+            f.write(file_str)
+            os.fsync(f.fileno())
+        file_bytes = file_str.encode('utf-8')
+        op.write(_gcs(idx), file_bytes)
+
+    # do test
+    print("do test")
+    do_test(config, FILE_NUM, ITEM_NUM_PER_FILE, run_id, fmt, credential_str)
+
+    # clean up gcs files
+    print("clean up gcs files")
+    for idx, _ in enumerate(formatted_files):
+        op.delete(_gcs(idx))

src/connector/Cargo.toml

@@ -78,6 +78,7 @@ mysql_common = { version = "0.31", default-features = false, features = [
 ] }
 nexmark = { version = "0.2", features = ["serde"] }
 num-bigint = "0.4"
+opendal = "0.41"
 parking_lot = "0.12"
 paste = "1"
 prometheus = { version = "0.13", features = ["process"] }

src/connector/src/macros.rs

@@ -32,6 +32,8 @@ macro_rules! for_all_classified_sources {
                 { GooglePubsub, $crate::source::google_pubsub::PubsubProperties, $crate::source::google_pubsub::PubsubSplit },
                 { Nats, $crate::source::nats::NatsProperties, $crate::source::nats::split::NatsSplit },
                 { S3, $crate::source::filesystem::S3Properties, $crate::source::filesystem::FsSplit },
+                { Gcs, $crate::source::filesystem::opendal_source::GcsProperties , $crate::source::filesystem::OpendalFsSplit<$crate::source::filesystem::opendal_source::GcsProperties> },
+                { OpenDalS3, $crate::source::filesystem::opendal_source::OpendalS3Properties, $crate::source::filesystem::OpendalFsSplit<$crate::source::filesystem::opendal_source::OpendalS3Properties> },
                 { Test, $crate::source::test_source::TestSourceProperties, $crate::source::test_source::TestSourceSplit}
             }
             $(

src/connector/src/source/base.rs

@@ -41,11 +41,12 @@ use super::google_pubsub::GooglePubsubMeta;
 use super::kafka::KafkaMeta;
 use super::monitor::SourceMetrics;
 use super::nexmark::source::message::NexmarkMeta;
+use super::{GCS_CONNECTOR, OPENDAL_S3_CONNECTOR};
 use crate::parser::ParserConfig;
 pub(crate) use crate::source::common::CommonSplitReader;
-use crate::source::filesystem::{FsPageItem, S3Properties, S3_V2_CONNECTOR};
+use crate::source::filesystem::opendal_source::OpendalS3Properties;
+use crate::source::filesystem::{FsPageItem, GcsProperties, S3Properties};
 use crate::source::monitor::EnumeratorMetrics;
-use crate::source::S3_CONNECTOR;
 use crate::{
     dispatch_source_prop, dispatch_split_impl, for_all_sources, impl_connector_properties,
     impl_split, match_source_name_str,
@@ -75,7 +76,7 @@ impl<P: DeserializeOwned> TryFromHashmap for P {
     }
 }
 
-pub async fn create_split_reader<P: SourceProperties>(
+pub async fn create_split_reader<P: SourceProperties + std::fmt::Debug>(
     prop: P,
     splits: Vec<SplitImpl>,
     parser_config: ParserConfig,
@@ -364,44 +365,51 @@ impl ConnectorProperties {
     pub fn is_new_fs_connector_b_tree_map(props: &BTreeMap<String, String>) -> bool {
         props
             .get(UPSTREAM_SOURCE_KEY)
-            .map(|s| s.eq_ignore_ascii_case(S3_V2_CONNECTOR))
+            .map(|s| s.eq_ignore_ascii_case(OPENDAL_S3_CONNECTOR))
             .unwrap_or(false)
+            || props
+                .get(UPSTREAM_SOURCE_KEY)
+                .map(|s| s.eq_ignore_ascii_case(GCS_CONNECTOR))
+                .unwrap_or(false)
     }
 
     pub fn is_new_fs_connector_hash_map(props: &HashMap<String, String>) -> bool {
         props
             .get(UPSTREAM_SOURCE_KEY)
-            .map(|s| s.eq_ignore_ascii_case(S3_V2_CONNECTOR))
+            .map(|s| s.eq_ignore_ascii_case(OPENDAL_S3_CONNECTOR))
             .unwrap_or(false)
-    }
-
-    pub fn rewrite_upstream_source_key_hash_map(props: &mut HashMap<String, String>) {
-        let connector = props.remove(UPSTREAM_SOURCE_KEY).unwrap();
-        match connector.as_str() {
-            S3_V2_CONNECTOR => {
-                tracing::info!(
-                    "using new fs source, rewrite connector from '{}' to '{}'",
-                    S3_V2_CONNECTOR,
-                    S3_CONNECTOR
-                );
-                props.insert(UPSTREAM_SOURCE_KEY.to_string(), S3_CONNECTOR.to_string());
-            }
-            _ => {
-                props.insert(UPSTREAM_SOURCE_KEY.to_string(), connector);
-            }
-        }
+            || props
+                .get(UPSTREAM_SOURCE_KEY)
+                .map(|s| s.eq_ignore_ascii_case(GCS_CONNECTOR))
+                .unwrap_or(false)
     }
 }
 
 impl ConnectorProperties {
     pub fn extract(mut props: HashMap<String, String>) -> Result<Self> {
         if Self::is_new_fs_connector_hash_map(&props) {
-            _ = props
+            let connector = props
                 .remove(UPSTREAM_SOURCE_KEY)
                 .ok_or_else(|| anyhow!("Must specify 'connector' in WITH clause"))?;
-            return Ok(ConnectorProperties::S3(Box::new(
-                S3Properties::try_from_hashmap(props)?,
-            )));
+            match connector.as_str() {
+                "s3_v2" => {
+                    let assume_role = props.get("s3.assume_role").cloned();
+                    return Ok(ConnectorProperties::OpenDalS3(Box::new(
+                        OpendalS3Properties {
+                            s3_properties: S3Properties::try_from_hashmap(props)?,
+                            assume_role,
+                        },
+                    )));
+                }
+                "gcs" => {
+                    return Ok(ConnectorProperties::Gcs(Box::new(
+                        GcsProperties::try_from_hashmap(props)?,
+                    )));
+                }
+                _ => {
+                    unreachable!()
+                }
+            }
         }
 
         let connector = props

src/connector/src/source/filesystem/file_common.rs

@@ -11,6 +11,10 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+use std::fmt::Debug;
+use std::hash::Hash;
+use std::marker::PhantomData;
+
 use anyhow::anyhow;
 use aws_sdk_s3::types::Object;
 use risingwave_common::types::{JsonbVal, Timestamp};
@@ -67,6 +71,70 @@ impl FsSplit {
     }
 }
 
+///  [`OpendalFsSplit`] Describes a file or a split of a file. A file is a generic concept,
+/// and can be a local file, a distributed file system, or am object in S3 bucket.
+#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
+pub struct OpendalFsSplit<C>
+where
+    C: Send + Clone + 'static,
+{
+    pub name: String,
+    pub offset: usize,
+    pub size: usize,
+    marker: PhantomData<C>,
+}
+
+impl<C> From<&Object> for OpendalFsSplit<C>
+where
+    C: Send + Clone + 'static,
+{
+    fn from(value: &Object) -> Self {
+        Self {
+            name: value.key().unwrap().to_owned(),
+            offset: 0,
+            size: value.size().unwrap_or_default() as usize,
+            marker: PhantomData,
+        }
+    }
+}
+
+impl<C> SplitMetaData for OpendalFsSplit<C>
+where
+    C: Sized + Send + Clone + 'static,
+{
+    fn id(&self) -> SplitId {
+        self.name.as_str().into()
+    }
+
+    fn restore_from_json(value: JsonbVal) -> anyhow::Result<Self> {
+        serde_json::from_value(value.take()).map_err(|e| anyhow!(e))
+    }
+
+    fn encode_to_json(&self) -> JsonbVal {
+        serde_json::to_value(self.clone()).unwrap().into()
+    }
+
+    fn update_with_offset(&mut self, start_offset: String) -> anyhow::Result<()> {
+        let offset = start_offset.parse().unwrap();
+        self.offset = offset;
+        Ok(())
+    }
+}
+
+impl<C> OpendalFsSplit<C>
+where
+    C: Send + Clone + 'static,
+{
+    pub fn new(name: String, start: usize, size: usize) -> Self {
+        Self {
+            name,
+            offset: start,
+            size,
+            marker: PhantomData,
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct FsPageItem {
     pub name: String,

src/connector/src/source/filesystem/mod.rs

@@ -12,11 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub use opendal_source::GcsProperties;
 pub use s3::{S3FileReader, S3Properties, S3SplitEnumerator, S3_CONNECTOR};
-
-mod file_common;
+pub mod file_common;
 pub mod nd_streaming;
-pub use file_common::{FsPage, FsPageItem, FsSplit};
+pub use file_common::{FsPage, FsPageItem, FsSplit, OpendalFsSplit};
+pub mod opendal_source;
 mod s3;
 pub mod s3_v2;
-pub const S3_V2_CONNECTOR: &str = "s3_v2";