github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(postgres): handle tsvector columns

Browse Source
This commit is contained in:
t8y2 committed 2026-06-15 16:59:01 +08:00
1 parent f2a0c65724
commit e0d481cc2e
5 files changed
+333 -15
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/data_grid_sql.rs
+87 -5
View File
@@ -272,7 +272,9 @@ pub fn build_data_grid_copy_insert_statement(options: DataGridCopyInsertStatemen
.iter()
.enumerate()
.filter_map(|(index, column)| Some((column.as_deref()?, index)))
.filter(|(column, _)| !is_oracle_row_id(options.database_type, Some(column)))
.filter(|(column, _)| {
!is_grid_insert_omitted_column(options.database_type, column_info_for(column_info, column), Some(column))
})
.collect();
let insert_columns: Vec<(&str, usize)> = insertable_columns
.iter()
@@ -411,6 +413,7 @@ fn validate_data_grid_save(options: &DataGridSaveStatementOptions) -> Option<Str
!column.is_nullable
&& column.column_default.is_none()
&& !is_auto_generated_column(column)
&& !is_non_identity_generated_column(Some(column))
&& !is_oracle_row_id(options.database_type, Some(&column.name))
})
.map(|column| normalize_column_name(&column.name))
@@ -510,7 +513,11 @@ fn build_data_grid_save_statements(options: &DataGridSaveStatementOptions) -> Ve
.iter()
.filter_map(|(column_index, value)| {
let column = save_columns.get(*column_index)?.as_deref()?;
if is_oracle_row_id(options.database_type, Some(column)) {
if is_grid_update_omitted_column(
options.database_type,
column_info_for(column_info, column),
Some(column),
) {
return None;
}
Some(format!(
@@ -557,7 +564,13 @@ fn build_data_grid_save_statements(options: &DataGridSaveStatementOptions) -> Ve
.iter()
.enumerate()
.filter_map(|(index, column)| Some((column.as_deref()?, row.get(index).unwrap_or(&Value::Null))))
.filter(|(column, _)| !is_oracle_row_id(options.database_type, Some(column)))
.filter(|(column, _)| {
!is_grid_insert_omitted_column(
options.database_type,
column_info_for(column_info, column),
Some(column),
)
})
.filter(|(_, value)| !value.is_null())
.collect();
if insert_pairs.is_empty() {
@@ -614,7 +627,13 @@ fn build_data_grid_rollback_statements(options: &DataGridSaveStatementOptions) -
.iter()
.enumerate()
.filter_map(|(index, column)| Some((column.as_deref()?, row.get(index).unwrap_or(&Value::Null))))
.filter(|(column, _)| !is_oracle_row_id(options.database_type, Some(column)))
.filter(|(column, _)| {
!is_grid_insert_omitted_column(
options.database_type,
column_info_for(column_info, column),
Some(column),
)
})
.collect();
let columns = insert_pairs
.iter()
@@ -648,7 +667,11 @@ fn build_data_grid_rollback_statements(options: &DataGridSaveStatementOptions) -
.iter()
.filter_map(|change @ (column_index, _)| {
let column = save_columns.get(*column_index)?.as_deref()?;
if is_oracle_row_id(options.database_type, Some(column)) {
if is_grid_update_omitted_column(
options.database_type,
column_info_for(column_info, column),
Some(column),
) {
return None;
}
Some((change, column))
@@ -1111,6 +1134,39 @@ fn is_auto_generated_column(column: &DataGridColumnInfo) -> bool {
.any(|part| matches!(part, "auto_increment" | "autoincrement" | "identity"))
}
fn is_grid_insert_omitted_column(
database_type: Option<DatabaseType>,
column_info: Option<&DataGridColumnInfo>,
name: Option<&str>,
) -> bool {
is_oracle_row_id(database_type, name)
|| is_postgres_tsvector_column(database_type, column_info)
|| is_non_identity_generated_column(column_info)
}
fn is_grid_update_omitted_column(
database_type: Option<DatabaseType>,
column_info: Option<&DataGridColumnInfo>,
name: Option<&str>,
) -> bool {
is_oracle_row_id(database_type, name) || is_non_identity_generated_column(column_info)
}
fn is_postgres_tsvector_column(database_type: Option<DatabaseType>, column_info: Option<&DataGridColumnInfo>) -> bool {
database_type == Some(DatabaseType::Postgres)
&& column_info.map(|column| is_postgres_tsvector_type(&column.data_type)).unwrap_or(false)
}
fn is_postgres_tsvector_type(data_type: &str) -> bool {
let normalized = data_type.trim().trim_matches('"').to_ascii_lowercase();
normalized == "tsvector" || normalized.ends_with(".tsvector")
}
fn is_non_identity_generated_column(column_info: Option<&DataGridColumnInfo>) -> bool {
let extra = column_info.and_then(|column| column.extra.as_deref()).unwrap_or("").to_ascii_lowercase();
extra.contains("generated always as") && !extra.contains("identity")
}
fn is_null_write_to_not_null_column(
database_type: Option<DatabaseType>,
not_null_columns: &[String],
@@ -1397,6 +1453,32 @@ mod tests {
);
}
#[test]
fn builds_copy_insert_statement_omits_postgres_tsvector_columns() {
let statement = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions {
database_type: Some(DatabaseType::Postgres),
table_meta: Some(DataGridTableMeta {
schema: Some("public".to_string()),
table_name: "articles".to_string(),
primary_keys: vec!["id".to_string()],
columns: Some(vec![
column("id", "integer", false, None),
column("title", "text", false, None),
column("search_vector", "tsvector", true, None),
]),
}),
columns: vec!["id".to_string(), "title".to_string(), "search_vector".to_string()],
source_columns: None,
rows: vec![vec![json!(1), json!("Hello"), json!("'hello':1A")]],
exclude_primary_keys: false,
});
assert_eq!(
statement.as_deref(),
Some("INSERT INTO \"public\".\"articles\" (\"id\", \"title\") VALUES (1, 'Hello');")
);
}
#[test]
fn builds_filter_conditions() {
assert_eq!(
crates/dbx-core/src/database_export.rs
+47 -7
View File
@@ -209,11 +209,24 @@ pub fn build_export_insert_statements(options: BuildExportInsertStatementsOption
options.table_name.as_deref(),
options.qualified_table_name.as_deref(),
)?;
let batch_size = options.batch_size.unwrap_or(DATABASE_EXPORT_INSERT_BATCH_SIZE).max(1);
let columns = options
let insert_columns = options
.columns
.iter()
.map(|column| quote_table_identifier(options.database_type, column))
.enumerate()
.filter(|(index, _)| {
!is_postgres_tsvector_export_column(
options.database_type,
options.column_types.get(*index).and_then(|value| value.as_deref()),
)
})
.collect::<Vec<_>>();
if insert_columns.is_empty() {
return Ok(Vec::new());
}
let batch_size = options.batch_size.unwrap_or(DATABASE_EXPORT_INSERT_BATCH_SIZE).max(1);
let columns = insert_columns
.iter()
.map(|(_, column)| quote_table_identifier(options.database_type, column))
.collect::<Vec<_>>()
.join(", ");
let mut statements = Vec::new();
@@ -222,14 +235,14 @@ pub fn build_export_insert_statements(options: BuildExportInsertStatementsOption
let values = rows
.iter()
.map(|row| {
let values = row
let values = insert_columns
.iter()
.enumerate()
.map(|(index, value)| {
.map(|(index, _)| {
let value = row.get(*index).unwrap_or(&Value::Null);
format_export_sql_literal_typed(
value,
options.database_type,
options.column_types.get(index).and_then(|value| value.as_deref()),
options.column_types.get(*index).and_then(|value| value.as_deref()),
)
})
.collect::<Vec<_>>()
@@ -244,6 +257,16 @@ pub fn build_export_insert_statements(options: BuildExportInsertStatementsOption
Ok(statements)
}
fn is_postgres_tsvector_export_column(database_type: Option<DatabaseType>, column_type: Option<&str>) -> bool {
database_type == Some(DatabaseType::Postgres)
&& column_type
.map(|column_type| {
let normalized = column_type.trim().trim_matches('"').to_ascii_lowercase();
normalized == "tsvector" || normalized.ends_with(".tsvector")
})
.unwrap_or(false)
}
pub fn build_export_sql_insert(options: BuildExportSqlInsertOptions) -> Result<String, String> {
build_export_insert_statements(options.insert).map(|statements| statements.join("\n"))
}
@@ -849,6 +872,23 @@ mod tests {
);
}
#[test]
fn postgres_tsvector_columns_are_omitted_from_sql_insert_export() {
let statements = build_export_insert_statements(BuildExportInsertStatementsOptions {
database_type: Some(DatabaseType::Postgres),
schema: Some("public".to_string()),
table_name: Some("articles".to_string()),
qualified_table_name: None,
columns: vec!["id".to_string(), "title".to_string(), "search_vector".to_string()],
column_types: vec![Some("integer".to_string()), Some("text".to_string()), Some("tsvector".to_string())],
rows: vec![vec![json!(1), json!("Hello"), json!("'hello':1A")]],
batch_size: Some(10),
})
.unwrap();
assert_eq!(statements, vec!["INSERT INTO \"public\".\"articles\" (\"id\", \"title\") VALUES (1, 'Hello');"]);
}
#[test]
fn builds_database_sql_export_with_ddl_before_data() {
let sql = build_database_sql_export(BuildDatabaseSqlExportOptions {
crates/dbx-core/src/db/postgres.rs
+90 -3
View File
@@ -538,6 +538,15 @@ fn pg_value_to_json(row: &Row, idx: usize, type_name: &str) -> serde_json::Value
.unwrap_or(serde_json::Value::Null);
}
if upper == "TSVECTOR" {
return row
.try_get::<_, PgRawBytes>(idx)
.ok()
.and_then(|raw| decode_tsvector_bytes(&raw.0))
.map(serde_json::Value::String)
.unwrap_or(serde_json::Value::Null);
}
if matches!(upper.as_str(), "OID" | "XID" | "CID") {
return pg_system_u32_to_json(row, idx).unwrap_or(serde_json::Value::Null);
}
@@ -602,6 +611,70 @@ fn pg_value_to_json(row: &Row, idx: usize, type_name: &str) -> serde_json::Value
.unwrap_or(serde_json::Value::Null)
}
fn decode_tsvector_bytes(raw: &[u8]) -> Option<String> {
let mut cursor = 0;
let count = read_i32_be(raw, &mut cursor)?;
if count < 0 {
return None;
}
let mut entries = Vec::with_capacity(count as usize);
for _ in 0..count {
let start = cursor;
while cursor < raw.len() && raw[cursor] != 0 {
cursor += 1;
}
if cursor >= raw.len() {
return None;
}
let lexeme = std::str::from_utf8(&raw[start..cursor]).ok()?;
cursor += 1;
let position_count = read_u16_be(raw, &mut cursor)? as usize;
let mut positions = Vec::with_capacity(position_count);
for _ in 0..position_count {
let encoded = read_u16_be(raw, &mut cursor)?;
let position = encoded & 0x3fff;
let weight = match encoded >> 14 {
3 => "A",
2 => "B",
1 => "C",
_ => "",
};
positions.push(format!("{position}{weight}"));
}
let mut entry = format!("'{}'", escape_tsvector_lexeme(lexeme));
if !positions.is_empty() {
entry.push(':');
entry.push_str(&positions.join(","));
}
entries.push(entry);
}
if cursor == raw.len() {
Some(entries.join(" "))
} else {
None
}
}
fn read_i32_be(raw: &[u8], cursor: &mut usize) -> Option<i32> {
let bytes: [u8; 4] = raw.get(*cursor..*cursor + 4)?.try_into().ok()?;
*cursor += 4;
Some(i32::from_be_bytes(bytes))
}
fn read_u16_be(raw: &[u8], cursor: &mut usize) -> Option<u16> {
let bytes: [u8; 2] = raw.get(*cursor..*cursor + 2)?.try_into().ok()?;
*cursor += 2;
Some(u16::from_be_bytes(bytes))
}
fn escape_tsvector_lexeme(value: &str) -> String {
value.replace('\\', "\\\\").replace('\'', "''")
}
fn pg_error_to_string(err: tokio_postgres::Error) -> String {
err.as_db_error().map(ToString::to_string).unwrap_or_else(|| err.to_string())
}
@@ -665,7 +738,7 @@ async fn execute_select_prepared(
Ok(QueryResult {
columns,
column_types: Vec::new(),
column_types,
column_sortables: Vec::new(),
rows: result_rows,
affected_rows: 0,
@@ -1284,7 +1357,7 @@ pub async fn list_schemas(pool: &Pool) -> Result<Vec<String>, String> {
const POSTGRES_COLUMNS_SQL: &str = "SELECT a.attname AS column_name, \
format_type(a.atttypid, a.atttypmod) AS full_type, \
NOT a.attnotnull AS is_nullable, \
pg_get_expr(ad.adbin, ad.adrelid) AS column_default, \
CASE WHEN a.attgenerated <> '' THEN NULL ELSE pg_get_expr(ad.adbin, ad.adrelid) END AS column_default, \
EXISTS ( \
SELECT 1 FROM pg_constraint co \
JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \
@@ -1295,7 +1368,11 @@ const POSTGRES_COLUMNS_SQL: &str = "SELECT a.attname AS column_name, \
CASE a.attidentity \
WHEN 'd' THEN 'generated by default as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \
WHEN 'a' THEN 'generated always as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \
ELSE NULL \
ELSE CASE a.attgenerated \
WHEN 's' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') stored' \
WHEN 'v' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') virtual' \
ELSE NULL \
END \
END AS column_extra, \
CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \
THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \
@@ -1977,6 +2054,16 @@ mod tests {
assert_eq!(raw.0, vec![0x01, 0xAB, 0xFF]);
}
#[test]
fn decodes_tsvector_binary_output() {
let raw = [
0, 0, 0, 2, b'b', b'a', b'c', b'k', b'\\', b's', b'l', b'a', b's', b'h', 0, 0, 1, 0x80, 0x03, b'o', b'\'',
b'c', b'l', b'o', b'c', b'k', 0, 0, 2, 0, 1, 0xc0, 0x02,
];
assert_eq!(decode_tsvector_bytes(&raw).as_deref(), Some("'back\\\\slash':3B 'o''clock':1,2A"));
}
fn decode_hex(hex: &str) -> Vec<u8> {
assert_eq!(hex.len() % 2, 0, "hex input must have an even number of chars");
(0..hex.len()).step_by(2).map(|idx| u8::from_str_radix(&hex[idx..idx + 2], 16).unwrap()).collect()