risingwave_connector/connector_common/iceberg/
mock_catalog.rs

1// Copyright 2024 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::collections::HashMap;
16
17use async_trait::async_trait;
18use iceberg::io::FileIO;
19use iceberg::spec::{
20    NestedField, PrimitiveType, Schema, TableMetadataBuilder, Transform, Type,
21    UnboundPartitionField, UnboundPartitionSpec,
22};
23use iceberg::table::Table;
24use iceberg::{
25    Catalog as CatalogV2, Namespace, NamespaceIdent, Runtime, TableCommit, TableCreation,
26    TableIdent,
27};
28
29#[cfg(any(test, feature = "test"))]
30pub use self::test_tables::MockTableGuard;
31
32#[cfg(any(test, feature = "test"))]
33mod test_tables {
34    use std::collections::HashMap;
35    use std::collections::hash_map::Entry;
36    use std::sync::LazyLock;
37
38    use iceberg::io::FileIO;
39    use iceberg::spec::{
40        FormatVersion, NestedField, PrimitiveType, Schema, SortOrder, TableMetadataBuilder, Type,
41        UnboundPartitionSpec,
42    };
43    use iceberg::table::Table;
44    use iceberg::{NamespaceIdent, Runtime, TableIdent};
45    use parking_lot::Mutex;
46
47    use super::MockCatalog;
48
49    static TABLES: LazyLock<Mutex<HashMap<TableIdent, Table>>> =
50        LazyLock::new(|| Mutex::new(HashMap::new()));
51
52    /// A registered table that is removed from the mock catalog when dropped.
53    #[must_use]
54    pub struct MockTableGuard {
55        ident: TableIdent,
56    }
57
58    impl MockTableGuard {
59        pub fn update(&self, table: Table) {
60            assert_eq!(table.identifier(), &self.ident);
61            *TABLES.lock().get_mut(&self.ident).unwrap() = table;
62        }
63    }
64
65    impl Drop for MockTableGuard {
66        fn drop(&mut self) {
67            TABLES.lock().remove(&self.ident);
68        }
69    }
70
71    impl MockCatalog {
72        /// Register a table shared by mock catalog instances. Use distinct identifiers in parallel tests.
73        pub fn register(table: Table) -> MockTableGuard {
74            let ident = table.identifier().clone();
75            match TABLES.lock().entry(ident.clone()) {
76                Entry::Vacant(entry) => {
77                    entry.insert(table);
78                }
79                Entry::Occupied(_) => panic!("mock table {ident} is already registered"),
80            }
81            MockTableGuard { ident }
82        }
83
84        /// Registered tables and the built-in `updates_v2`/`updates_v3` tables.
85        pub(super) fn test_table(ident: &TableIdent) -> Option<Table> {
86            if let Some(table) = TABLES.lock().get(ident).cloned() {
87                return Some(table);
88            }
89            match ident.name() {
90                "updates_v2" => Some(Self::updates_table(ident.name(), FormatVersion::V2)),
91                "updates_v3" => Some(Self::updates_table(ident.name(), FormatVersion::V3)),
92                _ => None,
93            }
94        }
95
96        pub(super) fn test_table_exists(ident: &TableIdent) -> bool {
97            TABLES.lock().contains_key(ident) || matches!(ident.name(), "updates_v2" | "updates_v3")
98        }
99
100        pub(super) fn updates_table(name: &str, version: FormatVersion) -> Table {
101            let schema = Schema::builder()
102                .with_fields(vec![
103                    NestedField::optional(10, "id", Type::Primitive(PrimitiveType::Int)).into(),
104                    NestedField::optional(20, "value", Type::Primitive(PrimitiveType::Long)).into(),
105                    NestedField::optional(30, "_row_id", Type::Primitive(PrimitiveType::Long))
106                        .into(),
107                ])
108                .build()
109                .unwrap();
110            let metadata = TableMetadataBuilder::new(
111                schema,
112                UnboundPartitionSpec::builder().build(),
113                SortOrder::unsorted_order(),
114                "memory://updates".to_owned(),
115                version,
116                HashMap::new(),
117            )
118            .unwrap()
119            .build()
120            .unwrap()
121            .metadata;
122            // No writer properties, like tables created by older PK-index sinks.
123            Table::builder()
124                .identifier(TableIdent::new(
125                    NamespaceIdent::new("mock_namespace".to_owned()),
126                    name.to_owned(),
127                ))
128                .file_io(FileIO::new_with_memory())
129                .runtime(Runtime::try_current().unwrap())
130                .metadata(metadata)
131                .build()
132                .unwrap()
133        }
134    }
135}
136
137/// A mock catalog for Iceberg planner and executor tests.
138#[derive(Debug)]
139pub struct MockCatalog;
140
141impl MockCatalog {
142    const RANGE_TABLE: &'static str = "range_table";
143    const SPARSE_TABLE: &'static str = "sparse_table";
144}
145
146impl MockCatalog {
147    fn build_table(name: &str, schema: Schema, partition_spec: UnboundPartitionSpec) -> Table {
148        let file_io = FileIO::new_with_memory();
149        let table_creation = TableCreation {
150            name: "ignore".to_owned(),
151            location: Some("1".to_owned()),
152            schema,
153            partition_spec: Some(partition_spec),
154            sort_order: None,
155            properties: HashMap::new(),
156            format_version: iceberg::spec::FormatVersion::V2,
157        };
158        Table::builder()
159            .identifier(TableIdent::new(
160                NamespaceIdent::new("mock_namespace".to_owned()),
161                name.to_owned(),
162            ))
163            .file_io(file_io)
164            .runtime(Runtime::try_current().unwrap())
165            .metadata(
166                TableMetadataBuilder::from_table_creation(table_creation)
167                    .unwrap()
168                    .build()
169                    .unwrap()
170                    .metadata,
171            )
172            .build()
173            .unwrap()
174    }
175
176    fn sparse_table() -> Table {
177        Self::build_table(
178            Self::SPARSE_TABLE,
179            Schema::builder()
180                .with_fields(vec![
181                    NestedField::new(1, "v1", Type::Primitive(PrimitiveType::Int), true).into(),
182                    NestedField::new(2, "v2", Type::Primitive(PrimitiveType::Long), true).into(),
183                    NestedField::new(3, "v3", Type::Primitive(PrimitiveType::String), true).into(),
184                    NestedField::new(4, "v4", Type::Primitive(PrimitiveType::Time), true).into(),
185                ])
186                .build()
187                .unwrap(),
188            UnboundPartitionSpec::builder()
189                .with_spec_id(1)
190                .add_partition_fields(vec![
191                    UnboundPartitionField {
192                        source_id: 1,
193                        field_id: Some(5),
194                        name: "f1".to_owned(),
195                        transform: Transform::Identity,
196                    },
197                    UnboundPartitionField {
198                        source_id: 2,
199                        field_id: Some(6),
200                        name: "f2".to_owned(),
201                        transform: Transform::Bucket(1),
202                    },
203                    UnboundPartitionField {
204                        source_id: 3,
205                        field_id: Some(7),
206                        name: "f3".to_owned(),
207                        transform: Transform::Truncate(1),
208                    },
209                    UnboundPartitionField {
210                        source_id: 4,
211                        field_id: Some(8),
212                        name: "f4".to_owned(),
213                        transform: Transform::Void,
214                    },
215                ])
216                .unwrap()
217                .build(),
218        )
219    }
220
221    fn range_table() -> Table {
222        Self::build_table(
223            Self::RANGE_TABLE,
224            Schema::builder()
225                .with_fields(vec![
226                    NestedField::new(1, "v1", Type::Primitive(PrimitiveType::Date), true).into(),
227                    NestedField::new(2, "v2", Type::Primitive(PrimitiveType::Timestamp), true)
228                        .into(),
229                    NestedField::new(3, "v3", Type::Primitive(PrimitiveType::Timestamptz), true)
230                        .into(),
231                    NestedField::new(4, "v4", Type::Primitive(PrimitiveType::Timestamptz), true)
232                        .into(),
233                ])
234                .build()
235                .unwrap(),
236            UnboundPartitionSpec::builder()
237                .with_spec_id(1)
238                .add_partition_fields(vec![
239                    UnboundPartitionField {
240                        source_id: 1,
241                        field_id: Some(5),
242                        name: "f1".to_owned(),
243                        transform: Transform::Year,
244                    },
245                    UnboundPartitionField {
246                        source_id: 2,
247                        field_id: Some(6),
248                        name: "f2".to_owned(),
249                        transform: Transform::Month,
250                    },
251                    UnboundPartitionField {
252                        source_id: 3,
253                        field_id: Some(7),
254                        name: "f3".to_owned(),
255                        transform: Transform::Day,
256                    },
257                    UnboundPartitionField {
258                        source_id: 4,
259                        field_id: Some(8),
260                        name: "f4".to_owned(),
261                        transform: Transform::Hour,
262                    },
263                ])
264                .unwrap()
265                .build(),
266        )
267    }
268}
269
270#[async_trait]
271impl CatalogV2 for MockCatalog {
272    /// List namespaces from table.
273    async fn list_namespaces(
274        &self,
275        _parent: Option<&NamespaceIdent>,
276    ) -> iceberg::Result<Vec<NamespaceIdent>> {
277        todo!()
278    }
279
280    /// Create a new namespace inside the catalog.
281    async fn create_namespace(
282        &self,
283        _namespace: &iceberg::NamespaceIdent,
284        _properties: HashMap<String, String>,
285    ) -> iceberg::Result<iceberg::Namespace> {
286        todo!()
287    }
288
289    /// Get a namespace information from the catalog.
290    async fn get_namespace(&self, _namespace: &NamespaceIdent) -> iceberg::Result<Namespace> {
291        todo!()
292    }
293
294    /// Check if namespace exists in catalog.
295    async fn namespace_exists(&self, _namespace: &NamespaceIdent) -> iceberg::Result<bool> {
296        todo!()
297    }
298
299    /// Drop a namespace from the catalog.
300    async fn drop_namespace(&self, _namespace: &NamespaceIdent) -> iceberg::Result<()> {
301        todo!()
302    }
303
304    /// List tables from namespace.
305    async fn list_tables(&self, _namespace: &NamespaceIdent) -> iceberg::Result<Vec<TableIdent>> {
306        todo!()
307    }
308
309    async fn update_namespace(
310        &self,
311        _namespace: &NamespaceIdent,
312        _properties: HashMap<String, String>,
313    ) -> iceberg::Result<()> {
314        todo!()
315    }
316
317    /// Create a new table inside the namespace.
318    async fn create_table(
319        &self,
320        _namespace: &NamespaceIdent,
321        _creation: TableCreation,
322    ) -> iceberg::Result<Table> {
323        todo!()
324    }
325
326    /// Load table from the catalog.
327    async fn load_table(&self, table: &TableIdent) -> iceberg::Result<Table> {
328        #[cfg(any(test, feature = "test"))]
329        if let Some(table) = Self::test_table(table) {
330            return Ok(table);
331        }
332        match table.name.as_ref() {
333            Self::SPARSE_TABLE => Ok(Self::sparse_table()),
334            Self::RANGE_TABLE => Ok(Self::range_table()),
335            _ => unimplemented!("table {} not found", table.name()),
336        }
337    }
338
339    /// Drop a table from the catalog.
340    async fn drop_table(&self, _table: &TableIdent) -> iceberg::Result<()> {
341        todo!()
342    }
343
344    async fn purge_table(&self, table: &TableIdent) -> iceberg::Result<()> {
345        self.drop_table(table).await
346    }
347
348    /// Check if a table exists in the catalog.
349    async fn table_exists(&self, table: &TableIdent) -> iceberg::Result<bool> {
350        #[cfg(any(test, feature = "test"))]
351        if Self::test_table_exists(table) {
352            return Ok(true);
353        }
354        match table.name.as_ref() {
355            Self::SPARSE_TABLE => Ok(true),
356            Self::RANGE_TABLE => Ok(true),
357            _ => Ok(false),
358        }
359    }
360
361    /// Rename a table in the catalog.
362    async fn rename_table(&self, _src: &TableIdent, _dest: &TableIdent) -> iceberg::Result<()> {
363        todo!()
364    }
365
366    /// Update a table to the catalog.
367    async fn update_table(&self, _commit: TableCommit) -> iceberg::Result<Table> {
368        todo!()
369    }
370
371    #[expect(
372        clippy::disallowed_types,
373        reason = "iceberg catalog trait requires returning iceberg::Error"
374    )]
375    async fn register_table(
376        &self,
377        _table_ident: &TableIdent,
378        _metadata_location: String,
379    ) -> iceberg::Result<Table> {
380        Err(iceberg::Error::new(
381            iceberg::ErrorKind::Unexpected,
382            "register_table is not supported in mock catalog",
383        ))
384    }
385}
386
387// Iceberg tables require a native Tokio runtime, which madsim does not provide.
388#[cfg(all(test, not(madsim)))]
389mod tests {
390    use super::*;
391
392    #[tokio::test]
393    async fn registered_tables_are_isolated_and_scoped() -> iceberg::Result<()> {
394        use iceberg::spec::FormatVersion;
395
396        let first = MockCatalog::updates_table(
397            &format!("registered_{}", uuid::Uuid::new_v4().simple()),
398            FormatVersion::V2,
399        );
400        let second = MockCatalog::updates_table(
401            &format!("registered_{}", uuid::Uuid::new_v4().simple()),
402            FormatVersion::V2,
403        );
404        let first_guard = MockCatalog::register(first.clone());
405        let second_guard = MockCatalog::register(second.clone());
406        let catalog = MockCatalog;
407        assert!(catalog.table_exists(first.identifier()).await?);
408        assert!(catalog.table_exists(second.identifier()).await?);
409        assert_eq!(
410            catalog
411                .load_table(first.identifier())
412                .await?
413                .metadata()
414                .uuid(),
415            first.metadata().uuid()
416        );
417
418        let updated = MockCatalog::updates_table(first.identifier().name(), FormatVersion::V3);
419        first_guard.update(updated.clone());
420        assert_eq!(
421            catalog
422                .load_table(first.identifier())
423                .await?
424                .metadata()
425                .uuid(),
426            updated.metadata().uuid()
427        );
428        assert_eq!(
429            catalog
430                .load_table(second.identifier())
431                .await?
432                .metadata()
433                .uuid(),
434            second.metadata().uuid()
435        );
436        drop(first_guard);
437        assert!(!catalog.table_exists(first.identifier()).await?);
438        assert!(catalog.table_exists(second.identifier()).await?);
439        drop(second_guard);
440        assert!(!catalog.table_exists(second.identifier()).await?);
441        Ok(())
442    }
443}
risingwave_connector/connector_common/iceberg/mock_catalog.rs

risingwave_connector/connector_common/iceberg/
mock_catalog.rs
