diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 03bdc169c..cb3055848 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -380,6 +380,9 @@ fn format_export_sql_literal_typed( if matches!(database_type, Some(DatabaseType::Mysql)) && column_type.is_some_and(is_mysql_bit_type) { return format_mysql_bit_literal(value); } + if let Some(literal) = format_mysql_spatial_export_literal(value, database_type, column_type) { + return literal; + } if is_mysql_compatible_export_literal_target(database_type) { if column_type.is_some_and(is_mysql_binary_export_type) { if let Some(literal) = format_mysql_binary_export_literal(value) { @@ -769,6 +772,72 @@ fn is_mysql_binary_export_type(column_type: &str) -> bool { matches!(base, "binary" | "varbinary" | "blob" | "tinyblob" | "mediumblob" | "longblob") } +pub(crate) fn is_mysql_spatial_export_type(column_type: &str) -> bool { + let base = column_type + .trim() + .to_ascii_lowercase() + .split(['(', ':', ' ', '\t', '\n']) + .next() + .unwrap_or("") + .trim() + .to_string(); + matches!( + base.as_str(), + "geometry" + | "point" + | "linestring" + | "polygon" + | "multipoint" + | "multilinestring" + | "multipolygon" + | "geometrycollection" + | "geomcollection" + ) +} + +/// Database exports encode MySQL spatial cells as `DBX_WKB::` while +/// reading them. Keeping this marker internal lets the normal JSON row shape +/// and all non-export query paths continue to expose readable WKT values. +fn format_mysql_spatial_export_literal( + value: &Value, + database_type: Option, + column_type: Option<&str>, +) -> Option { + if database_type != Some(DatabaseType::Mysql) || !column_type.is_some_and(is_mysql_spatial_export_type) { + return None; + } + let Value::String(value) = value else { + return value.is_null().then(|| "NULL".to_string()); + }; + let marker = value.strip_prefix("DBX_WKB:")?; + let (srid, hex) = marker.split_once(':')?; + if srid.is_empty() + || !srid.as_bytes().iter().all(u8::is_ascii_digit) + || hex.is_empty() + || hex.len() % 2 != 0 + || !hex.as_bytes().iter().all(u8::is_ascii_hexdigit) + { + return None; + } + let wkb = decode_mysql_spatial_export_wkb(hex)?; + crate::db::wkb::decode_wkb_geometry(&wkb)?; + let srid = srid.parse::().ok()?; + Some(if srid == 0 { format!("ST_GeomFromWKB(0x{hex})") } else { format!("ST_GeomFromWKB(0x{hex}, {srid})") }) +} + +fn decode_mysql_spatial_export_wkb(hex: &str) -> Option> { + fn nibble(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } + } + + hex.as_bytes().chunks_exact(2).map(|pair| Some((nibble(pair[0])? << 4) | nibble(pair[1])?)).collect() +} + fn format_mysql_binary_export_literal(value: &Value) -> Option { match value { Value::Null => Some("NULL".to_string()), @@ -1376,8 +1445,59 @@ fn record_export_error(file: &mut W, fail_on_error: bool, message: Str } } -fn database_export_select_sql(columns: &[String], table: &str, schema: &str, db_type: &DatabaseType) -> String { - let columns = columns.iter().map(|column| quote_identifier(column, db_type)).collect::>().join(", "); +fn mysql_spatial_export_marker_expression(column: &str) -> String { + let quoted = quote_identifier(column, &DatabaseType::Mysql); + format!( + "CASE WHEN {quoted} IS NULL THEN NULL ELSE CONCAT('DBX_WKB:', ST_SRID({quoted}), ':', HEX(ST_AsWKB({quoted}))) END AS {quoted}" + ) +} + +fn database_export_select_list(columns: &[String], column_types: &[Option], db_type: &DatabaseType) -> String { + columns + .iter() + .enumerate() + .map(|(index, column)| { + if *db_type == DatabaseType::Mysql + && column_types.get(index).and_then(|value| value.as_deref()).is_some_and(is_mysql_spatial_export_type) + { + mysql_spatial_export_marker_expression(column) + } else { + quote_identifier(column, db_type) + } + }) + .collect::>() + .join(", ") +} + +fn replace_database_export_select_list( + sql: String, + columns: &[String], + column_types: &[Option], + db_type: &DatabaseType, +) -> String { + let original = columns.iter().map(|column| quote_identifier(column, db_type)).collect::>().join(", "); + let replacement = database_export_select_list(columns, column_types, db_type); + if replacement == original { + return sql; + } + let prefix = format!("SELECT {original}"); + if !sql.starts_with(&prefix) { + log::warn!( + "MySQL spatial database export could not replace its SELECT list; geometry columns will be exported as WKT" + ); + return sql; + } + format!("SELECT {replacement}{}", &sql[prefix.len()..]) +} + +fn database_export_select_sql( + columns: &[String], + column_types: &[Option], + table: &str, + schema: &str, + db_type: &DatabaseType, +) -> String { + let columns = database_export_select_list(columns, column_types, db_type); let table = crate::transfer::qualified_table(table, schema, db_type, None); format!("SELECT {columns} FROM {table}") } @@ -1966,7 +2086,7 @@ pub async fn export_database_sql_core( if !col_names.is_empty() { if let Some(snapshot_session_id) = request.snapshot_session_id.as_deref() { - let sql = database_export_select_sql(&col_names, table_name, &request.schema, &db_type); + let sql = database_export_select_sql(&col_names, &col_types, table_name, &request.schema, &db_type); crate::query::stream_rows_in_manual_transaction( state, snapshot_session_id, @@ -2037,7 +2157,7 @@ pub async fn export_database_sql_core( } let sql = if use_keyset { - keyset_pagination_sql_with_identifier_quote( + let sql = keyset_pagination_sql_with_identifier_quote( &col_names, table_name, &request.schema, @@ -2046,16 +2166,18 @@ pub async fn export_database_sql_core( &last_primary_key_values, batch_size, None, - ) + ); + replace_database_export_select_list(sql, &col_names, &col_types, &db_type) } else { - crate::transfer::pagination_sql( + let sql = crate::transfer::pagination_sql( &col_names, table_name, &request.schema, &db_type, offset, batch_size, - ) + ); + replace_database_export_select_list(sql, &col_names, &col_types, &db_type) }; let result = match crate::transfer::execute_read_on_pool(state, &pool_key, &sql).await { Ok(result) => result, @@ -2359,15 +2481,16 @@ fn build_database_export_object_source_sql( mod tests { use super::{ build_database_export_object_source_sql, build_database_sql_export, build_export_insert_statements, - database_export_total_objects, drop_table_if_exists_sql, filter_export_table_infos, format_export_sql_literal, - format_export_table_ddl, generate_postgres_extension_ddl, generate_postgres_sequence_create_ddl, - generate_postgres_sequence_owner_ddl, generate_postgres_sequence_setval_sql, - is_postgres_extension_member_routine, mysql_database_export_preamble, mysql_view_dependencies_from_rows, - mysql_view_dependencies_sql, normalize_export_table_ddl, record_export_error, - sort_export_views_by_dependencies, write_database_export_rows, BuildDatabaseSqlExportOptions, - BuildExportInsertStatementsOptions, DatabaseExportObjectCounts, DatabaseExportRequest, DdlNormalizeOptions, - ExportedTableSql, PostgresExportExtension, PostgresExportSequence, PostgresExtensionMembers, - DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT, + database_export_select_sql, database_export_total_objects, drop_table_if_exists_sql, filter_export_table_infos, + format_export_sql_literal, format_export_table_ddl, format_mysql_spatial_export_literal, + generate_postgres_extension_ddl, generate_postgres_sequence_create_ddl, generate_postgres_sequence_owner_ddl, + generate_postgres_sequence_setval_sql, is_postgres_extension_member_routine, mysql_database_export_preamble, + mysql_view_dependencies_from_rows, mysql_view_dependencies_sql, normalize_export_table_ddl, + record_export_error, replace_database_export_select_list, sort_export_views_by_dependencies, + write_database_export_rows, BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, + DatabaseExportObjectCounts, DatabaseExportRequest, DdlNormalizeOptions, ExportedTableSql, + PostgresExportExtension, PostgresExportSequence, PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, + DATABASE_EXPORT_ROW_LIMIT, }; use super::{concurrent_metadata_prefetch_allowed, database_export_metadata_prefetch_concurrency}; use crate::models::connection::DatabaseType; @@ -2670,6 +2793,83 @@ mod tests { assert_eq!(format_export_sql_literal(&json!("O'Hara")), "'O''Hara'"); } + #[test] + fn mysql_spatial_export_uses_wkb_constructor_and_preserves_srid() { + let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { + database_type: Some(DatabaseType::Mysql), + schema: None, + table_name: Some("places".to_string()), + qualified_table_name: None, + columns: vec!["location".to_string(), "shape".to_string()], + column_types: vec![Some("point".to_string()), Some("geometry".to_string())], + column_extras: Vec::new(), + rows: vec![vec![ + json!("DBX_WKB:4326:0101000000AE47E17A14AE5C4052B81E85EBF34240"), + json!("DBX_WKB:0:0101000000000000000000F03F0000000000000040"), + ]], + batch_size: Some(10), + }) + .unwrap(); + + assert_eq!( + statements, + vec!["INSERT INTO `places` (`location`, `shape`) VALUES (ST_GeomFromWKB(0x0101000000AE47E17A14AE5C4052B81E85EBF34240, 4326), ST_GeomFromWKB(0x0101000000000000000000F03F0000000000000040));"] + ); + } + + #[test] + fn mysql_spatial_export_rejects_markers_for_unknown_or_nonspatial_types() { + let marker = json!("DBX_WKB:4326:0101000000AE47E17A14AE5C4052B81E85EBF34240"); + assert!(format_mysql_spatial_export_literal(&marker, Some(DatabaseType::Mysql), None).is_none()); + assert!(format_mysql_spatial_export_literal(&marker, Some(DatabaseType::Mysql), Some("varchar")).is_none()); + assert!(format_mysql_spatial_export_literal( + &json!("DBX_WKB:4326:0101000000"), + Some(DatabaseType::Mysql), + Some("geometry"), + ) + .is_none()); + } + + #[test] + fn mysql_spatial_export_select_normalizes_geometry_columns_to_wkb_markers() { + let sql = database_export_select_sql( + &["id".to_string(), "location".to_string(), "name".to_string()], + &[Some("int".to_string()), Some("point".to_string()), Some("varchar(32)".to_string())], + "places", + "app", + &DatabaseType::Mysql, + ); + + assert_eq!( + sql, + "SELECT `id`, CASE WHEN `location` IS NULL THEN NULL ELSE CONCAT('DBX_WKB:', ST_SRID(`location`), ':', HEX(ST_AsWKB(`location`))) END AS `location`, `name` FROM `places`" + ); + } + + #[test] + fn mysql_spatial_export_keeps_keyset_pagination_when_replacing_select_list() { + let columns = vec!["id".to_string(), "location".to_string()]; + let column_types = vec![Some("bigint".to_string()), Some("geometry".to_string())]; + let sql = crate::transfer::keyset_pagination_sql_with_identifier_quote( + &columns, + "places", + "app", + &DatabaseType::Mysql, + &["id".to_string()], + &[json!(7)], + 1000, + None, + ); + + let sql = replace_database_export_select_list(sql, &columns, &column_types, &DatabaseType::Mysql); + + assert!(sql.starts_with( + "SELECT `id`, CASE WHEN `location` IS NULL THEN NULL ELSE CONCAT('DBX_WKB:', ST_SRID(`location`), ':', HEX(ST_AsWKB(`location`))) END AS `location` FROM `places`" + )); + assert!(sql.contains("WHERE `id` > 7"), "sql: {sql}"); + assert!(sql.contains("ORDER BY `id` ASC LIMIT 1000"), "sql: {sql}"); + } + #[test] fn database_specific_boolean_export_literals() { let sqlserver_statements = build_export_insert_statements(BuildExportInsertStatementsOptions { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 8db33dc6b..4bf2df799 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -505,18 +505,59 @@ pub(crate) fn mysql_value_to_json(row: &mysql_async::Row, idx: usize) -> serde_j .unwrap_or(serde_json::Value::Null) } -fn decode_mysql_geometry(bytes: &[u8]) -> Option { +fn bytes_to_upper_hex(bytes: &[u8]) -> String { + const HEX: &[u8; 16] = b"0123456789ABCDEF"; + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + output.push(HEX[(byte >> 4) as usize] as char); + output.push(HEX[(byte & 0x0F) as usize] as char); + } + output +} + +// Native MySQL geometry values normally begin with a four-byte little-endian +// SRID prefix. A bare WKB value can have 0 or 1 at byte four as well, so only +// treat that prefix as present when the remaining bytes parse as WKB. +fn mysql_geometry_wkb_parts(bytes: &[u8]) -> Option<(&[u8], u32)> { if bytes.len() >= 5 && matches!(bytes[4], 0 | 1) { let prefix: [u8; 4] = bytes[..4].try_into().ok()?; - if let Some(mut geometry) = super::wkb::decode_wkb_geometry(&bytes[4..]) { - if geometry.srid.is_none() { - let srid = u32::from_le_bytes(prefix); - geometry.srid = (srid != 0).then_some(srid); - } - return Some(geometry); + if super::wkb::decode_wkb_geometry(&bytes[4..]).is_some() { + return Some((&bytes[4..], u32::from_le_bytes(prefix))); } } - super::wkb::decode_wkb_geometry(bytes) + super::wkb::decode_wkb_geometry(bytes).map(|_| (bytes, 0)) +} + +fn mysql_geometry_to_export_marker(bytes: &[u8]) -> Option { + let (wkb, srid) = mysql_geometry_wkb_parts(bytes)?; + Some(format!("DBX_WKB:{srid}:{}", bytes_to_upper_hex(wkb))) +} + +fn mysql_value_to_json_for_export( + row: &mysql_async::Row, + idx: usize, + spatial_as_wkb: bool, +) -> Result { + if !spatial_as_wkb + || !row.columns_ref().get(idx).is_some_and(|column| column.column_type() == ColumnType::MYSQL_TYPE_GEOMETRY) + { + return Ok(mysql_value_to_json(row, idx)); + } + let Some(bytes) = row_get::, _>(row, idx) else { + return Ok(mysql_value_to_json(row, idx)); + }; + mysql_geometry_to_export_marker(&bytes) + .map(serde_json::Value::String) + .ok_or_else(|| "Cannot export MySQL geometry value as WKB".to_string()) +} + +fn decode_mysql_geometry(bytes: &[u8]) -> Option { + let (wkb, srid) = mysql_geometry_wkb_parts(bytes)?; + let mut geometry = super::wkb::decode_wkb_geometry(wkb)?; + if geometry.srid.is_none() { + geometry.srid = (srid != 0).then_some(srid); + } + Some(geometry) } fn mysql_spatial_column_builder(columns: &[mysql_async::Column]) -> SpatialColumnBuilder { @@ -4011,7 +4052,7 @@ pub async fn stream_query_rows( mut on_row: impl FnMut(&[serde_json::Value]) -> Result<(), String>, ) -> Result { let mut conn = get_conn_with_health_check(pool).await?; - stream_query_result_on_conn(&mut conn, sql, bare, max_rows, dialect, cancelled, |item| { + stream_query_result_on_conn(&mut conn, sql, bare, max_rows, dialect, cancelled, false, |item| { if let MySqlQueryStreamItem::Row(row) = item { on_row(&row)?; } @@ -4027,17 +4068,18 @@ pub async fn stream_query_result_on_conn( max_rows: Option, dialect: MySqlQueryDialect, cancelled: &AtomicBool, + spatial_as_wkb: bool, mut on_item: impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, ) -> Result { let row_limit = max_rows.unwrap_or(usize::MAX); if bare || prefers_text_protocol_query(sql, dialect) { - stream_query_result_text(conn, sql, row_limit, cancelled, &mut on_item).await + stream_query_result_text(conn, sql, row_limit, cancelled, spatial_as_wkb, &mut on_item).await } else { - match stream_query_result_prepared(conn, sql, row_limit, cancelled, &mut on_item).await { + match stream_query_result_prepared(conn, sql, row_limit, cancelled, spatial_as_wkb, &mut on_item).await { Ok(rows) => Ok(rows), Err(err) if mysql_error_should_retry_with_text_protocol(&err) => { - stream_query_result_text(conn, sql, row_limit, cancelled, &mut on_item).await + stream_query_result_text(conn, sql, row_limit, cancelled, spatial_as_wkb, &mut on_item).await } Err(err) => Err(err), } @@ -4049,6 +4091,7 @@ async fn stream_query_result_text( sql: &str, row_limit: usize, cancelled: &AtomicBool, + spatial_as_wkb: bool, on_item: &mut impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, ) -> Result { let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?; @@ -4074,7 +4117,9 @@ async fn stream_query_result_text( break; } let row = row.map_err(|e| e.to_string())?; - let values: Vec = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect(); + let values: Vec = (0..row.len()) + .map(|i| mysql_value_to_json_for_export(&row, i, spatial_as_wkb)) + .collect::>()?; on_item(MySqlQueryStreamItem::Row(values))?; rows_exported += 1; } @@ -4087,6 +4132,7 @@ async fn stream_query_result_prepared( sql: &str, row_limit: usize, cancelled: &AtomicBool, + spatial_as_wkb: bool, on_item: &mut impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, ) -> Result { let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?; @@ -4112,7 +4158,9 @@ async fn stream_query_result_prepared( break; } let row = row.map_err(|e| e.to_string())?; - let values: Vec = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect(); + let values: Vec = (0..row.len()) + .map(|i| mysql_value_to_json_for_export(&row, i, spatial_as_wkb)) + .collect::>()?; on_item(MySqlQueryStreamItem::Row(values))?; rows_exported += 1; } @@ -5038,6 +5086,23 @@ mod tests { let decoded = decode_mysql_geometry(&raw).unwrap(); assert_eq!(decoded.wkt, "POINT(1 2)"); assert_eq!(decoded.srid, None); + assert_eq!( + mysql_geometry_to_export_marker(&raw).as_deref(), + Some("DBX_WKB:0:0101000000000000000000F03F0000000000000040") + ); + } + + #[test] + fn mysql_geometry_export_marker_strips_internal_srid_prefix() { + let mut raw = 4326_u32.to_le_bytes().to_vec(); + raw.extend_from_slice(&[ + 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0, 0x3f, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x40, + ]); + assert_eq!( + mysql_geometry_to_export_marker(&raw).as_deref(), + Some("DBX_WKB:4326:0101000000000000000000F03F0000000000000040") + ); } #[test] diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 2cd87ce56..b99515260 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -259,25 +259,27 @@ fn stream_export_was_cancelled(error: &str, token_cancelled: bool, export_cancel /// /// `column_types` are the types returned by the executed query (original column /// order). The request's overrides are expected to align 1:1 in the same order. -/// If the request provides fewer overrides than the result has columns the -/// extra columns are left untyped. Overrides that are `None` or empty are -/// treated as "infer from the query result". +/// Missing, `None`, or empty overrides infer only MySQL spatial types, which +/// preserves the historical literal formatting of other database types. fn sql_insert_column_types(request: &QueryResultExportRequest, column_types: &[String]) -> Vec> { - match request.export_column_types.as_ref() { - Some(overrides) => { - let mut result: Vec> = overrides - .iter() - .map(|t| match t { - Some(s) if !s.is_empty() => Some(s.clone()), - _ => None, + column_types + .iter() + .enumerate() + .map(|(index, inferred)| { + request + .export_column_types + .as_ref() + .and_then(|overrides| overrides.get(index)) + .and_then(|override_type| override_type.as_deref()) + .filter(|override_type| !override_type.is_empty()) + .map(str::to_string) + .or_else(|| { + (request.database_type == DatabaseType::Mysql + && crate::database_export::is_mysql_spatial_export_type(inferred)) + .then(|| inferred.clone()) }) - .collect(); - // Pad with None if fewer overrides than result columns - result.resize(column_types.len(), None); - result - } - None => vec![None; column_types.len()], - } + }) + .collect() } /// Bounded SQL INSERT writer with staged-file replacement safety. @@ -1197,6 +1199,7 @@ async fn try_export_mysql_query_result_stream( stream_row_limit, mysql_dialect, &export_cancelled, + format.eq_ignore_ascii_case("sql"), |item| { if export_cancelled.load(Ordering::SeqCst) || cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) @@ -1845,17 +1848,17 @@ mod tests { #[test] fn sql_insert_column_types_maps_request_types_to_option_vec() { let req = request("sql", None, None); - // No export_column_types set → every column becomes None + // Non-MySQL exports preserve their historical untyped behavior. let result = sql_insert_column_types(&req, &["int4".into(), "text".into()]); assert_eq!(result, vec![None, None]); - // Export_column_types provided → null becomes None, Some becomes Some + // Explicit non-empty overrides take precedence; missing values stay untyped. let mut req = req; req.export_column_types = Some(vec![Some("int4".into()), None, Some("jsonb".into())]); let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into()]); assert_eq!(result, vec![Some("int4".into()), None, Some("jsonb".into())]); - // Empty string in an override is treated as None + // Empty string in an override stays untyped for non-MySQL exports. req.export_column_types = Some(vec![Some("".into())]); let result = sql_insert_column_types(&req, &["int4".into()]); assert_eq!(result, vec![None]); @@ -1864,7 +1867,7 @@ mod tests { #[test] fn sql_insert_column_types_handles_partial_overrides_gracefully() { let req = request("sql", None, None); - // Fewer overrides than result columns → extra columns become None + // Fewer overrides than result columns → extra columns remain untyped. let mut req = req; req.export_column_types = Some(vec![Some("int4".into()), None]); let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into(), "bool".into()]); @@ -1885,7 +1888,7 @@ mod tests { #[test] fn sql_insert_column_types_handles_all_none_and_all_some() { let req = request("sql", None, None); - // All None + // All None remains untyped for non-MySQL exports. let mut req = req; req.export_column_types = Some(vec![None, None, None]); let result = sql_insert_column_types(&req, &["int4".into(), "text".into(), "json".into()]); @@ -1897,6 +1900,16 @@ mod tests { assert_eq!(result, vec![Some("int4".into()), Some("text".into()), Some("json".into())]); } + #[test] + fn sql_insert_column_types_infers_only_mysql_spatial_result_types() { + let mut req = request("sql", None, None); + req.database_type = DatabaseType::Mysql; + assert_eq!( + sql_insert_column_types(&req, &["int".into(), "geometry".into(), "varchar".into()]), + vec![None, Some("geometry".into()), None] + ); + } + #[test] fn sql_export_sql_file_is_initialized_only_for_sql_format() { // The sql_file variable is initialized at declaration only for "sql" format. diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index c712497b4..d91e952dc 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -140,10 +140,34 @@ fn resolve_requested_export_columns( (resolved_columns, resolved_column_types, resolved_primary_keys) } -fn requested_export_needs_column_extras(database_type: DatabaseType, format: &str) -> bool { +fn requested_mysql_sql_export_needs_column_metadata(database_type: DatabaseType, format: &str) -> bool { database_type == DatabaseType::Mysql && format.eq_ignore_ascii_case("sql") } +fn resolve_requested_export_column_types( + requested_columns: &[String], + requested_column_types: &[Option], + table_columns: &[crate::db::ColumnInfo], +) -> Vec> { + requested_columns + .iter() + .enumerate() + .map(|(index, requested)| { + requested_column_types + .get(index) + .cloned() + .flatten() + .filter(|column_type| !column_type.trim().is_empty()) + .or_else(|| { + table_columns + .iter() + .find(|column| column.name.eq_ignore_ascii_case(requested)) + .map(|column| column.data_type.clone()) + }) + }) + .collect() +} + fn resolve_requested_export_column_extras( requested_columns: &[String], table_columns: &[crate::db::ColumnInfo], @@ -213,13 +237,14 @@ fn table_page_sql( request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], + column_types: &[Option], primary_keys: &[String], use_keyset: bool, last_pk_values: &[Value], offset: u64, batch_size: usize, ) -> String { - if use_keyset { + let sql = if use_keyset { keyset_pagination_sql_with_identifier_quote( col_names, &request.table_name, @@ -243,6 +268,75 @@ fn table_page_sql( primary_keys, request.identifier_quote.as_deref(), ) + }; + replace_mysql_spatial_export_select_list(sql, request, db_type, col_names, column_types) +} + +fn mysql_spatial_export_column_expression(column: &str, identifier_quote: Option<&str>) -> String { + let quoted = crate::sql_dialect::quote_table_data_identifier(Some(DatabaseType::Mysql), column, identifier_quote); + format!( + "CASE WHEN {quoted} IS NULL THEN NULL ELSE CONCAT('DBX_WKB:', ST_SRID({quoted}), ':', HEX(ST_AsWKB({quoted}))) END AS {quoted}" + ) +} + +fn mysql_spatial_export_select_list( + request: &TableExportRequest, + db_type: &DatabaseType, + col_names: &[String], + column_types: &[Option], +) -> Option { + if *db_type != DatabaseType::Mysql || !request.format.eq_ignore_ascii_case("sql") { + return None; + } + let mut has_spatial_column = false; + let expressions = col_names + .iter() + .enumerate() + .map(|(index, column)| { + if column_types + .get(index) + .and_then(|column_type| column_type.as_deref()) + .is_some_and(crate::database_export::is_mysql_spatial_export_type) + { + has_spatial_column = true; + mysql_spatial_export_column_expression(column, request.identifier_quote.as_deref()) + } else { + crate::sql_dialect::quote_table_data_identifier( + Some(*db_type), + column, + request.identifier_quote.as_deref(), + ) + } + }) + .collect::>(); + has_spatial_column.then(|| expressions.join(", ")) +} + +fn replace_mysql_spatial_export_select_list( + sql: String, + request: &TableExportRequest, + db_type: &DatabaseType, + col_names: &[String], + column_types: &[Option], +) -> String { + let Some(replacement) = mysql_spatial_export_select_list(request, db_type, col_names, column_types) else { + return sql; + }; + let original = col_names + .iter() + .map(|column| { + crate::sql_dialect::quote_table_data_identifier(Some(*db_type), column, request.identifier_quote.as_deref()) + }) + .collect::>() + .join(", "); + let prefix = format!("SELECT {original}"); + if sql.starts_with(&prefix) { + format!("SELECT {replacement}{}", &sql[prefix.len()..]) + } else { + log::warn!( + "MySQL spatial table export could not replace its SELECT list; geometry columns will be exported as WKT" + ); + sql } } @@ -250,6 +344,7 @@ fn table_cursor_sql( request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], + column_types: &[Option], primary_keys: &[String], ) -> String { let full_table = crate::sql_dialect::table_data_qualified_table_name( @@ -258,13 +353,19 @@ fn table_cursor_sql( &request.table_name, request.identifier_quote.as_deref(), ); - let col_list = col_names - .iter() - .map(|column| { - crate::sql_dialect::quote_table_data_identifier(Some(*db_type), column, request.identifier_quote.as_deref()) - }) - .collect::>() - .join(", "); + let col_list = mysql_spatial_export_select_list(request, db_type, col_names, column_types).unwrap_or_else(|| { + col_names + .iter() + .map(|column| { + crate::sql_dialect::quote_table_data_identifier( + Some(*db_type), + column, + request.identifier_quote.as_deref(), + ) + }) + .collect::>() + .join(", ") + }); let predicate = crate::sql_dialect::normalize_where_input(request.where_input.as_deref()); let where_clause = if predicate.is_empty() { String::new() } else { format!(" WHERE ({predicate})") }; let order_by = request @@ -338,12 +439,13 @@ async fn execute_external_driver_export_page( request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], + column_types: &[Option], primary_keys: &[String], active_batch_size: usize, result_session_id: Option, cancel_token: CancellationToken, ) -> Result { - let sql = table_cursor_sql(request, db_type, col_names, primary_keys); + let sql = table_cursor_sql(request, db_type, col_names, column_types, primary_keys); let max_rows = request.row_limit.unwrap_or(i32::MAX as usize).min(i32::MAX as usize).max(1); let timeout_secs = table_export_query_timeout_secs(state, pool_key).await; execute_sql_statement_with_options( @@ -403,6 +505,7 @@ async fn fetch_table_export_batch( request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], + column_types: &[Option], primary_keys: &[String], use_keyset: bool, last_pk_values: &[Value], @@ -457,7 +560,7 @@ async fn fetch_table_export_batch( match table_export_cursor_kind(state, pool_key).await { Some(TableExportCursorKind::Agent) => { *table_read_attempted = true; - let sql = table_cursor_sql(request, db_type, col_names, primary_keys); + let sql = table_cursor_sql(request, db_type, col_names, column_types, primary_keys); let max_rows = request.row_limit.unwrap_or(i32::MAX as usize); let query_timeout = table_export_query_timeout_secs(state, pool_key).await; let params = AgentTableReadStartParams { @@ -498,6 +601,7 @@ async fn fetch_table_export_batch( request, db_type, col_names, + column_types, primary_keys, active_batch_size, None, @@ -553,6 +657,7 @@ async fn fetch_table_export_batch( request, db_type, col_names, + column_types, primary_keys, active_batch_size, Some(session_id.clone()), @@ -589,6 +694,7 @@ async fn fetch_table_export_batch( request, db_type, col_names, + column_types, primary_keys, use_keyset, last_pk_values, @@ -605,6 +711,7 @@ async fn fetch_paginated_table_export_batch( request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], + column_types: &[Option], primary_keys: &[String], use_keyset: bool, last_pk_values: &[Value], @@ -615,6 +722,7 @@ async fn fetch_paginated_table_export_batch( request, db_type, col_names, + column_types, primary_keys, use_keyset, last_pk_values, @@ -739,7 +847,7 @@ async fn try_export_native_table_stream( cancelled: Arc, cancel_token: CancellationToken, ) -> Result { - let sql = table_cursor_sql(request, db_type, col_names, primary_keys); + let sql = table_cursor_sql(request, db_type, col_names, column_types, primary_keys); let mut rows_exported = 0_u64; let progress_interval = batch_size.max(1) as u64; @@ -1169,25 +1277,29 @@ async fn export_table_data_core_inner( // directly, which avoids expensive metadata round-trips on JDBC drivers. let requested_columns = request.columns.as_ref().filter(|columns| !columns.is_empty()); let (col_names, column_types, column_extras, primary_keys) = if let Some(requested_columns) = requested_columns { - let (col_names, column_types, primary_keys) = resolve_requested_export_columns( + let (col_names, requested_column_types, primary_keys) = resolve_requested_export_columns( db_type, requested_columns, request.column_types.as_deref(), request.primary_keys.as_deref(), ); - let column_extras = if requested_export_needs_column_extras(db_type, &request.format) { - let table_columns = crate::schema::get_columns_core( - state, - &request.connection_id, - &request.database, - request.schema.as_deref().unwrap_or(""), - &request.table_name, - ) - .await?; - resolve_requested_export_column_extras(&col_names, &table_columns) - } else { - Vec::new() - }; + let (column_types, column_extras) = + if requested_mysql_sql_export_needs_column_metadata(db_type, &request.format) { + let table_columns = crate::schema::get_columns_core( + state, + &request.connection_id, + &request.database, + request.schema.as_deref().unwrap_or(""), + &request.table_name, + ) + .await?; + ( + resolve_requested_export_column_types(&col_names, &requested_column_types, &table_columns), + resolve_requested_export_column_extras(&col_names, &table_columns), + ) + } else { + (requested_column_types, Vec::new()) + }; (col_names, column_types, column_extras, primary_keys) } else { let columns = crate::schema::get_columns_core( @@ -1330,6 +1442,7 @@ async fn export_table_data_core_inner( request, &db_type, &col_names, + &column_types, &primary_keys, use_keyset, &last_pk_values, @@ -1418,6 +1531,7 @@ async fn export_table_data_core_inner( request, &db_type, &col_names, + &column_types, &primary_keys, use_keyset, &last_pk_values, @@ -1517,6 +1631,7 @@ async fn export_table_data_core_inner( request, &db_type, &col_names, + &column_types, &primary_keys, use_keyset, &last_pk_values, @@ -1610,6 +1725,7 @@ async fn export_table_data_core_inner( request, &db_type, &col_names, + &column_types, &primary_keys, use_keyset, &last_pk_values, @@ -1692,6 +1808,7 @@ async fn export_table_data_core_inner( request, &db_type, &col_names, + &column_types, &primary_keys, use_keyset, &last_pk_values, @@ -1773,6 +1890,7 @@ async fn export_table_data_core_inner( request, &db_type, &col_names, + &column_types, &primary_keys, use_keyset, &last_pk_values, @@ -2140,6 +2258,7 @@ mod tests { &request, &DatabaseType::Oracle, &[String::from("id"), String::from("status")], + &[], &[String::from("id")], ); @@ -2179,15 +2298,15 @@ mod tests { let primary_keys = vec!["id".to_string()]; assert_eq!( - table_cursor_sql(&request, &DatabaseType::Gaussdb, &columns, &primary_keys), + table_cursor_sql(&request, &DatabaseType::Gaussdb, &columns, &[], &primary_keys), "SELECT id, `DisplayName` FROM app_schema.`order` ORDER BY id ASC" ); assert_eq!( - table_page_sql(&request, &DatabaseType::Gaussdb, &columns, &primary_keys, false, &[], 100, 100), + table_page_sql(&request, &DatabaseType::Gaussdb, &columns, &[], &primary_keys, false, &[], 100, 100), "SELECT id, `DisplayName` FROM app_schema.`order` ORDER BY id LIMIT 100 OFFSET 100" ); assert_eq!( - table_page_sql(&request, &DatabaseType::Gaussdb, &columns, &primary_keys, true, &[json!(10)], 0, 100,), + table_page_sql(&request, &DatabaseType::Gaussdb, &columns, &[], &primary_keys, true, &[json!(10)], 0, 100,), "SELECT id, `DisplayName` FROM app_schema.`order` WHERE id > 10 ORDER BY id ASC LIMIT 100" ); assert_eq!( @@ -2203,6 +2322,49 @@ mod tests { ); } + #[test] + fn mysql_sql_table_export_selects_spatial_columns_as_wkb_markers() { + let request = TableExportRequest { + export_id: "export-mysql-spatial".to_string(), + connection_id: "conn-1".to_string(), + database: "app".to_string(), + schema: None, + identifier_quote: None, + table_name: "spatial_data".to_string(), + file_path: "spatial_data.sql".to_string(), + format: "sql".to_string(), + columns: None, + column_types: None, + primary_keys: None, + where_input: None, + order_by: None, + skip_count: true, + batch_size: Some(100), + row_limit: None, + date_time_format: None, + numeric_column_right_align: false, + column_comments: None, + }; + let columns = vec!["id".to_string(), "geom".to_string(), "name".to_string()]; + let column_types = vec![Some("int".to_string()), Some("geometry".to_string()), Some("varchar".to_string())]; + let primary_keys = vec!["id".to_string()]; + + let cursor_sql = table_cursor_sql(&request, &DatabaseType::Mysql, &columns, &column_types, &primary_keys); + assert!(cursor_sql.contains("ST_SRID(`geom`), ':', HEX(ST_AsWKB(`geom`))")); + assert!(cursor_sql.contains("AS `geom`")); + assert!(!cursor_sql.contains("SELECT `id`, `geom`, `name`")); + + let page_sql = + table_page_sql(&request, &DatabaseType::Mysql, &columns, &column_types, &primary_keys, false, &[], 0, 100); + assert!(page_sql.contains("ST_AsWKB(`geom`)")); + assert!(page_sql.contains("LIMIT 100 OFFSET 0")); + + let mut csv_request = request; + csv_request.format = "csv".to_string(); + let csv_sql = table_cursor_sql(&csv_request, &DatabaseType::Mysql, &columns, &column_types, &primary_keys); + assert_eq!(csv_sql, "SELECT `id`, `geom`, `name` FROM `spatial_data` ORDER BY `id` ASC"); + } + #[test] fn oracle_requested_export_columns_omit_synthetic_rowid_and_keep_metadata_aligned() { let columns = vec!["__DBX_ROWID".to_string(), "ID".to_string(), "NAME".to_string()]; @@ -2237,7 +2399,7 @@ mod tests { numeric_column_right_align: false, column_comments: None, }; - let sql = table_cursor_sql(&request, &DatabaseType::Oracle, &columns, &primary_keys); + let sql = table_cursor_sql(&request, &DatabaseType::Oracle, &columns, &[], &primary_keys); assert_eq!(sql, "SELECT \"ID\", \"NAME\" FROM \"APP\".\"USERS\""); let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { @@ -2268,26 +2430,36 @@ mod tests { } #[test] - fn requested_mysql_sql_export_resolves_generated_column_extras_only_for_sql() { + fn requested_mysql_sql_export_resolves_column_metadata_only_for_sql() { let table_columns = vec![ crate::db::ColumnInfo { name: "ID".to_string(), + data_type: "int".to_string(), extra: Some("auto_increment".to_string()), ..Default::default() }, crate::db::ColumnInfo { name: "virtual_total".to_string(), + data_type: "geometry".to_string(), extra: Some("VIRTUAL GENERATED".to_string()), ..Default::default() }, ]; let requested_columns = vec!["virtual_total".to_string(), "id".to_string(), "missing".to_string()]; - assert!(requested_export_needs_column_extras(DatabaseType::Mysql, "SQL")); + assert!(requested_mysql_sql_export_needs_column_metadata(DatabaseType::Mysql, "SQL")); for format in ["csv", "json", "xlsx"] { - assert!(!requested_export_needs_column_extras(DatabaseType::Mysql, format)); + assert!(!requested_mysql_sql_export_needs_column_metadata(DatabaseType::Mysql, format)); } - assert!(!requested_export_needs_column_extras(DatabaseType::Postgres, "sql")); + assert!(!requested_mysql_sql_export_needs_column_metadata(DatabaseType::Postgres, "sql")); + assert_eq!( + resolve_requested_export_column_types( + &requested_columns, + &[Some("".to_string()), Some("bigint".to_string())], + &table_columns, + ), + vec![Some("geometry".to_string()), Some("bigint".to_string()), None] + ); assert_eq!( resolve_requested_export_column_extras(&requested_columns, &table_columns), vec![Some("VIRTUAL GENERATED".to_string()), Some("auto_increment".to_string()), None] diff --git a/crates/dbx-core/tests/live_mysql_database_export.rs b/crates/dbx-core/tests/live_mysql_database_export.rs index d88f40361..a4a971de0 100644 --- a/crates/dbx-core/tests/live_mysql_database_export.rs +++ b/crates/dbx-core/tests/live_mysql_database_export.rs @@ -53,8 +53,8 @@ async fn live_mysql_database_export_restores_dependent_views() { for sql in [ format!("DROP DATABASE IF EXISTS `{database}`"), format!("CREATE DATABASE `{database}`"), - format!("CREATE TABLE `{database}`.`base_table` (id INT PRIMARY KEY)"), - format!("INSERT INTO `{database}`.`base_table` VALUES (7)"), + format!("CREATE TABLE `{database}`.`base_table` (id INT PRIMARY KEY, location POINT)"), + format!("INSERT INTO `{database}`.`base_table` VALUES (7, ST_GeomFromText('POINT(1 2)', 4326))"), format!("CREATE VIEW `{database}`.`z_view` AS SELECT id FROM `{database}`.`base_table`"), format!("CREATE VIEW `{database}`.`a_view` AS SELECT id FROM `{database}`.`z_view`"), ] { @@ -83,6 +83,9 @@ async fn live_mysql_database_export_restores_dependent_views() { let test_result = async { export_database_sql_core(&state, &export_request, |_| {}).await?; let exported = std::fs::read_to_string(&file_path).map_err(|error| error.to_string())?; + if !exported.contains("ST_GeomFromWKB(") { + return Err("MySQL geometry export did not use ST_GeomFromWKB".to_string()); + } let referenced_view = exported_view_position(&exported, "z_view") .ok_or_else(|| "exported z_view DDL was not found".to_string())?; let dependent_view = exported_view_position(&exported, "a_view") @@ -108,7 +111,16 @@ async fn live_mysql_database_export_restores_dependent_views() { let result = execute_sql_statement(&state, &connection_id, &database, "SELECT id FROM a_view", None, None).await?; - Ok::<_, String>((referenced_view, dependent_view, result)) + let spatial_result = execute_sql_statement( + &state, + &connection_id, + &database, + "SELECT ST_SRID(location), ST_Equals(location, ST_GeomFromText('POINT(1 2)', 4326)) FROM base_table", + None, + None, + ) + .await?; + Ok::<_, String>((referenced_view, dependent_view, result, spatial_result)) } .await; let cleanup = @@ -116,7 +128,8 @@ async fn live_mysql_database_export_restores_dependent_views() { cleanup.unwrap(); std::fs::remove_dir_all(dir).unwrap(); - let (referenced_view, dependent_view, result) = test_result.unwrap(); + let (referenced_view, dependent_view, result, spatial_result) = test_result.unwrap(); assert!(referenced_view < dependent_view); assert_eq!(result.rows, vec![vec![serde_json::json!("7")]]); + assert_eq!(spatial_result.rows, vec![vec![serde_json::json!(4326), serde_json::json!(1)]]); }