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

cli/data/import_v2/
packed.rs
