1use crate::engine::Engine;
18use crate::error::{ErrorCode, McpError};
19use crate::ingest::IngestResult;
20use crate::schema::ColumnSchema;
21use crate::stats::{IngestStats, StatsTimer};
22use hyperdb_api::escape_string_literal;
23use std::path::Path;
24
25#[derive(Debug, Clone)]
30pub struct IcebergIngestOptions {
31 pub table: String,
33 pub mode: String,
35 pub metadata_filename: Option<String>,
40 pub version_as_of: Option<i64>,
43}
44
45#[must_use]
55pub fn build_iceberg_ingest_sql(
56 table: &str,
57 path: &str,
58 opts: &IcebergIngestOptions,
59 is_replace: bool,
60) -> String {
61 let quoted_table = format!("\"{}\"", table.replace('"', "\"\""));
62 let quoted_path = escape_string_literal(path);
63
64 let mut external_args = vec![quoted_path, "format => 'iceberg'".to_string()];
65 if let Some(name) = &opts.metadata_filename {
66 external_args.push(format!(
67 "metadata_filename => {}",
68 escape_string_literal(name)
69 ));
70 }
71 if let Some(version) = opts.version_as_of {
72 external_args.push(format!("version_as_of => {version}"));
73 }
74 let external = format!("external({})", external_args.join(", "));
75
76 if is_replace {
77 format!("CREATE TABLE {quoted_table} AS SELECT * FROM {external}")
78 } else {
79 format!("INSERT INTO {quoted_table} SELECT * FROM {external}")
80 }
81}
82
83fn resolve_iceberg_path(path: &str) -> Result<String, McpError> {
86 let p = Path::new(path);
87 if !p.exists() {
88 return Err(McpError::new(
89 ErrorCode::FileNotFound,
90 format!("Iceberg path does not exist: {path}"),
91 ));
92 }
93 if !p.is_dir() {
94 return Err(McpError::new(
95 ErrorCode::UnsupportedFormat,
96 format!(
97 "Iceberg path must be a directory (the table root with a `metadata/` subdir), not a file: {path}"
98 ),
99 ));
100 }
101 Ok(std::fs::canonicalize(p)
102 .map_err(|e| {
103 McpError::new(
104 ErrorCode::FileNotFound,
105 format!("Cannot resolve Iceberg path {path}: {e}"),
106 )
107 })?
108 .to_string_lossy()
109 .into_owned())
110}
111
112fn describe_table(engine: &Engine, table: &str) -> Result<Vec<ColumnSchema>, McpError> {
118 let desc = engine.describe_table(table)?;
119 let cols = desc.get("columns").and_then(|v| v.as_array());
120 Ok(cols
121 .map(|arr| {
122 arr.iter()
123 .map(|c| ColumnSchema {
124 name: c
125 .get("name")
126 .and_then(|v| v.as_str())
127 .unwrap_or("")
128 .to_string(),
129 hyper_type: c
130 .get("type")
131 .and_then(|v| v.as_str())
132 .unwrap_or("")
133 .to_string(),
134 nullable: c
135 .get("nullable")
136 .and_then(serde_json::Value::as_bool)
137 .unwrap_or(true),
138 })
139 .collect()
140 })
141 .unwrap_or_default())
142}
143
144fn count_rows(engine: &Engine, table: &str) -> Result<u64, McpError> {
149 let quoted = format!("\"{}\"", table.replace('"', "\"\""));
150 let sql = format!("SELECT COUNT(*) FROM {quoted}");
151 let rows = engine.execute_query_to_json(&sql)?;
152 rows.first()
153 .and_then(|r| r.get("count"))
154 .and_then(serde_json::Value::as_u64)
155 .ok_or_else(|| {
156 McpError::new(
157 ErrorCode::InternalError,
158 "Could not read row count after Iceberg ingest",
159 )
160 })
161}
162
163pub fn ingest_iceberg_table(
182 engine: &Engine,
183 path: &str,
184 opts: &IcebergIngestOptions,
185) -> Result<IngestResult, McpError> {
186 let timer = StatsTimer::start();
187 let absolute_path = resolve_iceberg_path(path)?;
188
189 let is_replace = opts.mode != "append";
190 let sql = build_iceberg_ingest_sql(&opts.table, &absolute_path, opts, is_replace);
191
192 let affected = engine.execute_in_transaction(|engine| {
197 if is_replace {
198 let quoted_table = format!("\"{}\"", opts.table.replace('"', "\"\""));
199 engine.execute_command(&format!("DROP TABLE IF EXISTS {quoted_table}"))?;
200 }
201 engine.execute_command(&sql)
202 })?;
203
204 let row_count = if is_replace {
205 count_rows(engine, &opts.table)?
206 } else {
207 affected
208 };
209
210 let schema = describe_table(engine, &opts.table).unwrap_or_default();
211 let elapsed = timer.elapsed_ms();
212 let stats = IngestStats {
213 operation: "load_iceberg".into(),
214 rows: row_count,
215 elapsed_ms: elapsed,
216 bytes_read: 0,
217 bytes_stored: 0,
218 schema_inference_ms: Some(0),
219 table: opts.table.clone(),
220 file_format: Some("iceberg".into()),
221 warning: None,
222 schema_changed: false,
223 };
224
225 Ok(IngestResult {
226 rows: row_count,
227 schema,
228 stats,
229 })
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235
236 #[test]
237 fn builds_replace_sql_with_minimal_options() {
238 let opts = IcebergIngestOptions {
239 table: "my_table".into(),
240 mode: "replace".into(),
241 metadata_filename: None,
242 version_as_of: None,
243 };
244 let sql = build_iceberg_ingest_sql("my_table", "/abs/table", &opts, true);
245 assert_eq!(
246 sql,
247 "CREATE TABLE \"my_table\" AS SELECT * FROM external('/abs/table', format => 'iceberg')"
248 );
249 }
250
251 #[test]
252 fn builds_append_sql_with_metadata_filename() {
253 let opts = IcebergIngestOptions {
254 table: "t".into(),
255 mode: "append".into(),
256 metadata_filename: Some("v2.metadata.json".into()),
257 version_as_of: None,
258 };
259 let sql = build_iceberg_ingest_sql("t", "/abs/t", &opts, false);
260 assert_eq!(
261 sql,
262 "INSERT INTO \"t\" SELECT * FROM external('/abs/t', format => 'iceberg', metadata_filename => 'v2.metadata.json')"
263 );
264 }
265
266 #[test]
267 fn builds_sql_with_version_as_of() {
268 let opts = IcebergIngestOptions {
269 table: "t".into(),
270 mode: "replace".into(),
271 metadata_filename: None,
272 version_as_of: Some(7),
273 };
274 let sql = build_iceberg_ingest_sql("t", "/abs/t", &opts, true);
275 assert_eq!(
276 sql,
277 "CREATE TABLE \"t\" AS SELECT * FROM external('/abs/t', format => 'iceberg', version_as_of => 7)"
278 );
279 }
280
281 #[test]
282 fn builds_sql_with_both_metadata_and_version() {
283 let opts = IcebergIngestOptions {
284 table: "t".into(),
285 mode: "replace".into(),
286 metadata_filename: Some("v3.metadata.json".into()),
287 version_as_of: Some(42),
288 };
289 let sql = build_iceberg_ingest_sql("t", "/abs/t", &opts, true);
290 assert_eq!(
291 sql,
292 "CREATE TABLE \"t\" AS SELECT * FROM external('/abs/t', format => 'iceberg', metadata_filename => 'v3.metadata.json', version_as_of => 42)"
293 );
294 }
295
296 #[test]
297 fn escapes_single_quotes_in_path() {
298 let opts = IcebergIngestOptions {
299 table: "t".into(),
300 mode: "replace".into(),
301 metadata_filename: None,
302 version_as_of: None,
303 };
304 let sql = build_iceberg_ingest_sql("t", "/abs/it's/t", &opts, true);
306 assert!(
307 sql.contains("'/abs/it''s/t'"),
308 "single quote in path must be escaped; got: {sql}"
309 );
310 }
311
312 #[test]
313 fn escapes_single_quotes_in_metadata_filename() {
314 let opts = IcebergIngestOptions {
315 table: "t".into(),
316 mode: "replace".into(),
317 metadata_filename: Some("v1.metadata'.json".into()),
318 version_as_of: None,
319 };
320 let sql = build_iceberg_ingest_sql("t", "/abs/t", &opts, true);
321 assert!(
322 sql.contains("metadata_filename => 'v1.metadata''.json'"),
323 "single quote in metadata_filename must be escaped; got: {sql}"
324 );
325 }
326
327 #[test]
328 fn escapes_quotes_in_table_name() {
329 let opts = IcebergIngestOptions {
330 table: "weird\"name".into(),
331 mode: "replace".into(),
332 metadata_filename: None,
333 version_as_of: None,
334 };
335 let sql = build_iceberg_ingest_sql("weird\"name", "/abs/t", &opts, true);
336 assert!(
337 sql.contains("\"weird\"\"name\""),
338 "double-quote in table name must be escaped; got: {sql}"
339 );
340 }
341}