From e0d481cc2eae1d8d8eb675d7cb4d0d401569a1b1 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 15 Jun 2026 16:59:01 +0800 Subject: [PATCH] fix(postgres): handle tsvector columns --- .../src/components/sidebar/TreeItem.vue | 11 +++ crates/dbx-core/src/data_grid_sql.rs | 92 ++++++++++++++++- crates/dbx-core/src/database_export.rs | 54 ++++++++-- crates/dbx-core/src/db/postgres.rs | 93 +++++++++++++++++- .../dbx-core/tests/live_postgres_tsvector.rs | 98 +++++++++++++++++++ 5 files changed, 333 insertions(+), 15 deletions(-) create mode 100644 crates/dbx-core/tests/live_postgres_tsvector.rs diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 4eb25c6a7..f2694b49a 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -852,6 +852,9 @@ async function newQuery() { function isAutoGeneratedColumn(column: ColumnInfo, _dbType?: DatabaseType): boolean { const extra = (column.extra ?? "").toLowerCase().trim(); const colDefault = (column.column_default ?? "").toLowerCase().trim(); + const dataType = (column.data_type ?? "").trim(); + + if (typeLooksPostgresTextSearchVector(dataType)) return true; // MySQL / MariaDB: extra contains "auto_increment" if (extra.includes("auto_increment")) return true; @@ -919,6 +922,14 @@ function typeLooksArray(dataType: string): boolean { return /^(array|vector|_float\d*|_int\d*|_text|_varchar|_bool|_uuid)/i.test(dataType.trim()); } +function typeLooksPostgresTextSearchVector(dataType: string): boolean { + const normalized = dataType + .trim() + .replace(/^"+|"+$/g, "") + .toLowerCase(); + return normalized === "tsvector" || normalized.endsWith(".tsvector"); +} + /** * Return a type-appropriate placeholder value for a template column. * Inspects column.data_type and database type to produce the right diff --git a/crates/dbx-core/src/data_grid_sql.rs b/crates/dbx-core/src/data_grid_sql.rs index c8e29ddb1..b75552c0c 100644 --- a/crates/dbx-core/src/data_grid_sql.rs +++ b/crates/dbx-core/src/data_grid_sql.rs @@ -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 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, + 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, + 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, 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, 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!( diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 9302c03e0..a0e3d704d 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -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::>(); + 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::>() .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::>() @@ -244,6 +257,16 @@ pub fn build_export_insert_statements(options: BuildExportInsertStatementsOption Ok(statements) } +fn is_postgres_tsvector_export_column(database_type: Option, 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 { 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 { diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 2320517fe..d52961e51 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -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 { + 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 { + 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 { + 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, 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 { 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() diff --git a/crates/dbx-core/tests/live_postgres_tsvector.rs b/crates/dbx-core/tests/live_postgres_tsvector.rs new file mode 100644 index 000000000..e12ccd59e --- /dev/null +++ b/crates/dbx-core/tests/live_postgres_tsvector.rs @@ -0,0 +1,98 @@ +use std::time::Duration; + +use dbx_core::data_grid_sql::{ + build_data_grid_copy_insert_statement, DataGridColumnInfo, DataGridCopyInsertStatementOptions, DataGridTableMeta, +}; +use dbx_core::database_export::{build_export_insert_statements, BuildExportInsertStatementsOptions}; +use dbx_core::db::postgres; +use dbx_core::models::connection::DatabaseType; + +fn grid_column(column: &dbx_core::types::ColumnInfo) -> DataGridColumnInfo { + DataGridColumnInfo { + name: column.name.clone(), + data_type: column.data_type.clone(), + is_nullable: column.is_nullable, + is_primary_key: column.is_primary_key, + column_default: column.column_default.clone(), + extra: column.extra.clone(), + } +} + +#[tokio::test] +#[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] +async fn postgres_tsvector_generated_columns_are_readable_and_omitted_from_inserts() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = postgres::connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); + let schema = format!("dbx_tsvector_{}", std::process::id()); + let schema_ident = format!("\"{}\"", schema.replace('"', "\"\"")); + let table = format!("{schema_ident}.articles"); + + let _ = postgres::execute_query(&pool, &format!("DROP SCHEMA IF EXISTS {schema_ident} CASCADE")).await; + postgres::execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); + postgres::execute_query( + &pool, + &format!( + "CREATE TABLE {table} (\ + id integer PRIMARY KEY,\ + title text NOT NULL,\ + body text NOT NULL,\ + search_vector tsvector GENERATED ALWAYS AS \ + (to_tsvector('simple', coalesce(title, '') || ' ' || coalesce(body, ''))) STORED\ + )" + ), + ) + .await + .expect("create table"); + postgres::execute_query(&pool, &format!("INSERT INTO {table} (id, title, body) VALUES (1, 'Hello', 'World')")) + .await + .expect("insert row"); + + let result = + postgres::execute_query(&pool, &format!("SELECT id, title, body, search_vector FROM {table} ORDER BY id")) + .await + .expect("select row"); + assert_eq!(result.column_types, vec!["int4", "text", "text", "tsvector"]); + assert_eq!(result.rows[0][3].as_str(), Some("'hello':1 'world':2")); + + let columns = postgres::get_columns(&pool, &schema, "articles").await.expect("get columns"); + let search_vector = columns.iter().find(|column| column.name == "search_vector").expect("search_vector column"); + assert_eq!(search_vector.column_default, None); + assert!(search_vector.extra.as_deref().unwrap_or_default().contains("generated always as")); + + let table_meta = DataGridTableMeta { + schema: Some(schema.clone()), + table_name: "articles".to_string(), + primary_keys: vec!["id".to_string()], + columns: Some(columns.iter().map(grid_column).collect()), + }; + let copy_insert = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions { + database_type: Some(DatabaseType::Postgres), + table_meta: Some(table_meta), + columns: result.columns.clone(), + source_columns: None, + rows: result.rows.clone(), + exclude_primary_keys: false, + }) + .expect("copy insert"); + assert!(copy_insert.contains("\"id\", \"title\", \"body\"")); + assert!(!copy_insert.contains("search_vector")); + + let export_insert = build_export_insert_statements(BuildExportInsertStatementsOptions { + database_type: Some(DatabaseType::Postgres), + schema: Some(schema.clone()), + table_name: Some("articles".to_string()), + qualified_table_name: None, + columns: result.columns.clone(), + column_types: result.column_types.iter().map(|value| Some(value.clone())).collect(), + rows: result.rows.clone(), + batch_size: Some(10), + }) + .expect("export insert") + .join("\n"); + assert!(export_insert.contains("\"id\", \"title\", \"body\"")); + assert!(!export_insert.contains("search_vector")); + + postgres::execute_query(&pool, &format!("DROP SCHEMA IF EXISTS {schema_ident} CASCADE")) + .await + .expect("drop schema"); +}