1452 lines
52 KiB
Rust
1452 lines
52 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashSet;
|
|
use tokio::sync::RwLock;
|
|
|
|
use crate::connection::{AppState, PoolKind};
|
|
use crate::db;
|
|
use crate::models::connection::DatabaseType;
|
|
use crate::query::{agent_execute_query_params, QueryExecutionOptions};
|
|
|
|
static CANCELLED: std::sync::LazyLock<RwLock<HashSet<String>>> =
|
|
std::sync::LazyLock::new(|| RwLock::new(HashSet::new()));
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub enum TransferMode {
|
|
#[default]
|
|
Append,
|
|
Overwrite,
|
|
Upsert,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct TransferRequest {
|
|
pub transfer_id: String,
|
|
pub source_connection_id: String,
|
|
pub source_database: String,
|
|
pub source_schema: String,
|
|
pub target_connection_id: String,
|
|
pub target_database: String,
|
|
pub target_schema: String,
|
|
pub tables: Vec<String>,
|
|
pub create_table: bool,
|
|
#[serde(default)]
|
|
pub mode: TransferMode,
|
|
pub batch_size: usize,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct TransferProgress {
|
|
pub transfer_id: String,
|
|
pub table: String,
|
|
pub table_index: usize,
|
|
pub total_tables: usize,
|
|
pub rows_transferred: u64,
|
|
pub total_rows: Option<u64>,
|
|
pub status: TransferStatus,
|
|
pub error: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub enum TransferStatus {
|
|
Running,
|
|
TableDone,
|
|
Done,
|
|
Error,
|
|
Cancelled,
|
|
}
|
|
|
|
pub fn quote_identifier(name: &str, db_type: &DatabaseType) -> String {
|
|
match db_type {
|
|
DatabaseType::Mysql | DatabaseType::ClickHouse | DatabaseType::Doris | DatabaseType::StarRocks => {
|
|
format!("`{}`", name.replace('`', "``"))
|
|
}
|
|
DatabaseType::SqlServer => format!("[{}]", name.replace(']', "]]")),
|
|
_ => format!("\"{}\"", name.replace('"', "\"\"")),
|
|
}
|
|
}
|
|
|
|
pub fn qualified_table(table: &str, schema: &str, db_type: &DatabaseType) -> String {
|
|
let qt = quote_identifier(table, db_type);
|
|
if schema.is_empty() {
|
|
qt
|
|
} else {
|
|
format!("{}.{}", quote_identifier(schema, db_type), qt)
|
|
}
|
|
}
|
|
|
|
pub fn escape_value(val: &serde_json::Value, db_type: &DatabaseType) -> String {
|
|
escape_value_typed(val, db_type, None)
|
|
}
|
|
|
|
pub fn escape_value_typed(val: &serde_json::Value, db_type: &DatabaseType, column_type: Option<&str>) -> String {
|
|
match val {
|
|
serde_json::Value::Null => "NULL".to_string(),
|
|
serde_json::Value::Bool(b) => match db_type {
|
|
DatabaseType::Mysql
|
|
| DatabaseType::Sqlite
|
|
| DatabaseType::DuckDb
|
|
| DatabaseType::Doris
|
|
| DatabaseType::StarRocks => {
|
|
if *b {
|
|
"1".to_string()
|
|
} else {
|
|
"0".to_string()
|
|
}
|
|
}
|
|
_ => {
|
|
if *b {
|
|
"TRUE".to_string()
|
|
} else {
|
|
"FALSE".to_string()
|
|
}
|
|
}
|
|
},
|
|
serde_json::Value::Number(n) => n.to_string(),
|
|
serde_json::Value::String(s) => {
|
|
format!("'{}'", format_literal_string(s, db_type, column_type).replace('\\', "\\\\").replace('\'', "''"))
|
|
}
|
|
serde_json::Value::Array(arr) => format_pg_array_sql_literal(arr),
|
|
_ => {
|
|
let s = val.to_string();
|
|
format!("'{}'", s.replace('\\', "\\\\").replace('\'', "''"))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn format_pg_array_sql_literal(arr: &[serde_json::Value]) -> String {
|
|
if arr.is_empty() {
|
|
return "'{}'".to_string();
|
|
}
|
|
let elements: Vec<String> = arr.iter().map(format_pg_array_element).collect();
|
|
let inner = format!("{{{}}}", elements.join(","));
|
|
format!("'{}'", inner.replace('\\', "\\\\").replace('\'', "''"))
|
|
}
|
|
|
|
fn format_pg_array_element(val: &serde_json::Value) -> String {
|
|
match val {
|
|
serde_json::Value::Null => "NULL".to_string(),
|
|
serde_json::Value::Array(arr) => {
|
|
if arr.is_empty() {
|
|
return "{}".to_string();
|
|
}
|
|
let elements: Vec<String> = arr.iter().map(format_pg_array_element).collect();
|
|
format!("{{{}}}", elements.join(","))
|
|
}
|
|
serde_json::Value::String(s) => {
|
|
let escaped = s.replace('\\', "\\\\").replace('"', "\\\"");
|
|
format!("\"{}\"", escaped)
|
|
}
|
|
serde_json::Value::Number(n) => n.to_string(),
|
|
serde_json::Value::Bool(b) => {
|
|
if *b {
|
|
"true".to_string()
|
|
} else {
|
|
"false".to_string()
|
|
}
|
|
}
|
|
serde_json::Value::Object(o) => {
|
|
let json = serde_json::to_string(o).unwrap_or_default();
|
|
let escaped = json.replace('\\', "\\\\").replace('"', "\\\"");
|
|
format!("\"{}\"", escaped)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn format_literal_string(value: &str, db_type: &DatabaseType, column_type: Option<&str>) -> String {
|
|
if is_mysql_datetime_literal_database(db_type) && column_type.map(is_temporal_column_type).unwrap_or(true) {
|
|
normalize_mysql_temporal_literal(value, column_type).unwrap_or_else(|| value.to_string())
|
|
} else {
|
|
value.to_string()
|
|
}
|
|
}
|
|
|
|
fn is_mysql_datetime_literal_database(db_type: &DatabaseType) -> bool {
|
|
matches!(
|
|
db_type,
|
|
DatabaseType::Mysql
|
|
| DatabaseType::Doris
|
|
| DatabaseType::StarRocks
|
|
| DatabaseType::Goldendb
|
|
| DatabaseType::Sundb
|
|
)
|
|
}
|
|
|
|
fn normalize_mysql_temporal_literal(value: &str, column_type: Option<&str>) -> Option<String> {
|
|
let bytes = value.as_bytes();
|
|
if bytes.len() < 20 || !is_mysql_datetime_base(bytes) {
|
|
return None;
|
|
}
|
|
|
|
let rest = &value[19..];
|
|
let (fraction, offset) = if let Some(after_dot) = rest.strip_prefix('.') {
|
|
let digit_count = after_dot.bytes().take_while(|b| b.is_ascii_digit()).count();
|
|
if digit_count == 0 {
|
|
return None;
|
|
}
|
|
let fraction_len = 1 + digit_count;
|
|
(&rest[..fraction_len.min(7)], &rest[fraction_len..])
|
|
} else {
|
|
("", rest)
|
|
};
|
|
|
|
if !is_timezone_suffix(offset) {
|
|
return None;
|
|
}
|
|
|
|
match temporal_column_kind(column_type) {
|
|
Some("date") => Some(value[..10].to_string()),
|
|
Some("time") => Some(format!("{}{}", &value[11..19], fraction)),
|
|
_ => Some(format!("{} {}{}", &value[..10], &value[11..19], fraction)),
|
|
}
|
|
}
|
|
|
|
fn is_temporal_column_type(column_type: &str) -> bool {
|
|
temporal_column_kind(Some(column_type)).is_some()
|
|
}
|
|
|
|
fn temporal_column_kind(column_type: Option<&str>) -> Option<&'static str> {
|
|
let base = column_type?.trim().to_ascii_lowercase();
|
|
let base = base.split(['(', ':', ' ']).next().unwrap_or("");
|
|
match base {
|
|
"date" => Some("date"),
|
|
"time" => Some("time"),
|
|
"datetime" | "timestamp" => Some("datetime"),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
fn is_mysql_datetime_base(bytes: &[u8]) -> bool {
|
|
matches!(
|
|
bytes,
|
|
[
|
|
y0,
|
|
y1,
|
|
y2,
|
|
y3,
|
|
b'-',
|
|
m0,
|
|
m1,
|
|
b'-',
|
|
d0,
|
|
d1,
|
|
sep,
|
|
h0,
|
|
h1,
|
|
b':',
|
|
min0,
|
|
min1,
|
|
b':',
|
|
s0,
|
|
s1,
|
|
..
|
|
] if y0.is_ascii_digit()
|
|
&& y1.is_ascii_digit()
|
|
&& y2.is_ascii_digit()
|
|
&& y3.is_ascii_digit()
|
|
&& m0.is_ascii_digit()
|
|
&& m1.is_ascii_digit()
|
|
&& d0.is_ascii_digit()
|
|
&& d1.is_ascii_digit()
|
|
&& (*sep == b'T' || *sep == b' ')
|
|
&& h0.is_ascii_digit()
|
|
&& h1.is_ascii_digit()
|
|
&& min0.is_ascii_digit()
|
|
&& min1.is_ascii_digit()
|
|
&& s0.is_ascii_digit()
|
|
&& s1.is_ascii_digit()
|
|
)
|
|
}
|
|
|
|
fn is_timezone_suffix(value: &str) -> bool {
|
|
if value.eq_ignore_ascii_case("z") {
|
|
return true;
|
|
}
|
|
let bytes = value.as_bytes();
|
|
matches!(
|
|
bytes,
|
|
[sign, h0, h1, b':', m0, m1]
|
|
if (*sign == b'+' || *sign == b'-')
|
|
&& h0.is_ascii_digit()
|
|
&& h1.is_ascii_digit()
|
|
&& m0.is_ascii_digit()
|
|
&& m1.is_ascii_digit()
|
|
)
|
|
}
|
|
|
|
pub fn map_column_type(source_type: &str, _source_db: &DatabaseType, target_db: &DatabaseType) -> String {
|
|
let t = source_type.to_lowercase();
|
|
let base = t.split('(').next().unwrap_or(&t).trim();
|
|
|
|
match base {
|
|
"int" | "integer" | "int4" | "mediumint" => match target_db {
|
|
DatabaseType::Postgres => "INTEGER".into(),
|
|
DatabaseType::Mysql => "INT".into(),
|
|
DatabaseType::SqlServer => "INT".into(),
|
|
_ => "INTEGER".into(),
|
|
},
|
|
"bigint" | "int8" => "BIGINT".into(),
|
|
"smallint" | "int2" => "SMALLINT".into(),
|
|
"tinyint" => match target_db {
|
|
DatabaseType::Postgres => "SMALLINT".into(),
|
|
_ => "TINYINT".into(),
|
|
},
|
|
"serial" | "bigserial" | "smallserial" => match target_db {
|
|
DatabaseType::Postgres => source_type.to_uppercase(),
|
|
DatabaseType::Mysql => "BIGINT AUTO_INCREMENT".into(),
|
|
_ => "INTEGER".into(),
|
|
},
|
|
"float" | "float4" | "real" => match target_db {
|
|
DatabaseType::Postgres => "REAL".into(),
|
|
_ => "FLOAT".into(),
|
|
},
|
|
"double" | "double precision" | "float8" => match target_db {
|
|
DatabaseType::Postgres => "DOUBLE PRECISION".into(),
|
|
_ => "DOUBLE".into(),
|
|
},
|
|
"decimal" | "numeric" | "number" => {
|
|
if t.contains('(') {
|
|
match target_db {
|
|
DatabaseType::Mysql | DatabaseType::Postgres | DatabaseType::SqlServer | DatabaseType::Oracle => {
|
|
format!("DECIMAL{}", &t[t.find('(').unwrap()..])
|
|
}
|
|
_ => "NUMERIC".into(),
|
|
}
|
|
} else {
|
|
"NUMERIC".into()
|
|
}
|
|
}
|
|
"varchar" | "nvarchar" | "character varying" | "varchar2" => {
|
|
if t.contains('(') {
|
|
let len_part = &t[t.find('(').unwrap()..];
|
|
match target_db {
|
|
DatabaseType::Postgres => format!("VARCHAR{len_part}"),
|
|
DatabaseType::Mysql => format!("VARCHAR{len_part}"),
|
|
DatabaseType::SqlServer => format!("NVARCHAR{len_part}"),
|
|
_ => format!("VARCHAR{len_part}"),
|
|
}
|
|
} else {
|
|
"VARCHAR(255)".into()
|
|
}
|
|
}
|
|
"char" | "nchar" | "character" => {
|
|
if t.contains('(') {
|
|
let len_part = &t[t.find('(').unwrap()..];
|
|
format!("CHAR{len_part}")
|
|
} else {
|
|
"CHAR(1)".into()
|
|
}
|
|
}
|
|
"text" | "longtext" | "mediumtext" | "tinytext" | "clob" | "ntext" => "TEXT".into(),
|
|
"bool" | "boolean" => match target_db {
|
|
DatabaseType::Mysql => "TINYINT(1)".into(),
|
|
DatabaseType::SqlServer => "BIT".into(),
|
|
_ => "BOOLEAN".into(),
|
|
},
|
|
"date" => "DATE".into(),
|
|
"time" => "TIME".into(),
|
|
"datetime" => match target_db {
|
|
DatabaseType::Postgres => "TIMESTAMP".into(),
|
|
_ => "DATETIME".into(),
|
|
},
|
|
"timestamp" | "timestamptz" | "timestamp with time zone" | "timestamp without time zone" => match target_db {
|
|
DatabaseType::Mysql => "DATETIME".into(),
|
|
DatabaseType::SqlServer => "DATETIME2".into(),
|
|
_ => "TIMESTAMP".into(),
|
|
},
|
|
"blob" | "longblob" | "mediumblob" | "tinyblob" | "binary" | "varbinary" | "image" => match target_db {
|
|
DatabaseType::Postgres => "BYTEA".into(),
|
|
DatabaseType::Mysql => "BLOB".into(),
|
|
DatabaseType::SqlServer => "VARBINARY(MAX)".into(),
|
|
_ => "BLOB".into(),
|
|
},
|
|
"bytea" => match target_db {
|
|
DatabaseType::Postgres => "BYTEA".into(),
|
|
DatabaseType::Mysql => "BLOB".into(),
|
|
_ => "BLOB".into(),
|
|
},
|
|
"json" | "jsonb" => match target_db {
|
|
DatabaseType::Postgres => "JSONB".into(),
|
|
DatabaseType::Mysql => "JSON".into(),
|
|
_ => "TEXT".into(),
|
|
},
|
|
"uuid" => match target_db {
|
|
DatabaseType::Postgres => "UUID".into(),
|
|
_ => "VARCHAR(36)".into(),
|
|
},
|
|
"bit" => match target_db {
|
|
DatabaseType::Postgres => "BOOLEAN".into(),
|
|
_ => "BIT".into(),
|
|
},
|
|
_ => "TEXT".into(),
|
|
}
|
|
}
|
|
|
|
fn mysql_type_needs_key_prefix(mapped_type: &str) -> bool {
|
|
let base = mapped_type.split('(').next().unwrap_or(mapped_type).trim().to_ascii_lowercase();
|
|
matches!(
|
|
base.as_str(),
|
|
"text" | "tinytext" | "mediumtext" | "longtext" | "blob" | "tinyblob" | "mediumblob" | "longblob"
|
|
)
|
|
}
|
|
|
|
pub fn generate_create_table_ddl(
|
|
columns: &[db::ColumnInfo],
|
|
table: &str,
|
|
schema: &str,
|
|
target_db: &DatabaseType,
|
|
source_db: &DatabaseType,
|
|
table_comment: Option<&str>,
|
|
) -> String {
|
|
let full_table = qualified_table(table, schema, target_db);
|
|
|
|
let is_mysql_family = matches!(
|
|
target_db,
|
|
DatabaseType::Mysql
|
|
| DatabaseType::Doris
|
|
| DatabaseType::StarRocks
|
|
| DatabaseType::Goldendb
|
|
| DatabaseType::Sundb
|
|
);
|
|
|
|
let mut col_lines = Vec::with_capacity(columns.len());
|
|
for c in columns {
|
|
col_lines.push({
|
|
let mapped_type = map_column_type(&c.data_type, source_db, target_db);
|
|
let mut line = format!(" {} {}", quote_identifier(&c.name, target_db), mapped_type);
|
|
if !c.is_nullable {
|
|
line.push_str(" NOT NULL");
|
|
}
|
|
if is_mysql_family {
|
|
if let Some(ref comment) = c.comment {
|
|
let trimmed = comment.trim();
|
|
if !trimmed.is_empty() {
|
|
line.push_str(&format!(" COMMENT '{}'", trimmed.replace('\'', "''")));
|
|
}
|
|
}
|
|
}
|
|
line
|
|
});
|
|
}
|
|
|
|
let mut pks = Vec::new();
|
|
pks.reserve(columns.iter().filter(|c| c.is_primary_key).count());
|
|
for c in columns {
|
|
if c.is_primary_key {
|
|
let qname = quote_identifier(&c.name, target_db);
|
|
if is_mysql_family {
|
|
let mapped = map_column_type(&c.data_type, source_db, target_db);
|
|
if mysql_type_needs_key_prefix(&mapped) {
|
|
pks.push(format!("{qname}(255)"));
|
|
continue;
|
|
}
|
|
}
|
|
pks.push(qname);
|
|
}
|
|
}
|
|
|
|
let mut ddl = match target_db {
|
|
DatabaseType::SqlServer => {
|
|
format!("IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '{table}')\n")
|
|
}
|
|
_ => String::new(),
|
|
};
|
|
|
|
let create_prefix = match target_db {
|
|
DatabaseType::SqlServer => "CREATE TABLE",
|
|
_ => "CREATE TABLE IF NOT EXISTS",
|
|
};
|
|
|
|
ddl.push_str(&format!("{create_prefix} {full_table} (\n"));
|
|
ddl.push_str(&col_lines.join(",\n"));
|
|
|
|
if !pks.is_empty() {
|
|
ddl.push_str(&format!(",\n PRIMARY KEY ({})", pks.join(", ")));
|
|
}
|
|
|
|
ddl.push_str("\n)");
|
|
|
|
if is_mysql_family {
|
|
if let Some(ref comment) = table_comment {
|
|
let trimmed = comment.trim();
|
|
if !trimmed.is_empty() {
|
|
ddl.push_str(&format!(" COMMENT='{}'", trimmed.replace('\'', "''")));
|
|
}
|
|
}
|
|
}
|
|
|
|
if matches!(target_db, DatabaseType::ClickHouse) {
|
|
ddl.push_str(" ENGINE = MergeTree() ORDER BY tuple()");
|
|
}
|
|
|
|
ddl
|
|
}
|
|
|
|
/// Generate COMMENT ON COLUMN / ALTER TABLE COMMENT COLUMN / COMMENT ON TABLE
|
|
/// statements for databases that don't support inline comments in CREATE TABLE.
|
|
/// MySQL family uses inline syntax (handled in generate_create_table_ddl).
|
|
pub fn generate_comment_ddl(
|
|
columns: &[db::ColumnInfo],
|
|
table: &str,
|
|
schema: &str,
|
|
target_db: &DatabaseType,
|
|
table_comment: Option<&str>,
|
|
) -> Vec<String> {
|
|
if !matches!(target_db, DatabaseType::Postgres | DatabaseType::Oracle | DatabaseType::ClickHouse) {
|
|
return Vec::new();
|
|
}
|
|
|
|
let full_table = qualified_table(table, schema, target_db);
|
|
let mut statements = Vec::new();
|
|
|
|
// Table-level comment first (PostgreSQL/Oracle only; ClickHouse doesn't support COMMENT ON TABLE)
|
|
if matches!(target_db, DatabaseType::Postgres | DatabaseType::Oracle) {
|
|
if let Some(ref comment) = table_comment {
|
|
let trimmed = comment.trim();
|
|
if !trimmed.is_empty() {
|
|
let escaped = trimmed.replace('\'', "''");
|
|
statements.push(format!("COMMENT ON TABLE {full_table} IS '{escaped}'"));
|
|
}
|
|
}
|
|
}
|
|
|
|
for c in columns {
|
|
if let Some(ref comment) = c.comment {
|
|
let trimmed = comment.trim();
|
|
if trimmed.is_empty() {
|
|
continue;
|
|
}
|
|
let escaped = trimmed.replace('\'', "''");
|
|
let qcol = quote_identifier(&c.name, target_db);
|
|
|
|
match target_db {
|
|
DatabaseType::Postgres | DatabaseType::Oracle => {
|
|
statements.push(format!("COMMENT ON COLUMN {full_table}.{qcol} IS '{escaped}'"));
|
|
}
|
|
DatabaseType::ClickHouse => {
|
|
statements.push(format!("ALTER TABLE {full_table} COMMENT COLUMN {qcol} '{escaped}'"));
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
statements
|
|
}
|
|
|
|
pub fn generate_insert(
|
|
columns: &[String],
|
|
rows: &[Vec<serde_json::Value>],
|
|
table: &str,
|
|
schema: &str,
|
|
db_type: &DatabaseType,
|
|
) -> String {
|
|
generate_insert_typed(columns, &vec![None; columns.len()], rows, table, schema, db_type)
|
|
}
|
|
|
|
pub fn generate_insert_typed(
|
|
columns: &[String],
|
|
column_types: &[Option<String>],
|
|
rows: &[Vec<serde_json::Value>],
|
|
table: &str,
|
|
schema: &str,
|
|
db_type: &DatabaseType,
|
|
) -> String {
|
|
if rows.is_empty() {
|
|
return String::new();
|
|
}
|
|
|
|
let full_table = qualified_table(table, schema, db_type);
|
|
let col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
|
|
let value_rows = value_rows_sql(rows, column_types, db_type);
|
|
|
|
format!("INSERT INTO {full_table} ({col_list}) VALUES\n{}", value_rows.join(",\n"))
|
|
}
|
|
|
|
fn value_rows_sql(
|
|
rows: &[Vec<serde_json::Value>],
|
|
column_types: &[Option<String>],
|
|
db_type: &DatabaseType,
|
|
) -> Vec<String> {
|
|
let mut out = Vec::with_capacity(rows.len());
|
|
for row in rows {
|
|
let mut vals = Vec::with_capacity(row.len());
|
|
for (index, v) in row.iter().enumerate() {
|
|
vals.push(escape_value_typed(v, db_type, column_types.get(index).and_then(|value| value.as_deref())));
|
|
}
|
|
out.push(format!("({})", vals.join(", ")));
|
|
}
|
|
out
|
|
}
|
|
|
|
pub fn generate_upsert(
|
|
columns: &[String],
|
|
rows: &[Vec<serde_json::Value>],
|
|
table: &str,
|
|
schema: &str,
|
|
db_type: &DatabaseType,
|
|
pk_columns: &[String],
|
|
) -> String {
|
|
generate_upsert_typed(columns, &vec![None; columns.len()], rows, table, schema, db_type, pk_columns)
|
|
}
|
|
|
|
pub fn generate_upsert_typed(
|
|
columns: &[String],
|
|
column_types: &[Option<String>],
|
|
rows: &[Vec<serde_json::Value>],
|
|
table: &str,
|
|
schema: &str,
|
|
db_type: &DatabaseType,
|
|
pk_columns: &[String],
|
|
) -> String {
|
|
if rows.is_empty() || pk_columns.is_empty() {
|
|
return String::new();
|
|
}
|
|
|
|
let full_table = qualified_table(table, schema, db_type);
|
|
let col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
|
|
let value_rows = value_rows_sql(rows, column_types, db_type);
|
|
|
|
let mut non_pk_columns = Vec::with_capacity(columns.len().saturating_sub(pk_columns.len()));
|
|
for c in columns {
|
|
if !pk_columns.contains(c) {
|
|
non_pk_columns.push(c);
|
|
}
|
|
}
|
|
|
|
match db_type {
|
|
DatabaseType::Postgres | DatabaseType::Sqlite | DatabaseType::DuckDb => {
|
|
let pk_list = pk_columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
let mut sql = format!("INSERT INTO {full_table} ({col_list}) VALUES\n{}", value_rows.join(",\n"));
|
|
if non_pk_columns.is_empty() {
|
|
sql.push_str(&format!("\nON CONFLICT ({pk_list}) DO NOTHING"));
|
|
} else {
|
|
let update_set = non_pk_columns
|
|
.iter()
|
|
.map(|c| {
|
|
let qc = quote_identifier(c, db_type);
|
|
format!("{qc} = EXCLUDED.{qc}")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(", ");
|
|
sql.push_str(&format!("\nON CONFLICT ({pk_list}) DO UPDATE SET {update_set}"));
|
|
}
|
|
sql
|
|
}
|
|
DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks => {
|
|
let mut sql = format!("INSERT INTO {full_table} ({col_list}) VALUES\n{}", value_rows.join(",\n"));
|
|
if non_pk_columns.is_empty() {
|
|
sql.push_str("\nON DUPLICATE KEY UPDATE ");
|
|
let first_pk = quote_identifier(&pk_columns[0], db_type);
|
|
sql.push_str(&format!("{first_pk} = {first_pk}"));
|
|
} else {
|
|
let update_set = non_pk_columns
|
|
.iter()
|
|
.map(|c| {
|
|
let qc = quote_identifier(c, db_type);
|
|
format!("{qc} = VALUES({qc})")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(", ");
|
|
sql.push_str(&format!("\nON DUPLICATE KEY UPDATE {update_set}"));
|
|
}
|
|
sql
|
|
}
|
|
DatabaseType::SqlServer => {
|
|
let src_col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
let on_clause = pk_columns
|
|
.iter()
|
|
.map(|c| {
|
|
let qc = quote_identifier(c, db_type);
|
|
format!("target.{qc} = src.{qc}")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(" AND ");
|
|
|
|
let mut sql = format!(
|
|
"MERGE INTO {full_table} AS target USING (VALUES\n{}\n) AS src ({src_col_list}) ON {on_clause}",
|
|
value_rows.join(",\n")
|
|
);
|
|
|
|
if !non_pk_columns.is_empty() {
|
|
let update_set = non_pk_columns
|
|
.iter()
|
|
.map(|c| {
|
|
let qc = quote_identifier(c, db_type);
|
|
format!("target.{qc} = src.{qc}")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(", ");
|
|
sql.push_str(&format!("\nWHEN MATCHED THEN UPDATE SET {update_set}"));
|
|
}
|
|
|
|
let insert_cols = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
let insert_vals =
|
|
columns.iter().map(|c| format!("src.{}", quote_identifier(c, db_type))).collect::<Vec<_>>().join(", ");
|
|
sql.push_str(&format!("\nWHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({insert_vals});"));
|
|
sql
|
|
}
|
|
DatabaseType::Oracle => {
|
|
let mut using_rows = Vec::with_capacity(rows.len());
|
|
for row in rows {
|
|
let mut vals = Vec::with_capacity(row.len().min(columns.len()));
|
|
for (index, (v, c)) in row.iter().zip(columns.iter()).enumerate() {
|
|
vals.push(format!(
|
|
"{} AS {}",
|
|
escape_value_typed(v, db_type, column_types.get(index).and_then(|value| value.as_deref())),
|
|
quote_identifier(c, db_type)
|
|
));
|
|
}
|
|
using_rows.push(format!("SELECT {} FROM dual", vals.join(", ")));
|
|
}
|
|
|
|
let on_clause = pk_columns
|
|
.iter()
|
|
.map(|c| {
|
|
let qc = quote_identifier(c, db_type);
|
|
format!("t.{qc} = s.{qc}")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(" AND ");
|
|
|
|
let mut sql =
|
|
format!("MERGE INTO {full_table} t USING ({}) s ON ({on_clause})", using_rows.join(" UNION ALL "));
|
|
|
|
if !non_pk_columns.is_empty() {
|
|
let update_set = non_pk_columns
|
|
.iter()
|
|
.map(|c| {
|
|
let qc = quote_identifier(c, db_type);
|
|
format!("t.{qc} = s.{qc}")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(", ");
|
|
sql.push_str(&format!("\nWHEN MATCHED THEN UPDATE SET {update_set}"));
|
|
}
|
|
|
|
let insert_cols = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
let insert_vals =
|
|
columns.iter().map(|c| format!("s.{}", quote_identifier(c, db_type))).collect::<Vec<_>>().join(", ");
|
|
sql.push_str(&format!("\nWHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({insert_vals})"));
|
|
sql
|
|
}
|
|
_ => generate_insert_typed(columns, column_types, rows, table, schema, db_type),
|
|
}
|
|
}
|
|
|
|
pub fn pagination_sql(
|
|
columns: &[String],
|
|
table: &str,
|
|
schema: &str,
|
|
db_type: &DatabaseType,
|
|
offset: u64,
|
|
limit: usize,
|
|
) -> String {
|
|
let full_table = qualified_table(table, schema, db_type);
|
|
let col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
|
|
|
|
match db_type {
|
|
DatabaseType::SqlServer | DatabaseType::Oracle => {
|
|
format!(
|
|
"SELECT {col_list} FROM {full_table} ORDER BY (SELECT NULL) OFFSET {offset} ROWS FETCH NEXT {limit} ROWS ONLY"
|
|
)
|
|
}
|
|
_ => {
|
|
format!("SELECT {col_list} FROM {full_table} LIMIT {limit} OFFSET {offset}")
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn count_sql(table: &str, schema: &str, db_type: &DatabaseType) -> String {
|
|
let full_table = qualified_table(table, schema, db_type);
|
|
format!("SELECT COUNT(*) FROM {full_table}")
|
|
}
|
|
|
|
pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result<db::QueryResult, String> {
|
|
let connections = state.connections.read().await;
|
|
let pool = connections.get(pool_key).ok_or("Connection not found")?;
|
|
|
|
match pool {
|
|
PoolKind::Mysql(p, mode) => {
|
|
let p = p.clone();
|
|
let bare = *mode == crate::connection::MysqlMode::Bare;
|
|
drop(connections);
|
|
db::mysql::execute_query(&p, sql, bare).await
|
|
}
|
|
PoolKind::Postgres(p) => {
|
|
let p = p.clone();
|
|
drop(connections);
|
|
db::postgres::execute_query(&p, sql).await
|
|
}
|
|
PoolKind::Sqlite(p) => {
|
|
let p = p.clone();
|
|
drop(connections);
|
|
db::sqlite::execute_query(&p, sql).await
|
|
}
|
|
PoolKind::ClickHouse(client) => {
|
|
let client = client.clone();
|
|
let database = database_from_pool_key(pool_key).unwrap_or("default").to_string();
|
|
drop(connections);
|
|
db::clickhouse_driver::execute_query(&client, &database, sql).await
|
|
}
|
|
PoolKind::SqlServer(client) => {
|
|
let client = client.clone();
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
db::sqlserver::execute_query(&mut client, sql).await
|
|
}
|
|
PoolKind::Agent(client) => {
|
|
let client = client.clone();
|
|
let database = database_from_pool_key(pool_key).map(str::to_string);
|
|
let sql = sql.to_string();
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
let params = agent_execute_query_params(
|
|
&sql,
|
|
database.as_deref(),
|
|
None,
|
|
QueryExecutionOptions { max_rows: None, ..QueryExecutionOptions::default() },
|
|
);
|
|
client.execute_query(params).await
|
|
}
|
|
PoolKind::DuckDb(con) => {
|
|
let con = con.clone();
|
|
let sql = sql.to_string();
|
|
drop(connections);
|
|
tokio::task::spawn_blocking(move || {
|
|
let con = con.lock().map_err(|e| e.to_string())?;
|
|
let start = std::time::Instant::now();
|
|
let trimmed = sql.trim().to_uppercase();
|
|
if trimmed.starts_with("SELECT")
|
|
|| trimmed.starts_with("SHOW")
|
|
|| trimmed.starts_with("DESCRIBE")
|
|
|| trimmed.starts_with("WITH")
|
|
|| trimmed.starts_with("PRAGMA")
|
|
{
|
|
let mut stmt = con.prepare(&sql).map_err(|e| e.to_string())?;
|
|
let mut rows = stmt.query([]).map_err(|e| e.to_string())?;
|
|
let stmt_ref = rows.as_ref().ok_or("DuckDB statement unavailable")?;
|
|
let col_count = stmt_ref.column_count();
|
|
let columns: Vec<String> = (0..col_count)
|
|
.map(|i| stmt_ref.column_name(i).map(|s| s.to_string()).unwrap_or_else(|_| "?".to_string()))
|
|
.collect();
|
|
let mut result_rows = Vec::new();
|
|
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
|
|
let vals: Vec<serde_json::Value> = (0..col_count)
|
|
.map(|i| {
|
|
row.get::<_, String>(i)
|
|
.map(serde_json::Value::String)
|
|
.or_else(|_| row.get::<_, i64>(i).map(|v| serde_json::Value::Number(v.into())))
|
|
.or_else(|_| {
|
|
row.get::<_, f64>(i).map(|v| {
|
|
serde_json::Number::from_f64(v)
|
|
.map(serde_json::Value::Number)
|
|
.unwrap_or(serde_json::Value::Null)
|
|
})
|
|
})
|
|
.or_else(|_| row.get::<_, bool>(i).map(serde_json::Value::Bool))
|
|
.unwrap_or(serde_json::Value::Null)
|
|
})
|
|
.collect();
|
|
result_rows.push(vals);
|
|
}
|
|
Ok(db::QueryResult {
|
|
columns,
|
|
rows: result_rows,
|
|
affected_rows: 0,
|
|
execution_time_ms: start.elapsed().as_millis(),
|
|
truncated: false,
|
|
session_id: None,
|
|
has_more: false,
|
|
})
|
|
} else {
|
|
let affected = con.execute(&sql, []).map_err(|e| e.to_string())?;
|
|
Ok(db::QueryResult {
|
|
columns: vec![],
|
|
rows: vec![],
|
|
affected_rows: affected as u64,
|
|
execution_time_ms: start.elapsed().as_millis(),
|
|
truncated: false,
|
|
session_id: None,
|
|
has_more: false,
|
|
})
|
|
}
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?
|
|
}
|
|
PoolKind::ExternalTabular(ext_pool) => {
|
|
let con = ext_pool.cache.clone();
|
|
let sql = sql.to_string();
|
|
drop(connections);
|
|
tokio::task::spawn_blocking(move || {
|
|
let con = con.lock().map_err(|e| e.to_string())?;
|
|
crate::query::duckdb_execute(&con, &sql)
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?
|
|
}
|
|
_ => Err("Unsupported database type for transfer".to_string()),
|
|
}
|
|
}
|
|
|
|
fn database_from_pool_key(pool_key: &str) -> Option<&str> {
|
|
pool_key
|
|
.split_once(":session:")
|
|
.map(|(base, _)| base)
|
|
.unwrap_or(pool_key)
|
|
.split_once(':')
|
|
.map(|(_, database)| database)
|
|
.filter(|database| !database.is_empty())
|
|
}
|
|
|
|
pub async fn get_db_type(state: &AppState, connection_id: &str) -> Result<DatabaseType, String> {
|
|
let configs = state.configs.read().await;
|
|
configs
|
|
.get(connection_id)
|
|
.map(|c| c.db_type.clone())
|
|
.ok_or_else(|| format!("Connection config not found: {connection_id}"))
|
|
}
|
|
|
|
pub async fn get_columns_for_transfer(
|
|
state: &AppState,
|
|
pool_key: &str,
|
|
_connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<Vec<db::ColumnInfo>, String> {
|
|
let connections = state.connections.read().await;
|
|
|
|
if let Some(PoolKind::DuckDb(con)) = connections.get(pool_key) {
|
|
let con = con.clone();
|
|
drop(connections);
|
|
let table = table.to_string();
|
|
let schema = schema.to_string();
|
|
return tokio::task::spawn_blocking(move || {
|
|
let con = con.lock().map_err(|e| e.to_string())?;
|
|
crate::schema::duckdb_query_columns_in_database(&con, "main", &schema, &table)
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?;
|
|
}
|
|
|
|
if let Some(PoolKind::ExternalTabular(ext_pool)) = connections.get(pool_key) {
|
|
let con = ext_pool.cache.clone();
|
|
drop(connections);
|
|
let table = table.to_string();
|
|
let schema = schema.to_string();
|
|
return tokio::task::spawn_blocking(move || {
|
|
let con = con.lock().map_err(|e| e.to_string())?;
|
|
crate::schema::duckdb_query_columns_in_database(&con, "main", &schema, &table)
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?;
|
|
}
|
|
|
|
if let Some(PoolKind::ClickHouse(client)) = connections.get(pool_key) {
|
|
let client = client.clone();
|
|
let database = database.to_string();
|
|
let table = table.to_string();
|
|
drop(connections);
|
|
return db::clickhouse_driver::get_columns(&client, &database, &table).await;
|
|
}
|
|
if let Some(PoolKind::SqlServer(client)) = connections.get(pool_key) {
|
|
let client = client.clone();
|
|
let schema = schema.to_string();
|
|
let table = table.to_string();
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::get_columns(&mut client, &schema, &table).await;
|
|
}
|
|
if let Some(PoolKind::Agent(client)) = connections.get(pool_key) {
|
|
let client = client.clone();
|
|
let database = database.to_string();
|
|
let schema = schema.to_string();
|
|
let table = table.to_string();
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.get_columns(&database, &schema, &table).await;
|
|
}
|
|
let pool = connections.get(pool_key).ok_or("Pool not found")?;
|
|
let schema = schema.to_string();
|
|
let table = table.to_string();
|
|
match pool {
|
|
PoolKind::Mysql(p, _) => {
|
|
let p = p.clone();
|
|
drop(connections);
|
|
db::mysql::get_columns(&p, &schema, &table).await
|
|
}
|
|
PoolKind::Postgres(p) => {
|
|
let p = p.clone();
|
|
drop(connections);
|
|
db::postgres::get_columns(&p, &schema, &table).await
|
|
}
|
|
PoolKind::Sqlite(p) => {
|
|
let p = p.clone();
|
|
drop(connections);
|
|
db::sqlite::get_columns(&p, &schema, &table).await
|
|
}
|
|
_ => Err("Unsupported database type".to_string()),
|
|
}
|
|
}
|
|
|
|
pub async fn is_cancelled(transfer_id: &str) -> bool {
|
|
CANCELLED.read().await.contains(transfer_id)
|
|
}
|
|
|
|
pub async fn set_cancelled(transfer_id: &str) {
|
|
CANCELLED.write().await.insert(transfer_id.to_string());
|
|
}
|
|
|
|
pub async fn clear_cancelled(transfer_id: &str) {
|
|
CANCELLED.write().await.remove(transfer_id);
|
|
}
|
|
|
|
/// Transfer a single table. Returns rows transferred.
|
|
/// `progress_callback` is invoked for progress updates.
|
|
pub async fn transfer_table<F>(
|
|
state: &AppState,
|
|
request: &TransferRequest,
|
|
table: &str,
|
|
table_index: usize,
|
|
source_db_type: &DatabaseType,
|
|
target_db_type: &DatabaseType,
|
|
source_pool_key: &str,
|
|
target_pool_key: &str,
|
|
mut progress_callback: F,
|
|
) -> Result<u64, String>
|
|
where
|
|
F: FnMut(TransferProgress),
|
|
{
|
|
let total_tables = request.tables.len();
|
|
|
|
// Get source columns (deduplicate by name)
|
|
let columns = {
|
|
let raw = get_columns_for_transfer(
|
|
state,
|
|
source_pool_key,
|
|
&request.source_connection_id,
|
|
&request.source_database,
|
|
&request.source_schema,
|
|
table,
|
|
)
|
|
.await?;
|
|
let mut seen = std::collections::HashSet::new();
|
|
raw.into_iter().filter(|c| seen.insert(c.name.clone())).collect::<Vec<_>>()
|
|
};
|
|
|
|
if columns.is_empty() {
|
|
return Err(format!("No columns found for table {table}"));
|
|
}
|
|
|
|
let col_names: Vec<String> = columns.iter().map(|c| c.name.clone()).collect();
|
|
let col_types: Vec<Option<String>> = columns.iter().map(|c| Some(c.data_type.clone())).collect();
|
|
log::info!("[transfer] {} has {} columns, counting rows...", table, columns.len());
|
|
|
|
// Fetch source table comment
|
|
let table_comment: Option<String> = crate::schema::list_tables_core(
|
|
state,
|
|
&request.source_connection_id,
|
|
&request.source_database,
|
|
&request.source_schema,
|
|
Some(table),
|
|
Some(1),
|
|
)
|
|
.await
|
|
.unwrap_or_default()
|
|
.into_iter()
|
|
.next()
|
|
.and_then(|t| t.comment);
|
|
|
|
// Count source rows
|
|
let total_rows = {
|
|
let sql = count_sql(table, &request.source_schema, source_db_type);
|
|
match execute_on_pool(state, source_pool_key, &sql).await {
|
|
Ok(result) => result.rows.first().and_then(|r| r.first()).and_then(|v| match v {
|
|
serde_json::Value::Number(n) => n.as_u64(),
|
|
serde_json::Value::String(s) => s.parse::<u64>().ok(),
|
|
_ => None,
|
|
}),
|
|
Err(e) => {
|
|
log::warn!("[transfer] count failed for {}: {}", table, e);
|
|
None
|
|
}
|
|
}
|
|
};
|
|
log::info!("[transfer] {} total_rows={:?}", table, total_rows);
|
|
|
|
// Create table on target if requested
|
|
if request.create_table {
|
|
let ddl = generate_create_table_ddl(
|
|
&columns,
|
|
table,
|
|
&request.target_schema,
|
|
target_db_type,
|
|
source_db_type,
|
|
table_comment.as_deref(),
|
|
);
|
|
log::info!("[transfer] creating target table: {}", &ddl[..ddl.len().min(200)]);
|
|
let table_exists = match execute_on_pool(state, target_pool_key, &ddl).await {
|
|
Ok(_) => true,
|
|
Err(e) => {
|
|
let err_lower = e.to_lowercase();
|
|
if err_lower.contains("already exists") || err_lower.contains("there is already") {
|
|
true
|
|
} else {
|
|
return Err(format!("Failed to create table: {e}"));
|
|
}
|
|
}
|
|
};
|
|
if table_exists {
|
|
let comment_stmts =
|
|
generate_comment_ddl(&columns, table, &request.target_schema, target_db_type, table_comment.as_deref());
|
|
for stmt in &comment_stmts {
|
|
if let Err(e) = execute_on_pool(state, target_pool_key, stmt).await {
|
|
log::warn!("[transfer] failed to set column comment for {}: {}", table, e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Truncate target if overwrite mode
|
|
if request.mode == TransferMode::Overwrite {
|
|
let full_table = qualified_table(table, &request.target_schema, target_db_type);
|
|
let truncate_sql = match target_db_type {
|
|
DatabaseType::Sqlite | DatabaseType::DuckDb => format!("DELETE FROM {full_table}"),
|
|
_ => format!("TRUNCATE TABLE {full_table}"),
|
|
};
|
|
execute_on_pool(state, target_pool_key, &truncate_sql).await.map_err(|e| format!("Failed to truncate: {e}"))?;
|
|
}
|
|
|
|
// Determine effective mode and PK columns for upsert
|
|
let (effective_mode, pk_columns) = if request.mode == TransferMode::Upsert {
|
|
if matches!(target_db_type, DatabaseType::ClickHouse) {
|
|
log::warn!("[transfer] upsert not supported for ClickHouse, falling back to append");
|
|
(TransferMode::Append, vec![])
|
|
} else {
|
|
let target_columns = get_columns_for_transfer(
|
|
state,
|
|
target_pool_key,
|
|
&request.target_connection_id,
|
|
&request.target_database,
|
|
&request.target_schema,
|
|
table,
|
|
)
|
|
.await
|
|
.unwrap_or_default();
|
|
let pks: Vec<String> = target_columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.clone()).collect();
|
|
if pks.is_empty() {
|
|
log::warn!("[transfer] table {} has no primary key, falling back to append", table);
|
|
(TransferMode::Append, vec![])
|
|
} else {
|
|
(TransferMode::Upsert, pks)
|
|
}
|
|
}
|
|
} else {
|
|
(request.mode.clone(), vec![])
|
|
};
|
|
|
|
// Transfer data in batches
|
|
let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size };
|
|
let mut offset: u64 = 0;
|
|
let mut total_transferred: u64 = 0;
|
|
|
|
loop {
|
|
if is_cancelled(&request.transfer_id).await {
|
|
return Err("Cancelled".to_string());
|
|
}
|
|
|
|
let sql = pagination_sql(&col_names, table, &request.source_schema, source_db_type, offset, batch_size);
|
|
let result = execute_on_pool(state, source_pool_key, &sql).await?;
|
|
let row_count = result.rows.len();
|
|
|
|
if row_count == 0 {
|
|
break;
|
|
}
|
|
|
|
let batch_sql = match effective_mode {
|
|
TransferMode::Upsert => generate_upsert_typed(
|
|
&col_names,
|
|
&col_types,
|
|
&result.rows,
|
|
table,
|
|
&request.target_schema,
|
|
target_db_type,
|
|
&pk_columns,
|
|
),
|
|
_ => generate_insert_typed(
|
|
&col_names,
|
|
&col_types,
|
|
&result.rows,
|
|
table,
|
|
&request.target_schema,
|
|
target_db_type,
|
|
),
|
|
};
|
|
if !batch_sql.is_empty() {
|
|
execute_on_pool(state, target_pool_key, &batch_sql)
|
|
.await
|
|
.map_err(|e| format!("Insert failed at offset {offset}: {e}"))?;
|
|
}
|
|
|
|
total_transferred += row_count as u64;
|
|
log::info!("[transfer] {} batch +{} rows (total {})", table, row_count, total_transferred);
|
|
offset += row_count as u64;
|
|
|
|
progress_callback(TransferProgress {
|
|
transfer_id: request.transfer_id.clone(),
|
|
table: table.to_string(),
|
|
table_index,
|
|
total_tables,
|
|
rows_transferred: total_transferred,
|
|
total_rows,
|
|
status: TransferStatus::Running,
|
|
error: None,
|
|
});
|
|
|
|
if row_count < batch_size {
|
|
break;
|
|
}
|
|
}
|
|
|
|
Ok(total_transferred)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::connection::{AppState, PoolKind};
|
|
use crate::storage::Storage;
|
|
use serde_json::json;
|
|
use std::sync::Arc;
|
|
|
|
fn duckdb_test_config(id: &str) -> crate::models::connection::ConnectionConfig {
|
|
crate::models::connection::ConnectionConfig {
|
|
id: id.to_string(),
|
|
name: id.to_string(),
|
|
db_type: DatabaseType::DuckDb,
|
|
driver_profile: None,
|
|
driver_label: None,
|
|
url_params: None,
|
|
host: ":memory:".to_string(),
|
|
port: 0,
|
|
username: String::new(),
|
|
password: String::new(),
|
|
database: None,
|
|
visible_databases: None,
|
|
attached_databases: Vec::new(),
|
|
color: None,
|
|
ssh_enabled: false,
|
|
ssh_host: String::new(),
|
|
ssh_port: 22,
|
|
ssh_user: String::new(),
|
|
ssh_password: String::new(),
|
|
ssh_key_path: String::new(),
|
|
ssh_key_passphrase: String::new(),
|
|
ssh_expose_lan: false,
|
|
ssh_connect_timeout_secs: 5,
|
|
proxy_enabled: false,
|
|
proxy_type: crate::models::connection::ProxyType::Socks5,
|
|
proxy_host: String::new(),
|
|
proxy_port: 1080,
|
|
proxy_username: String::new(),
|
|
proxy_password: String::new(),
|
|
ssl: false,
|
|
ca_cert_path: String::new(),
|
|
sysdba: false,
|
|
oracle_connection_type: None,
|
|
connection_string: None,
|
|
external_config: None,
|
|
jdbc_driver_class: None,
|
|
jdbc_driver_paths: Vec::new(),
|
|
one_time: false,
|
|
}
|
|
}
|
|
|
|
fn test_column(name: &str, data_type: &str) -> db::ColumnInfo {
|
|
db::ColumnInfo {
|
|
name: name.to_string(),
|
|
data_type: data_type.to_string(),
|
|
is_nullable: true,
|
|
column_default: None,
|
|
is_primary_key: false,
|
|
extra: None,
|
|
comment: None,
|
|
numeric_precision: None,
|
|
numeric_scale: None,
|
|
character_maximum_length: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn mysql_create_table_includes_column_comments() {
|
|
let cols = vec![
|
|
db::ColumnInfo { comment: Some("用户ID".to_string()), is_primary_key: true, ..test_column("id", "int") },
|
|
db::ColumnInfo {
|
|
comment: Some("用户姓名".to_string()),
|
|
is_nullable: false,
|
|
..test_column("name", "varchar(100)")
|
|
},
|
|
db::ColumnInfo { comment: None, ..test_column("age", "int") },
|
|
];
|
|
|
|
let ddl = generate_create_table_ddl(&cols, "users", "", &DatabaseType::Mysql, &DatabaseType::Mysql, None);
|
|
|
|
assert!(ddl.contains("COMMENT '用户ID'"));
|
|
assert!(ddl.contains("COMMENT '用户姓名'"));
|
|
assert!(!ddl.contains("`age` INT COMMENT")); // no comment for age
|
|
assert!(ddl.contains("`name` VARCHAR(100) NOT NULL COMMENT '用户姓名'"));
|
|
assert!(ddl.contains("PRIMARY KEY (`id`)"));
|
|
}
|
|
|
|
#[test]
|
|
fn mysql_create_table_includes_table_comment() {
|
|
let cols = vec![db::ColumnInfo { is_primary_key: true, ..test_column("id", "int") }];
|
|
|
|
let ddl =
|
|
generate_create_table_ddl(&cols, "users", "", &DatabaseType::Mysql, &DatabaseType::Mysql, Some("用户表"));
|
|
|
|
assert!(ddl.contains(") COMMENT='用户表'"));
|
|
}
|
|
|
|
#[test]
|
|
fn mysql_text_pk_gets_key_prefix() {
|
|
let cols =
|
|
vec![db::ColumnInfo { data_type: "text".to_string(), is_primary_key: true, ..test_column("id", "text") }];
|
|
|
|
let ddl = generate_create_table_ddl(&cols, "logs", "", &DatabaseType::Mysql, &DatabaseType::Sqlite, None);
|
|
|
|
assert!(ddl.contains("PRIMARY KEY (`id`(255))"));
|
|
assert!(ddl.contains("`id` TEXT"));
|
|
}
|
|
|
|
#[test]
|
|
fn mysql_int_pk_no_prefix() {
|
|
let cols = vec![db::ColumnInfo { is_primary_key: true, ..test_column("id", "int") }];
|
|
|
|
let ddl = generate_create_table_ddl(&cols, "users", "", &DatabaseType::Mysql, &DatabaseType::Sqlite, None);
|
|
|
|
assert!(ddl.contains("PRIMARY KEY (`id`)"));
|
|
assert!(!ddl.contains("PRIMARY KEY (`id`(255))"));
|
|
}
|
|
|
|
#[test]
|
|
fn postgres_comment_ddl_generates_column_and_table_comments() {
|
|
let cols = vec![
|
|
db::ColumnInfo { comment: Some("主键".to_string()), ..test_column("id", "int") },
|
|
db::ColumnInfo { comment: Some("名称".to_string()), ..test_column("name", "varchar(100)") },
|
|
];
|
|
|
|
let stmts = generate_comment_ddl(&cols, "items", "public", &DatabaseType::Postgres, Some("项目表"));
|
|
|
|
assert_eq!(stmts.len(), 3);
|
|
assert!(stmts[0].contains("COMMENT ON TABLE \"public\".\"items\" IS '项目表'"));
|
|
assert!(stmts[1].contains("COMMENT ON COLUMN \"public\".\"items\".\"id\" IS '主键'"));
|
|
assert!(stmts[2].contains("COMMENT ON COLUMN \"public\".\"items\".\"name\" IS '名称'"));
|
|
}
|
|
|
|
#[test]
|
|
fn clickhouse_comment_ddl_uses_alter_table() {
|
|
let cols = vec![db::ColumnInfo { comment: Some("日志消息".to_string()), ..test_column("message", "text") }];
|
|
|
|
let stmts = generate_comment_ddl(&cols, "logs", "", &DatabaseType::ClickHouse, None);
|
|
|
|
assert_eq!(stmts.len(), 1);
|
|
assert!(stmts[0].contains("ALTER TABLE `logs` COMMENT COLUMN `message` '日志消息'"));
|
|
}
|
|
|
|
#[test]
|
|
fn pg_comment_ddl_skips_empty_comments() {
|
|
let cols = vec![
|
|
db::ColumnInfo { comment: None, ..test_column("id", "int") },
|
|
db::ColumnInfo { comment: Some(" ".to_string()), ..test_column("name", "varchar(100)") },
|
|
];
|
|
|
|
let stmts = generate_comment_ddl(&cols, "t", "", &DatabaseType::Postgres, None);
|
|
|
|
assert!(stmts.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn non_mysql_family_no_inline_comment() {
|
|
let cols = vec![db::ColumnInfo { comment: Some("test".to_string()), ..test_column("col", "text") }];
|
|
|
|
// PostgreSQL target should NOT have inline COMMENT
|
|
let ddl = generate_create_table_ddl(&cols, "t", "", &DatabaseType::Postgres, &DatabaseType::Postgres, None);
|
|
assert!(!ddl.contains("COMMENT"));
|
|
}
|
|
|
|
#[test]
|
|
fn mysql_insert_normalizes_rfc3339_datetime_strings() {
|
|
let sql = generate_insert_typed(
|
|
&[String::from("insurance_start_time")],
|
|
&[Some(String::from("datetime"))],
|
|
&[vec![json!("2026-05-12T00:00:00+00:00")]],
|
|
"policies",
|
|
"",
|
|
&DatabaseType::Mysql,
|
|
);
|
|
|
|
assert_eq!(sql, "INSERT INTO `policies` (`insurance_start_time`) VALUES\n('2026-05-12 00:00:00')");
|
|
}
|
|
|
|
#[test]
|
|
fn mysql_insert_uses_column_types_for_temporal_literals() {
|
|
let sql = generate_insert_typed(
|
|
&[String::from("dt"), String::from("raw_text"), String::from("d"), String::from("t")],
|
|
&[
|
|
Some(String::from("datetime")),
|
|
Some(String::from("varchar(64)")),
|
|
Some(String::from("date")),
|
|
Some(String::from("time")),
|
|
],
|
|
&[vec![
|
|
json!("2026-05-12T00:00:00+00:00"),
|
|
json!("2026-05-12T00:00:00+00:00"),
|
|
json!("2026-05-12T00:00:00+00:00"),
|
|
json!("2026-05-12T09:30:45+00:00"),
|
|
]],
|
|
"policies",
|
|
"",
|
|
&DatabaseType::Mysql,
|
|
);
|
|
|
|
assert_eq!(
|
|
sql,
|
|
"INSERT INTO `policies` (`dt`, `raw_text`, `d`, `t`) VALUES\n('2026-05-12 00:00:00', '2026-05-12T00:00:00+00:00', '2026-05-12', '09:30:45')"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn duckdb_transfer_columns_use_requested_schema() {
|
|
let dir = std::env::temp_dir().join(format!("dbx-transfer-test-{}", uuid::Uuid::new_v4()));
|
|
std::fs::create_dir_all(&dir).unwrap();
|
|
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
|
let con = duckdb::Connection::open_in_memory().unwrap();
|
|
con.execute_batch("CREATE SCHEMA analytics; CREATE TABLE analytics.items(id INTEGER);").unwrap();
|
|
|
|
let state = AppState::new(storage);
|
|
let con = Arc::new(std::sync::Mutex::new(con));
|
|
state.connections.write().await.insert("duckdb-1".to_string(), PoolKind::DuckDb(con));
|
|
state.configs.write().await.insert("duckdb-1".to_string(), duckdb_test_config("duckdb-1"));
|
|
|
|
let columns =
|
|
get_columns_for_transfer(&state, "duckdb-1", "duckdb-1", "main", "analytics", "items").await.unwrap();
|
|
|
|
assert_eq!(columns.iter().map(|c| c.name.as_str()).collect::<Vec<_>>(), vec!["id"]);
|
|
}
|
|
|
|
#[test]
|
|
fn database_from_pool_key_handles_session_scoped_keys() {
|
|
assert_eq!(database_from_pool_key("conn:analytics"), Some("analytics"));
|
|
assert_eq!(database_from_pool_key("conn:analytics:session:editor-1"), Some("analytics"));
|
|
assert_eq!(database_from_pool_key("conn"), None);
|
|
}
|
|
}
|