risingwave_connector/source/filesystem/opendal_source/
opendal_enumerator.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::marker::PhantomData;
16
17use anyhow::Context;
18use async_trait::async_trait;
19use chrono::{DateTime, Utc};
20use futures::stream::{self, BoxStream};
21use futures::{StreamExt, TryStreamExt};
22use opendal::Operator;
23use risingwave_common::types::Timestamptz;
24
25use super::OpendalSource;
26use crate::error::ConnectorResult;
27use crate::source::filesystem::file_common::CompressionFormat;
28use crate::source::filesystem::{FsPageItem, OpendalFsSplit};
29use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
30
31#[inline]
32fn to_timestamptz(ts: Option<opendal::raw::Timestamp>) -> Timestamptz {
33    let system_time = ts
34        .map(std::time::SystemTime::from)
35        .unwrap_or(std::time::UNIX_EPOCH);
36    Timestamptz::from(DateTime::<Utc>::from(system_time))
37}
38
39#[derive(Debug, Clone)]
40pub struct OpendalEnumerator<Src: OpendalSource> {
41    pub op: Operator,
42    // prefix is used to reduce the number of objects to be listed
43    pub(crate) prefix: Option<String>,
44    pub(crate) matcher: Option<glob::Pattern>,
45    pub(crate) marker: PhantomData<Src>,
46    pub(crate) compression_format: CompressionFormat,
47}
48
49#[async_trait]
50impl<Src: OpendalSource> SplitEnumerator for OpendalEnumerator<Src> {
51    type Properties = Src::Properties;
52    type Split = OpendalFsSplit<Src>;
53
54    async fn new(
55        properties: Src::Properties,
56        _context: SourceEnumeratorContextRef,
57    ) -> ConnectorResult<Self> {
58        Src::new_enumerator(properties)
59    }
60
61    async fn list_splits(&mut self) -> ConnectorResult<Vec<OpendalFsSplit<Src>>> {
62        let empty_split: OpendalFsSplit<Src> = OpendalFsSplit::empty_split();
63        let prefix = self.prefix.as_deref().unwrap_or("/");
64        let list_prefix = Self::extract_list_prefix(prefix);
65
66        let mut lister = self.op.lister(&list_prefix).await?;
67        // fetch one item as validation, no need to get all
68        lister
69            .try_next()
70            .await
71            .context("fail to create source, please check your config")?;
72        Ok(vec![empty_split])
73    }
74}
75
76impl<Src: OpendalSource> OpendalEnumerator<Src> {
77    /// Extract the directory to list from a prefix.
78    /// If prefix ends with "/", use it as-is (directory).
79    /// Otherwise, extract parent directory.
80    fn extract_list_prefix(prefix: &str) -> String {
81        if prefix.ends_with("/") {
82            prefix.to_owned()
83        } else if let Some(parent_pos) = prefix.rfind('/') {
84            prefix[..=parent_pos].to_owned()
85        } else {
86            "/".to_owned()
87        }
88    }
89
90    pub async fn list(&self) -> ConnectorResult<ObjectMetadataIter> {
91        let prefix = self.prefix.as_deref().unwrap_or("/");
92        let list_prefix = Self::extract_list_prefix(prefix);
93        let object_lister = self.op.lister_with(&list_prefix).recursive(true).await?;
94
95        let op = self.op.clone();
96        let stream = stream::unfold(object_lister, move |mut object_lister| {
97            let op = op.clone();
98
99            async move {
100                match object_lister.next().await {
101                    Some(Ok(object)) => {
102                        let name = object.path().to_owned();
103
104                        // OpenDAL 0.55 removed list metadata capability flags and reports
105                        // unknown content length as 0. Use listed metadata first, and call
106                        // stat() if timestamp is missing or size is 0 to avoid treating
107                        // unknown sizes as real zero-byte objects.
108                        let meta = object.metadata();
109                        let mut t = meta.last_modified();
110                        let mut size = meta.content_length() as i64;
111                        if t.is_none() || size == 0 {
112                            let stat_meta = match op
113                                .stat(&name)
114                                .await
115                                .with_context(|| format!("failed to stat listed object {name}"))
116                            {
117                                Ok(stat_meta) => stat_meta,
118                                Err(err) => return Some((Err(err.into()), object_lister)),
119                            };
120                            t = stat_meta.last_modified();
121                            size = stat_meta.content_length() as i64;
122                        }
123
124                        let timestamp = to_timestamptz(t);
125                        let metadata = FsPageItem {
126                            name,
127                            size,
128                            timestamp,
129                        };
130                        Some((Ok(metadata), object_lister))
131                    }
132                    Some(Err(err)) => Some((Err(err.into()), object_lister)),
133                    None => {
134                        tracing::info!("list object completed.");
135                        None
136                    }
137                }
138            }
139        });
140
141        Ok(stream.boxed())
142    }
143
144    pub fn get_matcher(&self) -> &Option<glob::Pattern> {
145        &self.matcher
146    }
147
148    pub fn get_prefix(&self) -> &str {
149        self.prefix.as_deref().unwrap_or("/")
150    }
151}
152pub type ObjectMetadataIter = BoxStream<'static, ConnectorResult<FsPageItem>>;
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::source::filesystem::opendal_source::OpendalS3;
158    use crate::source::filesystem::s3::enumerator::get_prefix;
159
160    fn calculate_list_prefix(prefix: &str) -> String {
161        OpendalEnumerator::<OpendalS3>::extract_list_prefix(prefix)
162    }
163
164    #[test]
165    fn test_prefix_logic() {
166        let test_cases = vec![
167            ("a/b/c/hello*/*.json", "a/b/c/hello", "a/b/c/"),
168            ("a/b/c.json", "a/b/c.json", "a/b/"),
169            ("a/b/c/", "a/b/c/", "a/b/c/"),
170            ("a/b/c", "a/b/c", "a/b/"),
171            ("file.json", "file.json", "/"),
172            ("*.json", "", "/"),
173            ("a/b/c/[h]ello*/*.json", "a/b/c/", "a/b/c/"),
174        ];
175
176        for (pattern, expected_prefix, expected_list_prefix) in test_cases {
177            let prefix = get_prefix(pattern);
178            let list_prefix = calculate_list_prefix(&prefix);
179
180            assert_eq!(
181                prefix, expected_prefix,
182                "get_prefix failed for: {}",
183                pattern
184            );
185            assert_eq!(
186                list_prefix, expected_list_prefix,
187                "list_prefix failed for: {}",
188                pattern
189            );
190        }
191    }
192
193    #[test]
194    fn test_bug_fix() {
195        let problematic_pattern = "a/b/c/hello*/*.json";
196        let prefix = get_prefix(problematic_pattern);
197        let list_prefix = calculate_list_prefix(&prefix);
198
199        // Before fix: would fallback to "/"
200        // After fix: should use parent directory "a/b/c/"
201        assert_eq!(prefix, "a/b/c/hello");
202        assert_eq!(list_prefix, "a/b/c/");
203    }
204}
risingwave_connector/source/filesystem/opendal_source/opendal_enumerator.rs

risingwave_connector/source/filesystem/opendal_source/
opendal_enumerator.rs
