cli/data/import_v2/
packed.rs

1// Copyright 2023 Greptime Team
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::HashSet;
16
17use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex};
18use snafu::ResultExt;
19use sql::dialect::GreptimeDbDialect;
20use sql::parser::{ParseOptions, ParserContext};
21use sql::statements::statement::Statement;
22
23use crate::data::export_v2::manifest::{ChunkStatus, Manifest};
24use crate::data::import_v2::error::{InvalidPackedSnapshotSnafu, Result, SnapshotStorageSnafu};
25use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema};
26use crate::data::snapshot_storage::SnapshotStorage;
27
28/// Matches every completed schema chunk to its source DDL before target writes.
29pub(crate) async fn validate_snapshot(
30    storage: &dyn SnapshotStorage,
31    manifest: &Manifest,
32    schemas: &[String],
33) -> Result<()> {
34    manifest
35        .validate_layout()
36        .map_err(|reason| InvalidPackedSnapshotSnafu { reason }.build())?;
37    if manifest.schema_only
38        && !storage
39            .list_files_recursive("data/")
40            .await
41            .context(SnapshotStorageSnafu)?
42            .is_empty()
43    {
44        return InvalidPackedSnapshotSnafu {
45            reason: "schema-only snapshot contains data objects",
46        }
47        .fail();
48    }
49    for schema in schemas {
50        let ddl = storage
51            .read_text(&ddl_path_for_schema(schema))
52            .await
53            .context(SnapshotStorageSnafu)?;
54        let statements = ParserContext::create_with_dialect(
55            &ddl,
56            &GreptimeDbDialect {},
57            ParseOptions::default(),
58        )
59        .map_err(|e| {
60            InvalidPackedSnapshotSnafu {
61                reason: e.to_string(),
62            }
63            .build()
64        })?;
65        let mut names = HashSet::new();
66        for statement in statements {
67            if let Statement::CreateTable(create) = statement {
68                if create.engine.eq_ignore_ascii_case("metric")
69                    && create.options.get("on_physical_table").is_none()
70                {
71                    continue;
72                }
73                let parts: Vec<_> = create
74                    .name
75                    .0
76                    .iter()
77                    .map(|p| p.as_ident().map(|i| i.value.as_str()))
78                    .collect();
79                let allowed = match parts.as_slice() {
80                    [Some(_)] => true,
81                    [Some(s), Some(_)] => *s == schema,
82                    [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema,
83                    _ => false,
84                };
85                if !allowed {
86                    return InvalidPackedSnapshotSnafu {
87                        reason: "snapshot table DDL escapes its catalog/schema",
88                    }
89                    .fail();
90                }
91                let name = parts.last().and_then(|p| *p).unwrap_or_default();
92                if !names.insert(name.to_string()) {
93                    return InvalidPackedSnapshotSnafu {
94                        reason: "duplicate snapshot table DDL",
95                    }
96                    .fail();
97                }
98            }
99        }
100        for chunk in &manifest.chunks {
101            if chunk.status == ChunkStatus::Skipped && !names.is_empty() {
102                return InvalidPackedSnapshotSnafu {
103                    reason: "packed chunks require zero-row streams for empty tables",
104                }
105                .fail();
106            }
107            if chunk.status != ChunkStatus::Completed {
108                continue;
109            }
110            let path = format!(
111                "{}{PACK_INDEX_FILE}",
112                data_dir_for_schema_chunk(schema, chunk.id)
113            );
114            let text = storage
115                .read_text(&path)
116                .await
117                .context(SnapshotStorageSnafu)?;
118            let index: PackIndex = serde_json::from_str(&text).map_err(|e| {
119                InvalidPackedSnapshotSnafu {
120                    reason: e.to_string(),
121                }
122                .build()
123            })?;
124            let prefix = data_dir_for_schema_chunk(schema, chunk.id);
125            let expected: HashSet<_> = std::iter::once(path)
126                .chain(index.objects.iter().map(|o| format!("{prefix}{}", o.path)))
127                .collect();
128            let actual: HashSet<_> = chunk
129                .files
130                .iter()
131                .filter(|file| file.starts_with(&prefix))
132                .cloned()
133                .collect();
134            if actual != expected {
135                return InvalidPackedSnapshotSnafu {
136                    reason: "chunk manifest inventory differs from indexed objects",
137                }
138                .fail();
139            }
140            index
141                .validate_membership(names.iter().map(String::as_str))
142                .map_err(|e| {
143                    InvalidPackedSnapshotSnafu {
144                        reason: e.to_string(),
145                    }
146                    .build()
147                })?;
148            for object in &index.objects {
149                let path = format!("{prefix}{}", object.path);
150                let length = storage
151                    .file_size(&path)
152                    .await
153                    .context(SnapshotStorageSnafu)?;
154                if length != Some(object.length) {
155                    return InvalidPackedSnapshotSnafu {
156                        reason: format!(
157                            "object {path}: expected length {}, actual {length:?}",
158                            object.length
159                        ),
160                    }
161                    .fail();
162                }
163            }
164        }
165    }
166    Ok(())
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172    use crate::data::export_v2::manifest::{ChunkMeta, TimeRange};
173    use crate::data::snapshot_storage::OpenDalStorage;
174
175    #[tokio::test]
176    async fn packed_membership_uses_literal_names_and_excludes_physical_and_views() {
177        let dir = tempfile::Builder::new()
178            .prefix("packed snapshot # ")
179            .tempdir()
180            .unwrap();
181        let uri = url::Url::from_file_path(dir.path()).unwrap();
182        let storage = OpenDalStorage::from_uri(uri.as_str(), &Default::default()).unwrap();
183        storage.write_text("schema/ddl/public.sql", "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\";").await.unwrap();
184        let path = "data/public/1/pack-index.json";
185        let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":12}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0}]});
186        storage.write_text(path, &index.to_string()).await.unwrap();
187        assert!(dir.path().join(path).is_file());
188        storage
189            .write_text("data/public/1/pack-0.bin", "abcdefghijkl")
190            .await
191            .unwrap();
192        let mut manifest = Manifest::new_schema_only("greptime".into(), vec!["public".into()]);
193        let mut chunk = ChunkMeta::new(1, TimeRange::unbounded());
194        chunk.mark_completed(vec![path.into(), "data/public/1/pack-0.bin".into()], None);
195        manifest.version = 2;
196        manifest.data_layout = Some("metric-parquet-packs".into());
197        manifest.schema_only = false;
198        manifest.chunks.push(chunk);
199        validate_snapshot(&storage, &manifest, &manifest.schemas)
200            .await
201            .unwrap();
202        let mut missing = index.clone();
203        missing["tables"] = serde_json::json!([]);
204        storage
205            .write_text(path, &missing.to_string())
206            .await
207            .unwrap();
208        assert!(
209            validate_snapshot(&storage, &manifest, &manifest.schemas)
210                .await
211                .is_err()
212        );
213        let mut escaped = index;
214        escaped["objects"][0]["path"] = "../pack-0.bin".into();
215        storage
216            .write_text(path, &escaped.to_string())
217            .await
218            .unwrap();
219        assert!(
220            validate_snapshot(&storage, &manifest, &manifest.schemas)
221                .await
222                .is_err()
223        );
224    }
225}
cli/data/import_v2/packed.rs

cli/data/import_v2/
packed.rs
