Files
dbx/crates/dbx-core/src/transfer.rs
T
e9ee982843 Preserve ordered SSH and proxy connection layers (#802)
* Preserve ordered SSH and proxy connection layers

Replace separate SSH/proxy connection fields with ordered transport_layers while keeping legacy migration and secret fallback paths intact. The UI now edits SSH tunnel/proxy layers in configured order without a global SSH enable gate.

Constraint: Existing saved SSH tunnels, proxy settings, and secret-store keys must continue to load through migration.
Rejected: Folding proxy fields into SshTunnelConfig | mixes proxy semantics into an SSH-specific structure.
Confidence: high
Scope-risk: moderate
Directive: Keep SSH/proxy structs provider-specific; put cross-layer chaining in transport_layer_tunnel orchestration.
Tested: git diff --check; cargo fmt --check; cargo check --workspace; cargo test -p dbx-core --lib; vue-tsc --noEmit --project apps/desktop/tsconfig.json; tsc -p packages/node-core/tsconfig.json --noEmit; oxlint --vue-plugin apps/desktop/src; tsx --tsconfig apps/desktop/tsconfig.json --test packages/app-tests/*.test.ts; tsx --test packages/node-core/tests/*.test.ts
Not-tested: Live external SSH/proxy/database endpoint integration.

* Keep ordered transport changes CI-format clean

Constraint: CI pnpm check failed only on oxfmt formatting for two desktop TypeScript files
Confidence: high
Scope-risk: narrow
Directive: Keep generated/editor config changes out of this PR fix commit
Tested: PATH="/Volumes/data/code/rust/dbx/node_modules/.bin:/Users/hekx/.codex/tmp/arg0/codex-arg0nuL34n:/Users/hekx/.cargo/bin:/Users/hekx/.local/bin:/opt/homebrew/opt/llvm/bin:/opt/homebrew/opt/libpq/bin:/Volumes/data/Users/hekx/.opencode/bin:/Users/hekx/.bun/bin:/Volumes/data/sdks/flutter/bin:/Volumes/data/Users/hekx/.cargo/bin:/Users/hekx/.local/bin:/opt/homebrew/opt/llvm/bin:/opt/homebrew/opt/libpq/bin:/Volumes/data/Users/hekx/.opencode/bin:/Users/hekx/.bun/bin:/Volumes/data/Users/hekx/Library/pnpm:/Volumes/data/sdks/flutter/bin:/Volumes/data/Users/hekx/.cargo/bin:/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/pkg/env/global/bin:/opt/X11/bin:/Library/Apple/usr/bin:/Applications/Wireshark.app/Contents/MacOS:/usr/local/go/bin:/opt/homebrew/bin:/opt/podman/bin:/Applications/Ghostty.app/Contents/MacOS" rtk node scripts/run-check.mjs
Not-tested: GitHub Actions rerun not observed locally

* Keep Rust checks warning-clean

Constraint: cargo clippy --workspace --all-targets --all-features -- -D warnings failed across core, web, and tauri crates
Rejected: Broad workspace-level clippy suppression | kept allows local to long-argument command/API boundaries
Confidence: high
Scope-risk: moderate
Directive: Preserve src-tauri/tauri.conf.json as an unrelated local change outside this commit
Tested: rtk cargo clippy --workspace --all-targets --all-features -- -D warnings
Not-tested: Full GitHub Actions rerun not observed locally

* Reuse existing proxy tunnels on retry

Proxy tunnel startup now mirrors SSH tunnel behavior by returning the existing local port for an active connection id instead of replacing the managed handle. A second map check aborts a just-spawned duplicate handle if a concurrent retry won the race before insertion, preventing orphaned listeners while keeping the change narrow.

Constraint: Reviewer requested proxy tunnel behavior align with SSH local-port reuse

Rejected: Always overwrite and abort the previous handle | less consistent with SSH behavior and churns listeners during retries

Confidence: high

Scope-risk: narrow

Tested: cargo fmt --check --all

Tested: cargo test -p dbx-core db::proxy_tunnel::tests::start_tunnel_reuses_existing_local_port

Tested: cargo clippy -p dbx-core --all-targets -- -D warnings

---------

Co-authored-by: hekx <hekx@momandeMac-mini.local>
Co-authored-by: caisin <caisin@caisins-Mac-mini.local>
2026-06-06 23:22:29 +08:00

3740 lines
139 KiB
Rust

use regex::Regex;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use tokio::sync::RwLock;
use crate::connection::{AppState, PoolKind};
use crate::db;
use crate::db::mongo_driver::MongoDocumentResult;
use crate::models::connection::DatabaseType;
use crate::object_source_sql::{build_executable_object_source_statements, EditableObjectSourceSqlInput};
use crate::query::{agent_execute_query_params, QueryExecutionOptions};
use crate::sql::starts_with_executable_sql_keyword;
static CANCELLED: std::sync::LazyLock<RwLock<HashSet<String>>> =
std::sync::LazyLock::new(|| RwLock::new(HashSet::new()));
const MAX_TRANSFER_WRITE_SQL_BYTES: usize = 512 * 1024;
const MAX_SQLSERVER_INSERT_ROWS: usize = 1000;
const MAX_ORACLE_MERGE_ROWS: usize = 500;
#[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
| DatabaseType::Hive => 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() || matches!(db_type, DatabaseType::Mysql | DatabaseType::MongoDb) {
qt
} else {
format!("{}.{}", quote_identifier(schema, db_type), qt)
}
}
fn quote_string_literal(value: &str) -> String {
format!("'{}'", value.replace('\'', "''"))
}
fn is_simple_identifier(value: &str) -> bool {
let mut chars = value.chars();
let Some(first) = chars.next() else {
return false;
};
if !(first == '_' || first.is_ascii_alphabetic()) {
return false;
}
chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
}
fn is_postgres_compat_transfer(source_db: &DatabaseType, target_db: &DatabaseType) -> bool {
matches!(source_db, DatabaseType::Postgres) && matches!(target_db, DatabaseType::Postgres)
}
fn is_postgres_integer_like_type(data_type: &str) -> bool {
let normalized = data_type.trim().to_ascii_lowercase();
matches!(
normalized.split(['(', ' ']).next().unwrap_or(""),
"smallint" | "integer" | "bigint" | "int2" | "int4" | "int8"
)
}
fn is_postgres_sequence_default(default_value: Option<&str>) -> bool {
default_value.is_some_and(|value| value.to_ascii_lowercase().contains("nextval("))
}
fn rewrite_postgres_schema_qualified_references(input: &str, source_schema: &str, target_schema: &str) -> String {
if source_schema.trim().is_empty() || source_schema == target_schema {
return input.to_string();
}
let quoted_source = format!("{}.", quote_identifier(source_schema, &DatabaseType::Postgres));
let quoted_target = format!("{}.", quote_identifier(target_schema, &DatabaseType::Postgres));
let rewritten = input.replace(&quoted_source, &quoted_target);
let unquoted_pattern =
Regex::new(&format!(r#"(^|[^"\w]){}\."#, regex::escape(source_schema))).expect("valid postgres schema regex");
unquoted_pattern
.replace_all(&rewritten, |captures: &regex::Captures| format!("{}{}", &captures[1], quoted_target))
.into_owned()
}
fn postgres_column_type_sql(
column: &db::ColumnInfo,
source_schema: &str,
target_schema: &str,
source_db: &DatabaseType,
target_db: &DatabaseType,
) -> String {
if is_postgres_compat_transfer(source_db, target_db) {
let trimmed = column.data_type.trim();
if !trimmed.is_empty() {
return rewrite_postgres_schema_qualified_references(trimmed, source_schema, target_schema);
}
}
map_column_type(&column.data_type, source_db, target_db)
}
fn postgres_default_clause(
column: &db::ColumnInfo,
source_schema: &str,
target_schema: &str,
source_db: &DatabaseType,
target_db: &DatabaseType,
) -> Option<String> {
if !is_postgres_compat_transfer(source_db, target_db) {
return None;
}
let default_value = column.column_default.as_deref()?.trim();
if default_value.is_empty() {
return None;
}
if is_postgres_sequence_default(Some(default_value)) && is_postgres_integer_like_type(&column.data_type) {
return Some("GENERATED BY DEFAULT AS IDENTITY".to_string());
}
Some(format!(
"DEFAULT {}",
rewrite_postgres_schema_qualified_references(default_value, source_schema, target_schema)
))
}
fn postgres_order_by_expression(columns: &[String], db_type: &DatabaseType) -> Option<String> {
if columns.is_empty() {
return None;
}
Some(columns.iter().map(|column| quote_identifier(column, db_type)).collect::<Vec<_>>().join(", "))
}
fn postgres_index_column_sql(column: &str) -> String {
if is_simple_identifier(column) {
quote_identifier(column, &DatabaseType::Postgres)
} else {
column.to_string()
}
}
fn generate_postgres_index_ddl(indexes: &[db::IndexInfo], table: &str, schema: &str) -> Vec<String> {
let full_table = qualified_table(table, schema, &DatabaseType::Postgres);
let mut statements = Vec::new();
for index in indexes.iter().filter(|index| !index.is_primary) {
if index.name.trim().is_empty() || index.columns.is_empty() {
continue;
}
let unique = if index.is_unique { "UNIQUE " } else { "" };
let using_clause = index
.index_type
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!(" USING {value}"))
.unwrap_or_default();
let columns =
index.columns.iter().map(|column| postgres_index_column_sql(column)).collect::<Vec<_>>().join(", ");
let include_clause = index
.included_columns
.as_ref()
.filter(|columns| !columns.is_empty())
.map(|columns| {
format!(
" INCLUDE ({})",
columns
.iter()
.map(|column| quote_identifier(column, &DatabaseType::Postgres))
.collect::<Vec<_>>()
.join(", ")
)
})
.unwrap_or_default();
let filter_clause = index
.filter
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!(" WHERE {value}"))
.unwrap_or_default();
statements.push(format!(
"CREATE {unique}INDEX IF NOT EXISTS {} ON {full_table}{using_clause} ({columns}){include_clause}{filter_clause}",
quote_identifier(&index.name, &DatabaseType::Postgres)
));
if let Some(comment) = index.comment.as_deref().map(str::trim).filter(|value| !value.is_empty()) {
let qualified_index = if schema.is_empty() {
quote_identifier(&index.name, &DatabaseType::Postgres)
} else {
format!(
"{}.{}",
quote_identifier(schema, &DatabaseType::Postgres),
quote_identifier(&index.name, &DatabaseType::Postgres)
)
};
statements.push(format!("COMMENT ON INDEX {qualified_index} IS {}", quote_string_literal(comment)));
}
}
statements
}
fn generate_postgres_foreign_key_ddl(foreign_keys: &[db::ForeignKeyInfo], table: &str, schema: &str) -> Vec<String> {
let full_table = qualified_table(table, schema, &DatabaseType::Postgres);
let mut grouped: HashMap<&str, Vec<&db::ForeignKeyInfo>> = HashMap::new();
let mut order = Vec::new();
for foreign_key in foreign_keys {
if !grouped.contains_key(foreign_key.name.as_str()) {
order.push(foreign_key.name.as_str());
}
grouped.entry(foreign_key.name.as_str()).or_default().push(foreign_key);
}
let mut statements = Vec::new();
for name in order {
let Some(group) = grouped.get(name) else {
continue;
};
let columns = group
.iter()
.map(|foreign_key| quote_identifier(&foreign_key.column, &DatabaseType::Postgres))
.collect::<Vec<_>>()
.join(", ");
let ref_columns = group
.iter()
.map(|foreign_key| quote_identifier(&foreign_key.ref_column, &DatabaseType::Postgres))
.collect::<Vec<_>>()
.join(", ");
let referenced_schema = group[0].ref_schema.as_deref().unwrap_or(schema);
let referenced_table = qualified_table(&group[0].ref_table, referenced_schema, &DatabaseType::Postgres);
statements.push(format!(
"ALTER TABLE {full_table} ADD CONSTRAINT {} FOREIGN KEY ({columns}) REFERENCES {referenced_table} ({ref_columns})",
quote_identifier(name, &DatabaseType::Postgres)
));
}
statements
}
fn generate_postgres_sequence_sync_sql(columns: &[db::ColumnInfo], table: &str, schema: &str) -> Vec<String> {
let full_table = qualified_table(table, schema, &DatabaseType::Postgres);
columns
.iter()
.filter(|column| is_postgres_sequence_default(column.column_default.as_deref()))
.map(|column| {
let quoted_column = quote_identifier(&column.name, &DatabaseType::Postgres);
format!(
"SELECT setval(pg_get_serial_sequence({}, {}), GREATEST(COALESCE(MAX({quoted_column}), 0), 1), MAX({quoted_column}) IS NOT NULL) FROM {full_table}",
quote_string_literal(&full_table),
quote_string_literal(&column.name)
)
})
.collect()
}
#[derive(Debug, Clone)]
struct PostgresTriggerSource {
table_name: String,
trigger_name: String,
source: String,
}
#[derive(Debug, Clone)]
struct PostgresExtensionSource {
extension_name: String,
}
#[derive(Debug, Clone)]
struct PostgresEnumSource {
type_name: String,
labels: Vec<String>,
}
#[derive(Debug, Clone)]
struct PostgresDomainSource {
domain_name: String,
base_type: String,
default_value: Option<String>,
not_null: bool,
checks: Vec<String>,
}
#[derive(Debug, Clone)]
struct PostgresMaterializedViewSource {
view_name: String,
source: String,
}
fn json_string_cell(row: &[serde_json::Value], index: usize) -> Option<String> {
row.get(index).and_then(|value| value.as_str().map(str::to_string))
}
fn result_rows_to_string_statements(rows: Vec<Vec<serde_json::Value>>) -> Vec<String> {
rows.into_iter().filter_map(|row| json_string_cell(&row, 0)).filter(|stmt| !stmt.trim().is_empty()).collect()
}
fn ensure_sql_statement_terminated(sql: &str) -> String {
let trimmed = sql.trim();
if trimmed.ends_with(';') {
trimmed.to_string()
} else {
format!("{trimmed};")
}
}
fn generate_postgres_extension_ddl(extension: &PostgresExtensionSource, target_schema: &str) -> String {
format!(
"CREATE EXTENSION IF NOT EXISTS {} WITH SCHEMA {}",
quote_identifier(&extension.extension_name, &DatabaseType::Postgres),
quote_identifier(target_schema, &DatabaseType::Postgres)
)
}
fn generate_postgres_enum_ddl(enum_type: &PostgresEnumSource, target_schema: &str) -> String {
let labels = enum_type.labels.iter().map(|label| quote_string_literal(label)).collect::<Vec<_>>().join(", ");
let create_sql = format!(
"CREATE TYPE {}.{} AS ENUM ({labels})",
quote_identifier(target_schema, &DatabaseType::Postgres),
quote_identifier(&enum_type.type_name, &DatabaseType::Postgres)
);
format!(
"DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace WHERE n.nspname = {} AND t.typname = {}) THEN {create_sql}; END IF; END $$",
quote_string_literal(target_schema),
quote_string_literal(&enum_type.type_name)
)
}
fn generate_postgres_domain_ddl(domain: &PostgresDomainSource, target_schema: &str) -> String {
let mut create_sql = format!(
"CREATE DOMAIN {}.{} AS {}",
quote_identifier(target_schema, &DatabaseType::Postgres),
quote_identifier(&domain.domain_name, &DatabaseType::Postgres),
domain.base_type
);
if let Some(default_value) = domain.default_value.as_deref().map(str::trim).filter(|value| !value.is_empty()) {
create_sql.push_str(&format!(" DEFAULT {default_value}"));
}
if domain.not_null {
create_sql.push_str(" NOT NULL");
}
for check in &domain.checks {
create_sql.push(' ');
create_sql.push_str(check);
}
format!(
"DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace WHERE n.nspname = {} AND t.typname = {}) THEN {}; END IF; END $$",
quote_string_literal(target_schema),
quote_string_literal(&domain.domain_name),
create_sql
)
}
fn generate_postgres_materialized_view_ddls(view: &PostgresMaterializedViewSource, target_schema: &str) -> Vec<String> {
let qualified_name = qualified_table(&view.view_name, target_schema, &DatabaseType::Postgres);
vec![
format!("DROP MATERIALIZED VIEW IF EXISTS {qualified_name}"),
format!("CREATE MATERIALIZED VIEW {qualified_name} AS\n{}", ensure_sql_statement_terminated(&view.source)),
]
}
fn rewrite_postgres_routine_schema(source: &str, target_schema: &str) -> Option<String> {
let re = Regex::new(
r#"(?is)^(\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:(?:NON)?EDITIONABLE\s+)?(?:FUNCTION|PROCEDURE)\s+)((?:"(?:""|[^"])+"|[A-Za-z_][\w$]*)(?:\s*\.\s*(?:"(?:""|[^"])+"|[A-Za-z_][\w$]*))?)"#,
)
.ok()?;
let captures = re.captures(source)?;
let full = captures.get(0)?;
let prefix = captures.get(1)?.as_str();
let existing_name = captures.get(2)?.as_str();
let name_re = Regex::new(r#""(?:""|[^"])+"|[A-Za-z_][\w$]*"#).ok()?;
let parts = name_re
.find_iter(existing_name)
.map(|part| part.as_str().trim().trim_matches('"').replace("\"\"", "\""))
.collect::<Vec<_>>();
let name = parts.last()?;
let replacement = format!(
"{}.{}",
quote_identifier(target_schema, &DatabaseType::Postgres),
quote_identifier(name, &DatabaseType::Postgres)
);
Some(format!("{}{}{}{}", &source[..full.start()], prefix, replacement, &source[full.end()..]))
}
fn rewrite_postgres_trigger_table_schema(
source: &str,
source_schema: &str,
table_name: &str,
target_schema: &str,
) -> String {
let qualified_target_table = qualified_table(table_name, target_schema, &DatabaseType::Postgres);
let candidate_patterns = [
format!(
" ON {}.{} ",
quote_identifier(source_schema, &DatabaseType::Postgres),
quote_identifier(table_name, &DatabaseType::Postgres)
),
format!(" ON {source_schema}.{table_name} "),
format!(" ON {} ", quote_identifier(table_name, &DatabaseType::Postgres)),
format!(" ON {table_name} "),
];
for pattern in candidate_patterns {
if source.contains(&pattern) {
return source.replacen(&pattern, &format!(" ON {qualified_target_table} "), 1);
}
}
source.to_string()
}
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();
if matches!(target_db, DatabaseType::Hive) {
return match base {
"tinyint" => "TINYINT".into(),
"smallint" | "int2" => "SMALLINT".into(),
"int" | "integer" | "int4" | "mediumint" | "serial" | "smallserial" => "INT".into(),
"bigint" | "int8" | "bigserial" => "BIGINT".into(),
"float" | "float4" | "real" => "FLOAT".into(),
"double" | "double precision" | "float8" => "DOUBLE".into(),
"decimal" | "numeric" | "number" => {
if let Some(index) = t.find('(') {
format!("DECIMAL{}", &t[index..])
} else {
"DECIMAL".into()
}
}
"bool" | "boolean" | "bit" => "BOOLEAN".into(),
"date" => "DATE".into(),
"datetime" | "timestamp" | "timestamptz" | "timestamp with time zone" | "timestamp without time zone" => {
"TIMESTAMP".into()
}
"binary" | "varbinary" | "blob" | "tinyblob" | "mediumblob" | "longblob" | "bytea" | "image" => {
"BINARY".into()
}
_ => "STRING".into(),
};
}
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()
}
}
"longtext" => match target_db {
DatabaseType::Mysql => "LONGTEXT".into(),
_ => "TEXT".into(),
},
"mediumtext" => match target_db {
DatabaseType::Mysql => "MEDIUMTEXT".into(),
_ => "TEXT".into(),
},
"text" | "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(),
},
"longblob" => match target_db {
DatabaseType::Mysql => "LONGBLOB".into(),
DatabaseType::Postgres => "BYTEA".into(),
DatabaseType::SqlServer => "VARBINARY(MAX)".into(),
_ => "BLOB".into(),
},
"mediumblob" => match target_db {
DatabaseType::Mysql => "MEDIUMBLOB".into(),
DatabaseType::Postgres => "BYTEA".into(),
DatabaseType::SqlServer => "VARBINARY(MAX)".into(),
_ => "BLOB".into(),
},
"blob" | "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"
)
}
fn parse_mysql_row_error(error: &str) -> Option<u64> {
let error = error.trim();
let at_row = error.rsplit("at row ").next()?;
at_row.trim().parse::<u64>().ok()
}
pub fn generate_create_table_ddl(
columns: &[db::ColumnInfo],
table: &str,
source_schema: &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 = postgres_column_type_sql(c, source_schema, schema, source_db, target_db);
let mut line = format!(" {} {}", quote_identifier(&c.name, target_db), mapped_type);
if let Some(default_clause) = postgres_default_clause(c, source_schema, schema, source_db, target_db) {
line.push(' ');
line.push_str(&default_clause);
}
if !c.is_nullable && !matches!(target_db, DatabaseType::Hive) {
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::with_capacity(columns.iter().filter(|c| c.is_primary_key).count());
if !matches!(target_db, DatabaseType::Hive) {
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(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(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),
}
}
fn max_transfer_write_rows(db_type: &DatabaseType, mode: &TransferMode) -> usize {
match (db_type, mode) {
(DatabaseType::SqlServer, TransferMode::Append | TransferMode::Overwrite) => MAX_SQLSERVER_INSERT_ROWS,
(DatabaseType::Hive, _) => 500,
(DatabaseType::Oracle, TransferMode::Upsert) => MAX_ORACLE_MERGE_ROWS,
_ => usize::MAX,
}
}
#[allow(clippy::too_many_arguments)]
fn generate_transfer_write_sql(
mode: &TransferMode,
columns: &[String],
column_types: &[Option<String>],
rows: &[Vec<serde_json::Value>],
table: &str,
schema: &str,
db_type: &DatabaseType,
pk_columns: &[String],
) -> String {
match mode {
TransferMode::Upsert => generate_upsert_typed(columns, column_types, rows, table, schema, db_type, pk_columns),
_ => generate_insert_typed(columns, column_types, rows, table, schema, db_type),
}
}
#[allow(clippy::too_many_arguments)]
fn generate_transfer_write_sql_batches(
mode: &TransferMode,
columns: &[String],
column_types: &[Option<String>],
rows: &[Vec<serde_json::Value>],
table: &str,
schema: &str,
db_type: &DatabaseType,
pk_columns: &[String],
) -> Vec<String> {
if rows.is_empty() {
return Vec::new();
}
let max_rows = max_transfer_write_rows(db_type, mode);
let mut statements = Vec::new();
let mut start = 0;
while start < rows.len() {
let mut end = start + 1;
let mut accepted = generate_transfer_write_sql(
mode,
columns,
column_types,
&rows[start..end],
table,
schema,
db_type,
pk_columns,
);
while end < rows.len() && end - start < max_rows {
let candidate = generate_transfer_write_sql(
mode,
columns,
column_types,
&rows[start..=end],
table,
schema,
db_type,
pk_columns,
);
if candidate.len() > MAX_TRANSFER_WRITE_SQL_BYTES && !accepted.is_empty() {
break;
}
accepted = candidate;
end += 1;
}
if !accepted.is_empty() {
statements.push(accepted);
}
start = end;
}
statements
}
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 pagination_sql_with_order(
columns: &[String],
table: &str,
schema: &str,
db_type: &DatabaseType,
offset: u64,
limit: usize,
order_by_columns: &[String],
) -> 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(", ");
let order_expression = postgres_order_by_expression(order_by_columns, db_type);
match db_type {
DatabaseType::SqlServer | DatabaseType::Oracle => {
let order_by = order_expression.unwrap_or_else(|| "(SELECT NULL)".to_string());
format!(
"SELECT {col_list} FROM {full_table} ORDER BY {order_by} OFFSET {offset} ROWS FETCH NEXT {limit} ROWS ONLY"
)
}
_ => {
let order_by = order_expression.map(|value| format!(" ORDER BY {value}")).unwrap_or_default();
format!("SELECT {col_list} FROM {full_table}{order_by} LIMIT {limit} OFFSET {offset}")
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn pagination_sql_with_filter_order(
columns: &[String],
table: &str,
schema: &str,
db_type: &DatabaseType,
offset: u64,
limit: usize,
where_input: Option<&str>,
order_by: Option<&str>,
default_order_columns: &[String],
) -> 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(", ");
let predicate = crate::sql_dialect::normalize_where_input(where_input);
let where_clause = if predicate.is_empty() { String::new() } else { format!(" WHERE ({predicate})") };
let order_expression = order_by
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| postgres_order_by_expression(default_order_columns, db_type));
match db_type {
DatabaseType::SqlServer | DatabaseType::Oracle => {
let order_by = order_expression.unwrap_or_else(|| "(SELECT NULL)".to_string());
format!(
"SELECT {col_list} FROM {full_table}{where_clause} ORDER BY {order_by} OFFSET {offset} ROWS FETCH NEXT {limit} ROWS ONLY"
)
}
_ => {
let order_by = order_expression.map(|value| format!(" ORDER BY {value}")).unwrap_or_default();
format!("SELECT {col_list} FROM {full_table}{where_clause}{order_by} LIMIT {limit} OFFSET {offset}")
}
}
}
pub fn count_sql(table: &str, schema: &str, db_type: &DatabaseType) -> String {
count_sql_with_where(table, schema, db_type, None)
}
pub fn count_sql_with_where(table: &str, schema: &str, db_type: &DatabaseType, where_input: Option<&str>) -> String {
let full_table = qualified_table(table, schema, db_type);
let predicate = crate::sql_dialect::normalize_where_input(where_input);
let where_clause = if predicate.is_empty() { String::new() } else { format!(" WHERE ({predicate})") };
format!("SELECT COUNT(*) FROM {full_table}{where_clause}")
}
pub fn keyset_pagination_sql(
columns: &[String],
table: &str,
schema: &str,
db_type: &DatabaseType,
primary_keys: &[String],
last_pk_values: &[serde_json::Value],
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(", ");
let order =
primary_keys.iter().map(|pk| format!("{} ASC", quote_identifier(pk, db_type))).collect::<Vec<_>>().join(", ");
let where_clause = keyset_where_clause(primary_keys, last_pk_values, db_type);
match db_type {
DatabaseType::SqlServer | DatabaseType::Oracle => {
format!(
"SELECT {col_list} FROM {full_table}{where_clause} ORDER BY {order} OFFSET 0 ROWS FETCH NEXT {limit} ROWS ONLY"
)
}
_ => {
format!("SELECT {col_list} FROM {full_table}{where_clause} ORDER BY {order} LIMIT {limit}")
}
}
}
fn keyset_where_clause(
primary_keys: &[String],
last_pk_values: &[serde_json::Value],
db_type: &DatabaseType,
) -> String {
if primary_keys.is_empty() || last_pk_values.is_empty() {
return String::new();
}
let quoted_keys = primary_keys.iter().map(|pk| quote_identifier(pk, db_type)).collect::<Vec<_>>();
let literals = last_pk_values.iter().map(|v| value_to_sql_literal(v, db_type)).collect::<Vec<_>>();
let comparison_count = quoted_keys.len().min(literals.len());
if comparison_count == 0 {
return String::new();
}
let mut clauses = Vec::with_capacity(comparison_count);
for index in 0..comparison_count {
let mut parts = Vec::with_capacity(index + 1);
for prefix_index in 0..index {
parts.push(format!("{} = {}", quoted_keys[prefix_index], literals[prefix_index]));
}
parts.push(format!("{} > {}", quoted_keys[index], literals[index]));
if parts.len() == 1 {
clauses.push(parts.remove(0));
} else {
clauses.push(format!("({})", parts.join(" AND ")));
}
}
if clauses.len() == 1 {
format!(" WHERE {}", clauses[0])
} else {
format!(" WHERE ({})", clauses.join(" OR "))
}
}
fn value_to_sql_literal(value: &serde_json::Value, _db_type: &DatabaseType) -> String {
match value {
serde_json::Value::Null => "NULL".to_string(),
serde_json::Value::Bool(b) => {
if *b {
"TRUE".to_string()
} else {
"FALSE".to_string()
}
}
serde_json::Value::Number(n) => n.to_string(),
serde_json::Value::String(s) => quote_string_literal(s),
_ => quote_string_literal(&value.to_string()),
}
}
fn is_mongodb_transfer_type(db_type: &DatabaseType) -> bool {
matches!(db_type, DatabaseType::MongoDb)
}
fn mongo_transfer_document_fields(documents: &[serde_json::Value]) -> Vec<String> {
let mut fields = Vec::new();
let mut seen = HashSet::new();
for document in documents {
let Some(object) = document.as_object() else {
continue;
};
for key in object.keys() {
if seen.insert(key.clone()) {
fields.push(key.clone());
}
}
}
fields
}
fn mongo_documents_to_rows(documents: &[serde_json::Value], columns: &[String]) -> Vec<Vec<serde_json::Value>> {
documents
.iter()
.map(|document| {
let object = document.as_object();
columns
.iter()
.map(|column| object.and_then(|values| values.get(column)).cloned().unwrap_or(serde_json::Value::Null))
.collect()
})
.collect()
}
fn sql_rows_to_mongo_documents(columns: &[String], rows: &[Vec<serde_json::Value>]) -> Vec<serde_json::Value> {
rows.iter()
.map(|row| {
let mut document = serde_json::Map::new();
for (index, column) in columns.iter().enumerate() {
document.insert(column.clone(), row.get(index).cloned().unwrap_or(serde_json::Value::Null));
}
serde_json::Value::Object(document)
})
.collect()
}
async fn find_mongo_documents_for_transfer(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
offset: u64,
batch_size: usize,
) -> Result<MongoDocumentResult, String> {
crate::mongo_ops::mongo_find_documents_core(
state,
connection_id,
database,
collection,
offset,
batch_size as i64,
None,
Some(r#"{"_id":1}"#),
)
.await
}
async fn insert_mongo_documents_for_transfer(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
documents: &[serde_json::Value],
) -> Result<u64, String> {
if documents.is_empty() {
return Ok(0);
}
let docs_json = serde_json::to_string(documents).map_err(|e| format!("Failed to encode MongoDB documents: {e}"))?;
match crate::mongo_ops::mongo_insert_documents_core(state, connection_id, database, collection, &docs_json).await {
Ok(count) => Ok(count),
Err(error) if error.to_ascii_lowercase().contains("legacy agent") => {
let mut inserted = 0;
for document in documents {
let doc_json =
serde_json::to_string(document).map_err(|e| format!("Failed to encode MongoDB document: {e}"))?;
crate::mongo_ops::mongo_insert_document_core(state, connection_id, database, collection, &doc_json)
.await?;
inserted += 1;
}
Ok(inserted)
}
Err(error) => Err(error),
}
}
async fn overwrite_mongo_collection_for_transfer(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
) -> Result<(), String> {
crate::mongo_ops::mongo_delete_documents_core(state, connection_id, database, collection, "{}", true)
.await
.map(|_| ())
}
fn mongo_value_column_type(value: Option<&serde_json::Value>) -> String {
match value {
Some(serde_json::Value::Bool(_)) => "boolean".to_string(),
Some(serde_json::Value::Number(number)) if number.is_i64() || number.is_u64() => "bigint".to_string(),
Some(serde_json::Value::Number(_)) => "double".to_string(),
Some(serde_json::Value::Array(_) | serde_json::Value::Object(_)) => "json".to_string(),
_ => "text".to_string(),
}
}
fn mongo_columns_from_documents(documents: &[serde_json::Value]) -> Vec<db::ColumnInfo> {
mongo_transfer_document_fields(documents)
.into_iter()
.map(|name| {
let sample =
documents.iter().filter_map(|document| document.as_object()?.get(&name)).find(|value| !value.is_null());
db::ColumnInfo {
name,
data_type: mongo_value_column_type(sample),
is_nullable: true,
column_default: None,
is_primary_key: false,
extra: None,
comment: None,
numeric_precision: None,
numeric_scale: None,
character_maximum_length: None,
}
})
.collect()
}
pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result<db::QueryResult, String> {
execute_on_pool_with_max_rows(state, pool_key, sql, None).await
}
pub async fn execute_on_pool_with_max_rows(
state: &AppState,
pool_key: &str,
sql: &str,
max_rows: Option<usize>,
) -> 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_with_max_rows(&p, sql, bare, max_rows, Default::default()).await
}
PoolKind::Postgres(p) => {
let p = p.clone();
drop(connections);
db::postgres::execute_query_with_max_rows(&p, sql, max_rows).await
}
PoolKind::Sqlite(p) => {
let p = p.clone();
drop(connections);
db::sqlite::execute_query_with_max_rows(&p, sql, max_rows).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_with_max_rows(&client, &database, sql, max_rows).await
}
PoolKind::SqlServer(client) => {
let client = client.clone();
drop(connections);
let mut client = client.lock().await;
db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows).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, fetch_size: max_rows, ..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())?;
if max_rows.is_some()
&& starts_with_executable_sql_keyword(&sql, &["SELECT", "SHOW", "DESCRIBE", "WITH", "PRAGMA"])
{
return crate::query::duckdb_execute_with_max_rows(&con, &sql, max_rows);
}
let start = std::time::Instant::now();
if starts_with_executable_sql_keyword(&sql, &["SELECT", "SHOW", "DESCRIBE", "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_with_max_rows(&con, &sql, max_rows)
})
.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).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()),
}
}
async fn get_postgres_indexes_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
table: &str,
) -> Result<Vec<db::IndexInfo>, String> {
let connections = state.connections.read().await;
let Some(PoolKind::Postgres(pool)) = connections.get(pool_key) else {
return Err("PostgreSQL pool not found".to_string());
};
let pool = pool.clone();
drop(connections);
db::postgres::list_indexes(&pool, schema, table).await
}
async fn get_postgres_foreign_keys_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
table: &str,
) -> Result<Vec<db::ForeignKeyInfo>, String> {
let connections = state.connections.read().await;
let Some(PoolKind::Postgres(pool)) = connections.get(pool_key) else {
return Err("PostgreSQL pool not found".to_string());
};
let pool = pool.clone();
drop(connections);
db::postgres::list_foreign_keys(&pool, schema, table).await
}
async fn get_postgres_schema_object_sources_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
) -> Result<Vec<db::ObjectSource>, String> {
let views_sql = format!(
"SELECT c.relname, pg_get_viewdef(c.oid, true) \
FROM pg_catalog.pg_class c \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = {} AND c.relkind = 'v' \
ORDER BY c.relname",
quote_string_literal(schema)
);
let routines_sql = format!(
"SELECT p.proname, CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, pg_get_functiondef(p.oid) \
FROM pg_catalog.pg_proc p \
JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \
WHERE n.nspname = {} AND p.prokind IN ('p', 'f') \
ORDER BY CASE p.prokind WHEN 'p' THEN 0 ELSE 1 END, p.proname, p.oid",
quote_string_literal(schema)
);
let mut sources = Vec::new();
for row in execute_on_pool(state, pool_key, &views_sql).await?.rows {
let Some(name) = json_string_cell(&row, 0) else {
continue;
};
let Some(source) = json_string_cell(&row, 1) else {
continue;
};
sources.push(db::ObjectSource {
name,
object_type: db::ObjectSourceKind::View,
schema: Some(schema.to_string()),
source,
});
}
for row in execute_on_pool(state, pool_key, &routines_sql).await?.rows {
let Some(name) = json_string_cell(&row, 0) else {
continue;
};
let kind = match json_string_cell(&row, 1).as_deref() {
Some("PROCEDURE") => db::ObjectSourceKind::Procedure,
_ => db::ObjectSourceKind::Function,
};
let Some(source) = json_string_cell(&row, 2) else {
continue;
};
sources.push(db::ObjectSource { name, object_type: kind, schema: Some(schema.to_string()), source });
}
Ok(sources)
}
async fn get_postgres_materialized_view_sources_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
) -> Result<Vec<PostgresMaterializedViewSource>, String> {
let sql = format!(
"SELECT c.relname, pg_get_viewdef(c.oid, true) \
FROM pg_catalog.pg_class c \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = {} AND c.relkind = 'm' \
ORDER BY c.relname",
quote_string_literal(schema)
);
let rows = execute_on_pool(state, pool_key, &sql).await?.rows;
Ok(rows
.into_iter()
.filter_map(|row| {
Some(PostgresMaterializedViewSource {
view_name: json_string_cell(&row, 0)?,
source: json_string_cell(&row, 1)?,
})
})
.collect())
}
async fn get_postgres_trigger_sources_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
tables: &[String],
) -> Result<Vec<PostgresTriggerSource>, String> {
if tables.is_empty() {
return Ok(Vec::new());
}
let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::<Vec<_>>().join(", ");
let sql = format!(
"SELECT c.relname, t.tgname, pg_get_triggerdef(t.oid, true) \
FROM pg_catalog.pg_trigger t \
JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = {} AND NOT t.tgisinternal AND c.relname IN ({table_list}) \
ORDER BY c.relname, t.tgname",
quote_string_literal(schema)
);
let rows = execute_on_pool(state, pool_key, &sql).await?.rows;
Ok(rows
.into_iter()
.filter_map(|row| {
Some(PostgresTriggerSource {
table_name: json_string_cell(&row, 0)?,
trigger_name: json_string_cell(&row, 1)?,
source: json_string_cell(&row, 2)?,
})
})
.collect())
}
async fn get_postgres_extension_sources_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
) -> Result<Vec<PostgresExtensionSource>, String> {
let sql = format!(
"SELECT e.extname \
FROM pg_extension e \
JOIN pg_namespace n ON n.oid = e.extnamespace \
WHERE n.nspname = {} \
ORDER BY e.extname",
quote_string_literal(schema)
);
let rows = execute_on_pool(state, pool_key, &sql).await?.rows;
Ok(rows
.into_iter()
.filter_map(|row| json_string_cell(&row, 0).map(|extension_name| PostgresExtensionSource { extension_name }))
.collect())
}
async fn get_postgres_enum_sources_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
) -> Result<Vec<PostgresEnumSource>, String> {
let sql = format!(
"SELECT t.typname, COALESCE(array_to_json(array_agg(e.enumlabel ORDER BY e.enumsortorder))::text, '[]') \
FROM pg_type t \
JOIN pg_namespace n ON n.oid = t.typnamespace \
LEFT JOIN pg_enum e ON e.enumtypid = t.oid \
WHERE n.nspname = {} AND t.typtype = 'e' \
GROUP BY t.typname \
ORDER BY t.typname",
quote_string_literal(schema)
);
let rows = execute_on_pool(state, pool_key, &sql).await?.rows;
Ok(rows
.into_iter()
.filter_map(|row| {
let type_name = json_string_cell(&row, 0)?;
let labels_json = json_string_cell(&row, 1)?;
let labels = serde_json::from_str::<Vec<String>>(&labels_json).ok()?;
Some(PostgresEnumSource { type_name, labels })
})
.collect())
}
async fn get_postgres_domain_sources_for_transfer(
state: &AppState,
pool_key: &str,
schema: &str,
) -> Result<Vec<PostgresDomainSource>, String> {
let sql = format!(
"SELECT t.typname, \
pg_catalog.format_type(t.typbasetype, t.typtypmod), \
NULLIF(t.typdefault, ''), \
t.typnotnull, \
COALESCE(( \
SELECT array_to_json(array_agg(pg_get_constraintdef(c.oid, true) ORDER BY c.conname))::text \
FROM pg_constraint c \
WHERE c.contypid = t.oid AND c.contype = 'c' \
), '[]') \
FROM pg_type t \
JOIN pg_namespace n ON n.oid = t.typnamespace \
WHERE n.nspname = {} AND t.typtype = 'd' \
ORDER BY t.typname",
quote_string_literal(schema)
);
let rows = execute_on_pool(state, pool_key, &sql).await?.rows;
Ok(rows
.into_iter()
.filter_map(|row| {
let domain_name = json_string_cell(&row, 0)?;
let base_type = json_string_cell(&row, 1)?;
let default_value = json_string_cell(&row, 2);
let not_null = row.get(3).and_then(|value| value.as_bool()).unwrap_or(false);
let checks = json_string_cell(&row, 4)
.and_then(|json| serde_json::from_str::<Vec<String>>(&json).ok())
.unwrap_or_default();
Some(PostgresDomainSource { domain_name, base_type, default_value, not_null, checks })
})
.collect())
}
async fn get_postgres_policy_statements_for_transfer(
state: &AppState,
pool_key: &str,
source_schema: &str,
target_schema: &str,
tables: &[String],
) -> Result<Vec<String>, String> {
if tables.is_empty() {
return Ok(Vec::new());
}
let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::<Vec<_>>().join(", ");
let sql = format!(
"WITH selected_tables AS ( \
SELECT c.oid, c.relname, c.relrowsecurity, c.relforcerowsecurity \
FROM pg_catalog.pg_class c \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = {source_schema} AND c.relkind IN ('r','p') AND c.relname IN ({table_list}) \
), \
policy_rows AS ( \
SELECT t.relname, t.relrowsecurity, t.relforcerowsecurity, p.polname, p.polpermissive, p.polcmd, \
COALESCE((SELECT string_agg(CASE WHEN role_oid = 0 THEN 'PUBLIC' ELSE quote_ident(r.rolname) END, ', ' ORDER BY CASE WHEN role_oid = 0 THEN '' ELSE r.rolname END) \
FROM unnest(p.polroles) AS role_oid LEFT JOIN pg_roles r ON r.oid = role_oid), '') AS role_list, \
pg_get_expr(p.polqual, p.polrelid) AS using_expr, \
pg_get_expr(p.polwithcheck, p.polrelid) AS with_check_expr \
FROM selected_tables t \
JOIN pg_catalog.pg_policy p ON p.polrelid = t.oid \
) \
SELECT stmt FROM ( \
SELECT format('ALTER TABLE %I.%I ENABLE ROW LEVEL SECURITY', {target_schema}, relname) AS stmt, relname, 0 AS sort_order \
FROM selected_tables WHERE relrowsecurity \
UNION ALL \
SELECT format('ALTER TABLE %I.%I FORCE ROW LEVEL SECURITY', {target_schema}, relname) AS stmt, relname, 1 AS sort_order \
FROM selected_tables WHERE relforcerowsecurity \
UNION ALL \
SELECT format('DROP POLICY IF EXISTS %I ON %I.%I', polname, {target_schema}, relname) AS stmt, relname, 2 AS sort_order \
FROM policy_rows \
UNION ALL \
SELECT format( \
'CREATE POLICY %I ON %I.%I AS %s FOR %s%s%s%s', \
polname, {target_schema}, relname, \
CASE WHEN polpermissive THEN 'PERMISSIVE' ELSE 'RESTRICTIVE' END, \
CASE polcmd WHEN 'r' THEN 'SELECT' WHEN 'a' THEN 'INSERT' WHEN 'w' THEN 'UPDATE' WHEN 'd' THEN 'DELETE' ELSE 'ALL' END, \
CASE WHEN role_list <> '' THEN ' TO ' || role_list ELSE '' END, \
CASE WHEN using_expr IS NOT NULL THEN ' USING (' || using_expr || ')' ELSE '' END, \
CASE WHEN with_check_expr IS NOT NULL THEN ' WITH CHECK (' || with_check_expr || ')' ELSE '' END \
) AS stmt, relname, 3 AS sort_order \
FROM policy_rows \
) statements \
ORDER BY relname, sort_order, stmt",
source_schema = quote_string_literal(source_schema),
target_schema = quote_string_literal(target_schema),
);
Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows))
}
async fn get_postgres_ownership_statements_for_transfer(
state: &AppState,
pool_key: &str,
source_schema: &str,
target_schema: &str,
tables: &[String],
) -> Result<Vec<String>, String> {
let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::<Vec<_>>().join(", ");
let table_filter = if tables.is_empty() { "FALSE".to_string() } else { format!("c.relname IN ({table_list})") };
let sql = format!(
"WITH relation_owners AS ( \
SELECT CASE c.relkind \
WHEN 'm' THEN format('ALTER MATERIALIZED VIEW %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \
WHEN 'v' THEN format('ALTER VIEW %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \
WHEN 'f' THEN format('ALTER FOREIGN TABLE %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \
WHEN 'S' THEN format('ALTER SEQUENCE %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \
ELSE format('ALTER TABLE %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \
END AS stmt \
FROM pg_catalog.pg_class c \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = {source_schema} AND (c.relkind IN ('v','m') OR ({table_filter} AND c.relkind IN ('r','p','f','S'))) \
), \
routine_owners AS ( \
SELECT format('ALTER %s %I.%I(%s) OWNER TO %I', \
CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \
{target_schema}, p.proname, pg_get_function_identity_arguments(p.oid), pg_get_userbyid(p.proowner)) AS stmt \
FROM pg_catalog.pg_proc p \
JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \
WHERE n.nspname = {source_schema} AND p.prokind IN ('p','f') \
), \
type_owners AS ( \
SELECT format('ALTER %s %I.%I OWNER TO %I', \
CASE t.typtype WHEN 'd' THEN 'DOMAIN' ELSE 'TYPE' END, \
{target_schema}, t.typname, pg_get_userbyid(t.typowner)) AS stmt \
FROM pg_catalog.pg_type t \
JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \
WHERE n.nspname = {source_schema} AND t.typtype IN ('e','d') \
) \
SELECT stmt FROM ( \
SELECT format('ALTER SCHEMA %I OWNER TO %I', {target_schema}, pg_get_userbyid(n.nspowner)) AS stmt \
FROM pg_catalog.pg_namespace n WHERE n.nspname = {source_schema} \
UNION ALL SELECT stmt FROM relation_owners \
UNION ALL SELECT stmt FROM routine_owners \
UNION ALL SELECT stmt FROM type_owners \
) statements",
source_schema = quote_string_literal(source_schema),
target_schema = quote_string_literal(target_schema),
table_filter = table_filter,
);
Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows))
}
async fn get_postgres_grant_statements_for_transfer(
state: &AppState,
pool_key: &str,
source_schema: &str,
target_schema: &str,
tables: &[String],
) -> Result<Vec<String>, String> {
let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::<Vec<_>>().join(", ");
let table_filter = if tables.is_empty() { "FALSE".to_string() } else { format!("c.relname IN ({table_list})") };
let sql = format!(
"WITH schema_grants AS ( \
SELECT format( \
'GRANT %s ON SCHEMA %I TO %s%s', \
string_agg(a.privilege_type, ', ' ORDER BY a.privilege_type), \
{target_schema}, \
CASE WHEN a.grantee = 0 THEN 'PUBLIC' ELSE quote_ident(grantee.rolname) END, \
CASE WHEN bool_or(a.is_grantable) THEN ' WITH GRANT OPTION' ELSE '' END \
) AS stmt \
FROM pg_catalog.pg_namespace n \
JOIN LATERAL aclexplode(n.nspacl) a ON true \
LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \
WHERE n.nspname = {source_schema} \
GROUP BY a.grantee, grantee.rolname \
), \
relation_grants AS ( \
SELECT format( \
'GRANT %s ON %s %I.%I TO %s%s', \
string_agg(a.privilege_type, ', ' ORDER BY a.privilege_type), \
CASE WHEN relkind = 'S' THEN 'SEQUENCE' ELSE 'TABLE' END, \
{target_schema}, relname, \
CASE WHEN a.grantee = 0 THEN 'PUBLIC' ELSE quote_ident(grantee.rolname) END, \
CASE WHEN bool_or(a.is_grantable) THEN ' WITH GRANT OPTION' ELSE '' END \
) AS stmt \
FROM ( \
SELECT c.relname, c.relkind, a.grantee, a.privilege_type, a.is_grantable, grantee.rolname \
FROM pg_catalog.pg_class c \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
JOIN LATERAL aclexplode(c.relacl) a ON true \
LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \
WHERE n.nspname = {source_schema} AND (c.relkind IN ('v','m') OR ({table_filter} AND c.relkind IN ('r','p','f','S'))) \
) rels \
GROUP BY relname, relkind, grantee, rolname \
), \
routine_grants AS ( \
SELECT format( \
'GRANT %s ON %s %I.%I(%s) TO %s%s', \
string_agg(a.privilege_type, ', ' ORDER BY a.privilege_type), \
CASE WHEN prokind = 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \
{target_schema}, proname, identity_args, \
CASE WHEN a.grantee = 0 THEN 'PUBLIC' ELSE quote_ident(grantee.rolname) END, \
CASE WHEN bool_or(a.is_grantable) THEN ' WITH GRANT OPTION' ELSE '' END \
) AS stmt \
FROM ( \
SELECT p.proname, p.prokind, pg_get_function_identity_arguments(p.oid) AS identity_args, a.grantee, a.privilege_type, a.is_grantable, grantee.rolname \
FROM pg_catalog.pg_proc p \
JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \
JOIN LATERAL aclexplode(p.proacl) a ON true \
LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \
WHERE n.nspname = {source_schema} AND p.prokind IN ('p','f') \
) routines \
GROUP BY proname, prokind, identity_args, grantee, rolname \
) \
SELECT stmt FROM ( \
SELECT stmt FROM schema_grants \
UNION ALL SELECT stmt FROM relation_grants \
UNION ALL SELECT stmt FROM routine_grants \
) statements \
WHERE stmt IS NOT NULL",
source_schema = quote_string_literal(source_schema),
target_schema = quote_string_literal(target_schema),
table_filter = table_filter,
);
Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows))
}
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);
}
#[allow(clippy::too_many_arguments)]
async fn transfer_mongodb_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();
let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size };
let mut offset: u64 = 0;
let mut total_transferred: u64 = 0;
let mut total_rows = None;
if request.mode == TransferMode::Upsert {
log::warn!("[transfer] MongoDB upsert is not supported yet, falling back to append");
}
if is_mongodb_transfer_type(target_db_type) && request.mode == TransferMode::Overwrite {
overwrite_mongo_collection_for_transfer(state, &request.target_connection_id, &request.target_database, table)
.await
.map_err(|e| format!("Failed to clear MongoDB collection '{table}': {e}"))?;
}
let mut sql_target_column_names: Vec<String> = Vec::new();
let mut sql_target_column_types: Vec<Option<String>> = Vec::new();
let mut sql_target_prepared = false;
loop {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
let documents = if is_mongodb_transfer_type(source_db_type) {
let result = find_mongo_documents_for_transfer(
state,
&request.source_connection_id,
&request.source_database,
table,
offset,
batch_size,
)
.await?;
total_rows = Some(result.total);
result.documents
} else {
let columns = get_columns_for_transfer(
state,
source_pool_key,
&request.source_connection_id,
&request.source_database,
&request.source_schema,
table,
)
.await?;
let col_names = columns.iter().map(|column| column.name.clone()).collect::<Vec<_>>();
let primary_key_columns = columns
.iter()
.filter(|column| column.is_primary_key)
.map(|column| column.name.clone())
.collect::<Vec<_>>();
let sql = pagination_sql_with_order(
&col_names,
table,
&request.source_schema,
source_db_type,
offset,
batch_size,
&primary_key_columns,
);
let result = execute_on_pool(state, source_pool_key, &sql).await?;
sql_rows_to_mongo_documents(&col_names, &result.rows)
};
let row_count = documents.len();
if row_count == 0 {
break;
}
if is_mongodb_transfer_type(target_db_type) {
insert_mongo_documents_for_transfer(
state,
&request.target_connection_id,
&request.target_database,
table,
&documents,
)
.await
.map_err(|e| format!("Insert failed for MongoDB collection '{table}' at offset {offset}: {e}"))?;
} else {
if !sql_target_prepared {
let mut sql_target_columns = mongo_columns_from_documents(&documents);
if sql_target_columns.is_empty() {
sql_target_columns.push(db::ColumnInfo {
name: "document".to_string(),
data_type: "json".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,
});
}
sql_target_column_names = sql_target_columns.iter().map(|column| column.name.clone()).collect();
sql_target_column_types =
sql_target_columns.iter().map(|column| Some(column.data_type.clone())).collect();
if request.create_table {
let ddl = generate_create_table_ddl(
&sql_target_columns,
table,
&request.source_schema,
&request.target_schema,
target_db_type,
source_db_type,
None,
);
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 from MongoDB collection '{table}': {e}"));
}
}
};
if table_exists {
for stmt in generate_comment_ddl(
&sql_target_columns,
table,
&request.target_schema,
target_db_type,
None,
) {
if let Err(e) = execute_on_pool(state, target_pool_key, &stmt).await {
log::warn!(
"[transfer] failed to set MongoDB transfer column comment for {}: {}",
table,
e
);
}
}
}
}
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 MongoDB transfer target table: {e}"))?;
}
sql_target_prepared = true;
}
let rows = if sql_target_column_names.len() == 1 && sql_target_column_names[0] == "document" {
documents.iter().map(|document| vec![document.clone()]).collect::<Vec<_>>()
} else {
mongo_documents_to_rows(&documents, &sql_target_column_names)
};
let write_statements = generate_transfer_write_sql_batches(
&TransferMode::Append,
&sql_target_column_names,
&sql_target_column_types,
&rows,
table,
&request.target_schema,
target_db_type,
&[],
);
for (statement_index, batch_sql) in write_statements.iter().enumerate() {
execute_on_pool(state, target_pool_key, batch_sql).await.map_err(|e| {
format!(
"Insert failed for MongoDB collection '{table}' at offset {offset}, chunk {} of {}: {e}",
statement_index + 1,
write_statements.len()
)
})?;
}
}
total_transferred += row_count as u64;
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)
}
/// Transfer a single table. Returns rows transferred.
/// `progress_callback` is invoked for progress updates.
#[allow(clippy::too_many_arguments)]
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),
{
if is_mongodb_transfer_type(source_db_type) || is_mongodb_transfer_type(target_db_type) {
return transfer_mongodb_table(
state,
request,
table,
table_index,
source_db_type,
target_db_type,
source_pool_key,
target_pool_key,
progress_callback,
)
.await;
}
let total_tables = request.tables.len();
let pg_compat_transfer = is_postgres_compat_transfer(source_db_type, target_db_type);
// 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();
let primary_key_columns: Vec<String> =
columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.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);
let target_table_preexisting = crate::schema::list_tables_core(
state,
&request.target_connection_id,
&request.target_database,
&request.target_schema,
Some(table),
Some(1),
)
.await
.map(|tables| !tables.is_empty())
.unwrap_or(false);
let source_indexes = if request.create_table && pg_compat_transfer && !target_table_preexisting {
get_postgres_indexes_for_transfer(state, source_pool_key, &request.source_schema, table).await?
} else {
Vec::new()
};
let source_foreign_keys = if request.create_table && pg_compat_transfer && !target_table_preexisting {
get_postgres_foreign_keys_for_transfer(state, source_pool_key, &request.source_schema, table).await?
} else {
Vec::new()
};
// 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 {
if matches!(target_db_type, DatabaseType::Postgres) && !request.target_schema.trim().is_empty() {
let create_schema_sql =
format!("CREATE SCHEMA IF NOT EXISTS {}", quote_identifier(&request.target_schema, target_db_type));
execute_on_pool(state, target_pool_key, &create_schema_sql)
.await
.map_err(|e| format!("Failed to ensure schema exists: {e}"))?;
}
let ddl = generate_create_table_ddl(
&columns,
table,
&request.source_schema,
&request.target_schema,
target_db_type,
source_db_type,
table_comment.as_deref(),
);
log::info!("[transfer] creating target table: {}", ddl.chars().take(200).collect::<String>());
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 | DatabaseType::Hive) {
log::warn!("[transfer] upsert not supported for {:?}, falling back to append", target_db_type);
(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_with_order(
&col_names,
table,
&request.source_schema,
source_db_type,
offset,
batch_size,
&primary_key_columns,
);
let result = execute_on_pool(state, source_pool_key, &sql).await?;
let row_count = result.rows.len();
if row_count == 0 {
break;
}
let write_statements = generate_transfer_write_sql_batches(
&effective_mode,
&col_names,
&col_types,
&result.rows,
table,
&request.target_schema,
target_db_type,
&pk_columns,
);
for (statement_index, batch_sql) in write_statements.iter().enumerate() {
execute_on_pool(state, target_pool_key, batch_sql).await.map_err(|e| {
let absolute_row = parse_mysql_row_error(&e).map(|row| offset + row);
match absolute_row {
Some(row) => format!(
"Insert failed for table '{table}' at row {row} (chunk {} of {}): {e}",
statement_index + 1,
write_statements.len()
),
None => format!(
"Insert failed for table '{table}' at offset {offset}, chunk {} of {}: {e}",
statement_index + 1,
write_statements.len()
),
}
})?;
}
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;
}
}
if pg_compat_transfer {
for statement in generate_postgres_sequence_sync_sql(&columns, table, &request.target_schema) {
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to sync PostgreSQL sequence for {table}: {e}"))?;
}
}
if request.create_table && pg_compat_transfer && !target_table_preexisting {
for statement in generate_postgres_index_ddl(&source_indexes, table, &request.target_schema) {
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to create PostgreSQL index for {table}: {e}"))?;
}
for statement in generate_postgres_foreign_key_ddl(&source_foreign_keys, table, &request.target_schema) {
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to create PostgreSQL foreign key for {table}: {e}"))?;
}
}
Ok(total_transferred)
}
pub async fn transfer_postgres_schema_dependencies<F>(
state: &AppState,
request: &TransferRequest,
source_pool_key: &str,
target_pool_key: &str,
mut progress_callback: F,
) -> Result<(), String>
where
F: FnMut(TransferProgress),
{
let source_db_type = get_db_type(state, &request.source_connection_id).await?;
let target_db_type = get_db_type(state, &request.target_connection_id).await?;
if !request.create_table || !is_postgres_compat_transfer(&source_db_type, &target_db_type) {
return Ok(());
}
if !request.target_schema.trim().is_empty() {
let create_schema_sql = format!(
"CREATE SCHEMA IF NOT EXISTS {}",
quote_identifier(&request.target_schema, &DatabaseType::Postgres)
);
execute_on_pool(state, target_pool_key, &create_schema_sql)
.await
.map_err(|e| format!("Failed to ensure PostgreSQL target schema exists: {e}"))?;
}
let extensions =
get_postgres_extension_sources_for_transfer(state, source_pool_key, &request.source_schema).await?;
let enum_types = get_postgres_enum_sources_for_transfer(state, source_pool_key, &request.source_schema).await?;
let domains = get_postgres_domain_sources_for_transfer(state, source_pool_key, &request.source_schema).await?;
let total_steps = extensions.len() + enum_types.len() + domains.len();
let table_index = 0;
let mut completed_steps = 0_u64;
for extension in extensions {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("extension: {}", extension.extension_name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &generate_postgres_extension_ddl(&extension, &request.target_schema))
.await
.map_err(|e| format!("Failed to create PostgreSQL extension {}: {e}", extension.extension_name))?;
}
for enum_type in enum_types {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("enum: {}", enum_type.type_name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &generate_postgres_enum_ddl(&enum_type, &request.target_schema))
.await
.map_err(|e| format!("Failed to create PostgreSQL enum {}: {e}", enum_type.type_name))?;
}
for domain in domains {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("domain: {}", domain.domain_name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &generate_postgres_domain_ddl(&domain, &request.target_schema))
.await
.map_err(|e| format!("Failed to create PostgreSQL domain {}: {e}", domain.domain_name))?;
}
Ok(())
}
pub async fn transfer_postgres_schema_objects<F>(
state: &AppState,
request: &TransferRequest,
source_pool_key: &str,
target_pool_key: &str,
mut progress_callback: F,
) -> Result<(), String>
where
F: FnMut(TransferProgress),
{
let source_db_type = get_db_type(state, &request.source_connection_id).await?;
let target_db_type = get_db_type(state, &request.target_connection_id).await?;
if !request.create_table || !is_postgres_compat_transfer(&source_db_type, &target_db_type) {
return Ok(());
}
let object_sources =
get_postgres_schema_object_sources_for_transfer(state, source_pool_key, &request.source_schema).await?;
let materialized_views =
get_postgres_materialized_view_sources_for_transfer(state, source_pool_key, &request.source_schema).await?;
let trigger_sources =
get_postgres_trigger_sources_for_transfer(state, source_pool_key, &request.source_schema, &request.tables)
.await?;
let policy_statements = get_postgres_policy_statements_for_transfer(
state,
source_pool_key,
&request.source_schema,
&request.target_schema,
&request.tables,
)
.await?;
let ownership_statements = get_postgres_ownership_statements_for_transfer(
state,
source_pool_key,
&request.source_schema,
&request.target_schema,
&request.tables,
)
.await?;
let grant_statements = get_postgres_grant_statements_for_transfer(
state,
source_pool_key,
&request.source_schema,
&request.target_schema,
&request.tables,
)
.await?;
let materialized_view_step_count = materialized_views
.iter()
.map(|view| generate_postgres_materialized_view_ddls(view, &request.target_schema).len())
.sum::<usize>();
let trigger_step_count = trigger_sources.len() * 2;
let total_steps = object_sources.len()
+ materialized_view_step_count
+ trigger_step_count
+ policy_statements.len()
+ ownership_statements.len()
+ grant_statements.len();
let table_index = request.tables.len();
let mut completed_steps = 0_u64;
for object in object_sources {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("schema object: {}", object.name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
let rewritten_source = match object.object_type {
db::ObjectSourceKind::View => object.source.clone(),
db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function => {
rewrite_postgres_routine_schema(&object.source, &request.target_schema)
.unwrap_or_else(|| object.source.clone())
}
db::ObjectSourceKind::Package | db::ObjectSourceKind::PackageBody => object.source.clone(),
};
let statements = build_executable_object_source_statements(EditableObjectSourceSqlInput {
database_type: DatabaseType::Postgres,
object_type: object.object_type.clone(),
schema: Some(request.target_schema.clone()),
name: object.name.clone(),
source: rewritten_source,
})?;
for statement in statements {
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to create PostgreSQL {:?} {}: {e}", object.object_type, object.name))?;
}
}
for view in materialized_views {
for statement in generate_postgres_materialized_view_ddls(&view, &request.target_schema) {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("materialized view: {}", view.view_name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to create PostgreSQL materialized view {}: {e}", view.view_name))?;
}
}
for trigger in trigger_sources {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("trigger: {}", trigger.trigger_name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
let full_table = qualified_table(&trigger.table_name, &request.target_schema, &DatabaseType::Postgres);
let drop_sql = format!(
"DROP TRIGGER IF EXISTS {} ON {full_table}",
quote_identifier(&trigger.trigger_name, &DatabaseType::Postgres)
);
execute_on_pool(state, target_pool_key, &drop_sql)
.await
.map_err(|e| format!("Failed to drop PostgreSQL trigger {}: {e}", trigger.trigger_name))?;
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: format!("trigger: {}", trigger.trigger_name),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
let create_sql = rewrite_postgres_trigger_table_schema(
&ensure_sql_statement_terminated(&trigger.source),
&request.source_schema,
&trigger.table_name,
&request.target_schema,
);
execute_on_pool(state, target_pool_key, &create_sql)
.await
.map_err(|e| format!("Failed to create PostgreSQL trigger {}: {e}", trigger.trigger_name))?;
}
for statement in policy_statements {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: "row security policies".to_string(),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to apply PostgreSQL row security statement: {e}"))?;
}
for statement in ownership_statements {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: "ownership".to_string(),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to apply PostgreSQL ownership statement: {e}"))?;
}
for statement in grant_statements {
if is_cancelled(&request.transfer_id).await {
return Err("Cancelled".to_string());
}
completed_steps += 1;
progress_callback(TransferProgress {
transfer_id: request.transfer_id.clone(),
table: "grants".to_string(),
table_index,
total_tables: request.tables.len(),
rows_transferred: completed_steps,
total_rows: Some(total_steps as u64),
status: TransferStatus::Running,
error: None,
});
execute_on_pool(state, target_pool_key, &statement)
.await
.map_err(|e| format!("Failed to apply PostgreSQL grant statement: {e}"))?;
}
Ok(())
}
#[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,
transport_layers: Vec::new(),
connect_timeout_secs: 5,
query_timeout_secs: 30,
ssl: false,
ca_cert_path: String::new(),
sysdba: false,
oracle_connection_type: None,
connection_string: None,
redis_connection_mode: None,
redis_sentinel_master: String::new(),
redis_sentinel_nodes: String::new(),
redis_sentinel_username: String::new(),
redis_sentinel_password: String::new(),
redis_sentinel_tls: false,
redis_cluster_nodes: String::new(),
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 postgres_create_table_preserves_defaults_identity_and_exact_types() {
let cols = vec![
db::ColumnInfo {
data_type: "integer".to_string(),
column_default: Some("nextval('public.users_id_seq'::regclass)".to_string()),
is_primary_key: true,
is_nullable: false,
..test_column("id", "integer")
},
db::ColumnInfo {
data_type: "timestamp with time zone".to_string(),
column_default: Some("now()".to_string()),
is_nullable: false,
..test_column("created_at", "timestamp with time zone")
},
db::ColumnInfo {
data_type: "character varying(120)".to_string(),
column_default: Some("'guest'::character varying".to_string()),
..test_column("name", "character varying(120)")
},
];
let ddl = generate_create_table_ddl(
&cols,
"users",
"public",
"public",
&DatabaseType::Postgres,
&DatabaseType::Postgres,
None,
);
assert!(ddl.contains("\"id\" integer GENERATED BY DEFAULT AS IDENTITY NOT NULL"));
assert!(ddl.contains("\"created_at\" timestamp with time zone DEFAULT now() NOT NULL"));
assert!(ddl.contains("\"name\" character varying(120) DEFAULT 'guest'::character varying"));
assert!(ddl.contains("PRIMARY KEY (\"id\")"));
}
#[test]
fn postgres_create_table_rewrites_schema_qualified_custom_types_and_defaults() {
let cols = vec![db::ColumnInfo {
data_type: "\"public\".\"user_status\"".to_string(),
column_default: Some("'active'::public.user_status".to_string()),
is_nullable: false,
..test_column("status", "\"public\".\"user_status\"")
}];
let ddl = generate_create_table_ddl(
&cols,
"users",
"public",
"archive",
&DatabaseType::Postgres,
&DatabaseType::Postgres,
None,
);
assert!(
ddl.contains("\"status\" \"archive\".\"user_status\" DEFAULT 'active'::\"archive\".user_status NOT NULL")
);
}
#[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 hive_create_table_uses_hive_friendly_columns() {
let cols = vec![
db::ColumnInfo { is_primary_key: true, is_nullable: false, ..test_column("id", "bigint") },
db::ColumnInfo { is_nullable: false, ..test_column("payload", "json") },
];
let ddl = generate_create_table_ddl(
&cols,
"events",
"public",
"warehouse",
&DatabaseType::Hive,
&DatabaseType::Postgres,
None,
);
assert!(ddl.contains("CREATE TABLE IF NOT EXISTS `warehouse`.`events`"));
assert!(ddl.contains("`id` BIGINT"));
assert!(ddl.contains("`payload` STRING"));
assert!(!ddl.contains("PRIMARY KEY"));
assert!(!ddl.contains("NOT NULL"));
}
#[test]
fn hive_transfer_uses_backticks_and_hive_type_mapping() {
assert_eq!(quote_identifier("user`events", &DatabaseType::Hive), "`user``events`");
assert_eq!(map_column_type("jsonb", &DatabaseType::Postgres, &DatabaseType::Hive), "STRING");
assert_eq!(
map_column_type("timestamp with time zone", &DatabaseType::Postgres, &DatabaseType::Hive),
"TIMESTAMP"
);
}
#[test]
fn mongo_transfer_document_fields_preserve_first_seen_order() {
let documents = vec![json!({"b": 1}), json!({"a": 2, "c": 3}), json!({"b": 4, "d": 5})];
assert_eq!(mongo_transfer_document_fields(&documents), vec!["b", "a", "c", "d"]);
}
#[test]
fn mongo_transfer_rows_fill_missing_fields_with_null() {
let rows = mongo_documents_to_rows(
&[json!({"id": 1, "name": "Ada"}), json!({"id": 2})],
&[String::from("id"), String::from("name")],
);
assert_eq!(rows, vec![vec![json!(1), json!("Ada")], vec![json!(2), serde_json::Value::Null]]);
}
#[test]
fn sql_rows_to_mongo_documents_maps_columns_to_fields() {
let documents = sql_rows_to_mongo_documents(
&[String::from("id"), String::from("name"), String::from("active")],
&[vec![json!(1), json!("Ada")], vec![json!(2), json!("Grace"), json!(true)]],
);
assert_eq!(
documents,
vec![json!({"id": 1, "name": "Ada", "active": null}), json!({"id": 2, "name": "Grace", "active": true})]
);
}
#[test]
fn postgres_pagination_uses_stable_primary_key_order() {
let sql = pagination_sql_with_order(
&[String::from("id"), String::from("name")],
"users",
"public",
&DatabaseType::Postgres,
200,
100,
&[String::from("id")],
);
assert_eq!(sql, "SELECT \"id\", \"name\" FROM \"public\".\"users\" ORDER BY \"id\" LIMIT 100 OFFSET 200");
}
#[test]
fn filtered_pagination_preserves_where_and_order() {
let sql = pagination_sql_with_filter_order(
&[String::from("id"), String::from("status")],
"users",
"public",
&DatabaseType::SapHana,
10_000,
2_000,
Some("WHERE status = 'active'"),
Some("\"id\" DESC"),
&[String::from("id")],
);
assert_eq!(
sql,
"SELECT \"id\", \"status\" FROM \"public\".\"users\" WHERE (status = 'active') ORDER BY \"id\" DESC LIMIT 2000 OFFSET 10000"
);
}
#[test]
fn filtered_count_preserves_where() {
let sql = count_sql_with_where("users", "public", &DatabaseType::SapHana, Some("WHERE status = 'active'"));
assert_eq!(sql, "SELECT COUNT(*) FROM \"public\".\"users\" WHERE (status = 'active')");
}
#[test]
fn sqlserver_keyset_pagination_includes_offset_fetch() {
let sql = keyset_pagination_sql(
&[String::from("id"), String::from("name")],
"users",
"dbo",
&DatabaseType::SqlServer,
&[String::from("id")],
&[],
100,
);
assert_eq!(
sql,
"SELECT [id], [name] FROM [dbo].[users] ORDER BY [id] ASC OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"
);
}
#[test]
fn composite_keyset_pagination_uses_portable_lexicographic_predicate() {
let sql = keyset_pagination_sql(
&[String::from("tenant_id"), String::from("id"), String::from("name")],
"users",
"dbo",
&DatabaseType::SqlServer,
&[String::from("tenant_id"), String::from("id")],
&[json!(10), json!(25)],
100,
);
assert_eq!(
sql,
"SELECT [tenant_id], [id], [name] FROM [dbo].[users] WHERE ([tenant_id] > 10 OR ([tenant_id] = 10 AND [id] > 25)) ORDER BY [tenant_id] ASC, [id] ASC OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"
);
}
#[test]
fn postgres_generates_index_and_foreign_key_sql() {
let indexes = vec![db::IndexInfo {
name: "users_name_idx".to_string(),
columns: vec!["lower(name)".to_string()],
is_unique: false,
is_primary: false,
filter: Some("name IS NOT NULL".to_string()),
index_type: Some("btree".to_string()),
included_columns: Some(vec!["created_at".to_string()]),
comment: Some("lookup index".to_string()),
}];
let foreign_keys = vec![
db::ForeignKeyInfo {
name: "orders_user_id_fkey".to_string(),
column: "user_id".to_string(),
ref_schema: None,
ref_table: "users".to_string(),
ref_column: "id".to_string(),
},
db::ForeignKeyInfo {
name: "orders_user_id_fkey".to_string(),
column: "tenant_id".to_string(),
ref_schema: None,
ref_table: "users".to_string(),
ref_column: "tenant_id".to_string(),
},
];
let index_sql = generate_postgres_index_ddl(&indexes, "users", "public");
let foreign_key_sql = generate_postgres_foreign_key_ddl(&foreign_keys, "orders", "public");
assert_eq!(
index_sql,
vec![
"CREATE INDEX IF NOT EXISTS \"users_name_idx\" ON \"public\".\"users\" USING btree (lower(name)) INCLUDE (\"created_at\") WHERE name IS NOT NULL".to_string(),
"COMMENT ON INDEX \"public\".\"users_name_idx\" IS 'lookup index'".to_string(),
]
);
assert_eq!(
foreign_key_sql,
vec![
"ALTER TABLE \"public\".\"orders\" ADD CONSTRAINT \"orders_user_id_fkey\" FOREIGN KEY (\"user_id\", \"tenant_id\") REFERENCES \"public\".\"users\" (\"id\", \"tenant_id\")".to_string()
]
);
}
#[test]
fn postgres_sequence_sync_sql_uses_table_max_values() {
let sql = generate_postgres_sequence_sync_sql(
&[db::ColumnInfo {
name: "id".to_string(),
data_type: "integer".to_string(),
is_nullable: false,
column_default: Some("nextval('public.users_id_seq'::regclass)".to_string()),
is_primary_key: true,
extra: None,
comment: None,
numeric_precision: None,
numeric_scale: None,
character_maximum_length: None,
}],
"users",
"public",
);
assert_eq!(
sql,
vec![
"SELECT setval(pg_get_serial_sequence('\"public\".\"users\"', 'id'), GREATEST(COALESCE(MAX(\"id\"), 0), 1), MAX(\"id\") IS NOT NULL) FROM \"public\".\"users\"".to_string()
]
);
}
#[test]
fn postgres_routine_schema_rewrite_targets_destination_schema() {
let rewritten = rewrite_postgres_routine_schema(
"CREATE OR REPLACE FUNCTION public.bump_counter(id integer)\nRETURNS integer\nLANGUAGE plpgsql\nAS $$ BEGIN RETURN id + 1; END; $$",
"archive",
)
.unwrap();
assert!(rewritten.starts_with("CREATE OR REPLACE FUNCTION \"archive\".\"bump_counter\"("));
}
#[test]
fn postgres_trigger_schema_rewrite_targets_destination_table() {
let rewritten = rewrite_postgres_trigger_table_schema(
"CREATE TRIGGER bump BEFORE INSERT ON public.users FOR EACH ROW EXECUTE FUNCTION public.bump_counter()",
"public",
"users",
"archive",
);
assert!(rewritten.contains(" ON \"archive\".\"users\" "));
}
#[test]
fn postgres_extension_enum_and_domain_ddl_is_repeatable() {
let extension_sql = generate_postgres_extension_ddl(
&PostgresExtensionSource { extension_name: "pgcrypto".to_string() },
"archive",
);
let enum_sql = generate_postgres_enum_ddl(
&PostgresEnumSource {
type_name: "status".to_string(),
labels: vec!["pending".to_string(), "done".to_string()],
},
"archive",
);
let domain_sql = generate_postgres_domain_ddl(
&PostgresDomainSource {
domain_name: "email".to_string(),
base_type: "text".to_string(),
default_value: Some("'unknown@example.com'::text".to_string()),
not_null: true,
checks: vec!["CHECK ((VALUE ~* '^[^@]+@[^@]+$'::text))".to_string()],
},
"archive",
);
assert_eq!(extension_sql, "CREATE EXTENSION IF NOT EXISTS \"pgcrypto\" WITH SCHEMA \"archive\"");
assert!(enum_sql.contains("DO $$ BEGIN IF NOT EXISTS"));
assert!(enum_sql.contains("CREATE TYPE \"archive\".\"status\" AS ENUM ('pending', 'done')"));
assert!(domain_sql.contains("CREATE DOMAIN \"archive\".\"email\" AS text DEFAULT 'unknown@example.com'::text NOT NULL CHECK ((VALUE ~* '^[^@]+@[^@]+$'::text))"));
}
#[test]
fn postgres_materialized_view_ddls_drop_and_recreate_in_target_schema() {
let ddls = generate_postgres_materialized_view_ddls(
&PostgresMaterializedViewSource {
view_name: "active_users".to_string(),
source: "SELECT id, name FROM public.users WHERE active".to_string(),
},
"archive",
);
assert_eq!(ddls.len(), 2);
assert_eq!(ddls[0], "DROP MATERIALIZED VIEW IF EXISTS \"archive\".\"active_users\"");
assert_eq!(
ddls[1],
"CREATE MATERIALIZED VIEW \"archive\".\"active_users\" AS\nSELECT id, name FROM public.users WHERE active;"
);
}
#[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_omits_database_qualified_table_name() {
let sql = generate_insert_typed(
&[String::from("id")],
&[Some(String::from("int"))],
&[vec![json!(1)]],
"users",
"app",
&DatabaseType::Mysql,
);
assert_eq!(sql, "INSERT INTO `users` (`id`) VALUES\n(1)");
}
#[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')"
);
}
#[test]
fn transfer_write_sql_batches_split_large_insert_statements() {
let rows = (0..4).map(|index| vec![json!(index), json!("x".repeat(180 * 1024))]).collect::<Vec<_>>();
let statements = generate_transfer_write_sql_batches(
&TransferMode::Append,
&[String::from("id"), String::from("payload")],
&[Some(String::from("int")), Some(String::from("text"))],
&rows,
"events",
"",
&DatabaseType::Mysql,
&[],
);
assert!(statements.len() > 1);
assert!(statements.iter().all(|sql| sql.starts_with("INSERT INTO `events`")));
}
#[test]
fn transfer_write_sql_batches_keep_existing_upsert_sql_shape() {
let statements = generate_transfer_write_sql_batches(
&TransferMode::Upsert,
&[String::from("id"), String::from("name")],
&[Some(String::from("int")), Some(String::from("varchar(64)"))],
&[vec![json!(1), json!("Ada")]],
"users",
"",
&DatabaseType::Mysql,
&[String::from("id")],
);
assert_eq!(statements.len(), 1);
assert!(statements[0].contains("ON DUPLICATE KEY UPDATE"));
}
#[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);
}
#[test]
fn map_column_type_preserves_longtext_for_mysql_target() {
assert_eq!(map_column_type("longtext", &DatabaseType::Mysql, &DatabaseType::Mysql), "LONGTEXT");
}
#[test]
fn map_column_type_preserves_mediumtext_for_mysql_target() {
assert_eq!(map_column_type("mediumtext", &DatabaseType::Mysql, &DatabaseType::Mysql), "MEDIUMTEXT");
}
#[test]
fn map_column_type_preserves_longblob_for_mysql_target() {
assert_eq!(map_column_type("longblob", &DatabaseType::Mysql, &DatabaseType::Mysql), "LONGBLOB");
}
#[test]
fn map_column_type_preserves_mediumblob_for_mysql_target() {
assert_eq!(map_column_type("mediumblob", &DatabaseType::Mysql, &DatabaseType::Mysql), "MEDIUMBLOB");
}
#[test]
fn map_column_type_longtext_falls_back_to_text_for_non_mysql_target() {
assert_eq!(map_column_type("longtext", &DatabaseType::Mysql, &DatabaseType::Postgres), "TEXT");
}
#[test]
fn map_column_type_longblob_falls_back_for_non_mysql_target() {
assert_eq!(map_column_type("longblob", &DatabaseType::Mysql, &DatabaseType::Postgres), "BYTEA");
}
#[test]
fn parse_mysql_row_error_extracts_row_number() {
let err = "ERROR 22001 (1406): Data too long column 'content' at row 8";
assert_eq!(parse_mysql_row_error(err), Some(8));
}
#[test]
fn parse_mysql_row_error_returns_none_for_non_mysql_error() {
assert_eq!(parse_mysql_row_error("some other error"), None);
}
}