xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(mysql): correct geometry export lon/lat order via WKB round-trip

This commit is contained in:
amwps290 2026-08-06 23:15:59 +08:00 committed by GitHub
parent 00bbfad055
commit c835deb3f9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 553 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

232
crates/dbx-core/src/database_export.rs
View File

@ -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:<srid>:<hex>` 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<DatabaseType>,
column_type: Option<&str>,
) -> Option<String> {
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::<u32>().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<Vec<u8>> {
fn nibble(byte: u8) -> Option<u8> {
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<String> {
match value {
Value::Null => Some("NULL".to_string()),
@ -1376,8 +1445,59 @@ fn record_export_error<W: Write>(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::<Vec<_>>().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<String>], 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::<Vec<_>>()
.join(", ")
}
fn replace_database_export_select_list(
sql: String,
columns: &[String],
column_types: &[Option<String>],
db_type: &DatabaseType,
) -> String {
let original = columns.iter().map(|column| quote_identifier(column, db_type)).collect::<Vec<_>>().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<String>],
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 {

93
crates/dbx-core/src/db/mysql.rs
View File

@ -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<super::wkb::DecodedGeometry> {
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<String> {
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<serde_json::Value, String> {
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::<Vec<u8>, _>(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<super::wkb::DecodedGeometry> {
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<u64, String> {
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<usize>,
dialect: MySqlQueryDialect,
cancelled: &AtomicBool,
spatial_as_wkb: bool,
mut on_item: impl FnMut(MySqlQueryStreamItem) -> Result<(), String>,
) -> Result<u64, String> {
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<u64, String> {
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<serde_json::Value> = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect();
let values: Vec<serde_json::Value> = (0..row.len())
.map(|i| mysql_value_to_json_for_export(&row, i, spatial_as_wkb))
.collect::<Result<_, _>>()?;
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<u64, String> {
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<serde_json::Value> = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect();
let values: Vec<serde_json::Value> = (0..row.len())
.map(|i| mysql_value_to_json_for_export(&row, i, spatial_as_wkb))
.collect::<Result<_, _>>()?;
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]

57
crates/dbx-core/src/query_result_export.rs
View File

@ -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<Option<String>> {
match request.export_column_types.as_ref() {
Some(overrides) => {
let mut result: Vec<Option<String>> = 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.

240
crates/dbx-core/src/table_export.rs
View File

@ -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<String>],
table_columns: &[crate::db::ColumnInfo],
) -> Vec<Option<String>> {
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<String>],
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<String>],
) -> Option<String> {
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::<Vec<_>>();
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>],
) -> 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::<Vec<_>>()
.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<String>],
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::<Vec<_>>()
.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::<Vec<_>>()
.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<String>],
primary_keys: &[String],
active_batch_size: usize,
result_session_id: Option<String>,
cancel_token: CancellationToken,
) -> Result<QueryResult, String> {
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<String>],
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<String>],
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<AtomicBool>,
cancel_token: CancellationToken,
) -> Result<bool, String> {
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]

21
crates/dbx-core/tests/live_mysql_database_export.rs
View File

@ -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)]]);
}