diff --git a/.sisyphus/run-continuation/ses_1c1cddfcfffenEHXE36U5julMT.json b/.sisyphus/run-continuation/ses_1c1cddfcfffenEHXE36U5julMT.json deleted file mode 100644 index c59380b..0000000 --- a/.sisyphus/run-continuation/ses_1c1cddfcfffenEHXE36U5julMT.json +++ /dev/null @@ -1,10 +0,0 @@ -{ - "sessionID": "ses_1c1cddfcfffenEHXE36U5julMT", - "updatedAt": "2026-05-19T09:40:24.440Z", - "sources": { - "background-task": { - "state": "idle", - "updatedAt": "2026-05-19T09:40:24.440Z" - } - } -} \ No newline at end of file diff --git a/build/index.html b/build/index.html index c606a09..0f47129 100644 --- a/build/index.html +++ b/build/index.html @@ -5,10 +5,10 @@ veloxdb - - + + - +
diff --git a/public/DuckDB_icon-darkmode.svg b/public/DuckDB_icon-darkmode.svg new file mode 100644 index 0000000..a345735 --- /dev/null +++ b/public/DuckDB_icon-darkmode.svg @@ -0,0 +1,12 @@ + + + + + + + + + + + + diff --git a/public/mongodb-icon-light.svg b/public/mongodb-icon-light.svg new file mode 100644 index 0000000..dc317f7 --- /dev/null +++ b/public/mongodb-icon-light.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/public/redis.svg b/public/redis.svg new file mode 100644 index 0000000..d40f137 --- /dev/null +++ b/public/redis.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 52cf787..13c8cba 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -19,6 +19,20 @@ dependencies = [ "version_check", ] +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -81,6 +95,24 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] + [[package]] name = "arrayref" version = "0.3.9" @@ -93,6 +125,169 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "arrow" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" +dependencies = [ + "ahash 0.8.12", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown 0.17.1", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + +[[package]] +name = "arrow-cast" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ord", + "arrow-schema", + "arrow-select", + "atoi", + "base64 0.22.1", + "chrono", + "comfy-table", + "half", + "lexical-core", + "num-traits", + "ryu", +] + +[[package]] +name = "arrow-data" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-ord" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-row" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" +dependencies = [ + "bitflags 2.11.0", +] + +[[package]] +name = "arrow-select" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" +dependencies = [ + "ahash 0.8.12", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num-traits", +] + +[[package]] +name = "arrow-string" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num-traits", + "regex", + "regex-syntax", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -270,6 +465,29 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "bson" +version = "2.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969a9ba84b0ff843813e7249eed1678d9b6607ce5a3b8f0a47af3fcf7978e6e" +dependencies = [ + "ahash 0.8.12", + "base64 0.22.1", + "bitvec", + "getrandom 0.2.17", + "getrandom 0.3.4", + "hex", + "indexmap 2.13.0", + "js-sys", + "once_cell", + "rand 0.9.2", + "serde", + "serde_bytes", + "serde_json", + "time", + "uuid", +] + [[package]] name = "bstr" version = "1.12.1" @@ -415,6 +633,12 @@ dependencies = [ "toml 0.9.12+spec-1.1.0", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.57" @@ -422,6 +646,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a0dd1ca384932ff3641c8718a02769f1698e7563dc6974ffd03346116310423" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -464,6 +690,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.45" @@ -491,7 +728,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", +] + +[[package]] +name = "comfy-table" +version = "7.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a" +dependencies = [ + "crossterm", + "unicode-segmentation", + "unicode-width", ] [[package]] @@ -509,12 +761,41 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + [[package]] name = "convert_case" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e" +[[package]] +name = "convert_case" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "cookie" version = "0.18.1" @@ -525,6 +806,16 @@ dependencies = [ "version_check", ] +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation" version = "0.10.1" @@ -548,7 +839,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "064badf302c3194842cf2c5d61f56cc88e54a759313879cdf03abdd27d0c3b97" dependencies = [ "bitflags 2.11.0", - "core-foundation", + "core-foundation 0.10.1", "core-graphics-types", "foreign-types", "libc", @@ -561,7 +852,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb" dependencies = [ "bitflags 2.11.0", - "core-foundation", + "core-foundation 0.10.1", "libc", ] @@ -583,6 +874,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc" version = "3.4.0" @@ -607,6 +907,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "critical-section" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -616,6 +922,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.12" @@ -631,6 +946,34 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crossterm" +version = "0.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" +dependencies = [ + "bitflags 2.11.0", + "crossterm_winapi", + "parking_lot", + "rustix 0.38.44", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +dependencies = [ + "winapi", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -752,6 +1095,12 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "data-url" version = "0.3.2" @@ -838,13 +1187,46 @@ dependencies = [ "serde_core", ] +[[package]] +name = "derive-syn-parse" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d65d7ce8132b7c0e54497a4d9a55a1c2a0912a0d786cf894472ba818fba45762" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "derive-where" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d08b3a0bcc0d079199cd476b2cae8435016ec11d1c0986c6901c5ac223041534" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "derive_more" version = "0.99.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6edb4b64a43d977b8e99788fe3a04d483834fba1215a7e02caa415b626497f7f" dependencies = [ - "convert_case", + "convert_case 0.4.0", "proc-macro2", "quote", "rustc_version", @@ -866,10 +1248,12 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb" dependencies = [ + "convert_case 0.10.0", "proc-macro2", "quote", "rustc_version", "syn 2.0.117", + "unicode-xid", ] [[package]] @@ -1011,6 +1395,24 @@ version = "0.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f678cf4a922c215c63e0de95eb1ff08a958a81d47e485cf9da1e27bf6305cfa5" +[[package]] +name = "duckdb" +version = "1.10504.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b9b997383221efd999a362448a0866c54b9a2cb7d4da8cd4903ead4df9f8eaa" +dependencies = [ + "arrow", + "cast", + "comfy-table", + "fallible-iterator 0.3.0", + "fallible-streaming-iterator", + "hashlink", + "libduckdb-sys", + "num-integer", + "rust_decimal", + "strum", +] + [[package]] name = "dunce" version = "1.0.5" @@ -1135,6 +1537,18 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fastrand" version = "2.3.0" @@ -1170,12 +1584,22 @@ dependencies = [ ] [[package]] -name = "find-msvc-tools" -version = "0.1.9" +name = "filetime" +version = "0.2.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" - -[[package]] +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] name = "flagset" version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1189,6 +1613,7 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -1301,6 +1726,21 @@ dependencies = [ "new_debug_unreachable", ] +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -1374,6 +1814,7 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -1549,6 +1990,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -1711,13 +2153,25 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" dependencies = [ - "ahash", + "ahash 0.7.8", ] [[package]] @@ -1737,6 +2191,12 @@ version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "hashlink" version = "0.10.0" @@ -1770,6 +2230,76 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hickory-net" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2295ed2f9c31e471e1428a8f88a3f0e1f4b27c15049592138d1eebe9c35b183" +dependencies = [ + "async-trait", + "cfg-if", + "data-encoding", + "futures-channel", + "futures-io", + "futures-util", + "hickory-proto", + "idna", + "ipnet", + "jni 0.22.4", + "rand 0.10.1", + "thiserror 2.0.18", + "tinyvec", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "hickory-proto" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bab31817bfb44672a252e97fe81cd0c18d1b2cf892108922f6818820df8c643" +dependencies = [ + "data-encoding", + "idna", + "ipnet", + "jni 0.22.4", + "once_cell", + "prefix-trie", + "rand 0.10.1", + "ring", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "url", +] + +[[package]] +name = "hickory-resolver" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d58d28879ceecde6607729660c2667a081ccdc082e082675042793960f178c" +dependencies = [ + "cfg-if", + "futures-util", + "hickory-net", + "hickory-proto", + "ipconfig", + "ipnet", + "jni 0.22.4", + "moka", + "ndk-context", + "once_cell", + "parking_lot", + "rand 0.10.1", + "resolv-conf", + "smallvec", + "system-configuration", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "hkdf" version = "0.12.4" @@ -1912,7 +2442,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.3", "tokio", "tower-service", "tracing", @@ -2114,11 +2644,27 @@ dependencies = [ "cfb", ] +[[package]] +name = "ipconfig" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d40460c0ce33d6ce4b0630ad68ff63d6661961c48b6dba35e5a4d81cfb48222" +dependencies = [ + "socket2 0.6.3", + "widestring", + "windows-registry", + "windows-result 0.4.1", + "windows-sys 0.61.2", +] + [[package]] name = "ipnet" version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +dependencies = [ + "serde", +] [[package]] name = "iri-string" @@ -2168,19 +2714,78 @@ dependencies = [ "cesu8", "cfg-if", "combine", - "jni-sys", + "jni-sys 0.3.0", "log", "thiserror 1.0.69", "walkdir", "windows-sys 0.45.0", ] +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys 0.4.1", + "log", + "simd_cesu8", + "thiserror 2.0.18", + "walkdir", + "windows-link 0.2.1", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.117", +] + [[package]] name = "jni-sys" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8eaf4bc02d17cbdd7ff4c7438cafcdf7fb9a4613313ad11b4f8fefe7d3fa0130" +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.117", +] + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + [[package]] name = "js-sys" version = "0.3.91" @@ -2272,6 +2877,63 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + [[package]] name = "libappindicator" version = "0.9.0" @@ -2311,6 +2973,23 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "libduckdb-sys" +version = "1.10504.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94e67523a40ee3da30411e00e243990b3760d91877202ab2f39c03ab34f8dbfc" +dependencies = [ + "cc", + "flate2", + "pkg-config", + "reqwest 0.12.28", + "serde", + "serde_json", + "tar", + "vcpkg", + "zip", +] + [[package]] name = "libloading" version = "0.7.4" @@ -2356,6 +3035,18 @@ version = "0.5.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f" +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.1" @@ -2409,6 +3100,54 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" +[[package]] +name = "macro_magic" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc33f9f0351468d26fbc53d9ce00a096c8522ecb42f19b50f34f2c422f76d21d" +dependencies = [ + "macro_magic_core", + "macro_magic_macros", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "macro_magic_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1687dc887e42f352865a393acae7cf79d98fab6351cde1f58e9e057da89bf150" +dependencies = [ + "const-random", + "derive-syn-parse", + "macro_magic_core_macros", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "macro_magic_core_macros" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b02abfe41815b5bd98dbd4260173db2c116dda171dc0fe7838cb206333b83308" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "macro_magic_macros" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73ea28ee64b88876bf45277ed9a5817c1817df061a74f2b988971a12570e5869" +dependencies = [ + "macro_magic_core", + "quote", + "syn 2.0.117", +] + [[package]] name = "markup5ever" version = "0.14.1" @@ -2508,14 +3247,107 @@ dependencies = [ ] [[package]] -name = "mio" -version = "1.1.1" +name = "mio" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" +dependencies = [ + "libc", + "wasi 0.11.1+wasi-snapshot-preview1", + "windows-sys 0.61.2", +] + +[[package]] +name = "moka" +version = "0.12.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "mongocrypt" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8da0cd419a51a5fb44819e290fbdb0665a54f21dead8923446a799c7f4d26ad9" +dependencies = [ + "bson", + "mongocrypt-sys", + "once_cell", + "serde", +] + +[[package]] +name = "mongocrypt-sys" +version = "0.1.5+1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224484c5d09285a7b8cb0a0c117e847ebd14cb6e4470ecf68cdb89c503b0edb9" + +[[package]] +name = "mongodb" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "276ba0cd571553d1f6936c6f180964776ece6ab7507dc8765f8a9c9c49d8cd00" +dependencies = [ + "base64 0.22.1", + "bitflags 2.11.0", + "bson", + "derive-where", + "derive_more 2.1.1", + "futures-core", + "futures-io", + "futures-util", + "hex", + "hickory-net", + "hickory-proto", + "hickory-resolver", + "hmac", + "macro_magic", + "md-5", + "mongocrypt", + "mongodb-internal-macros", + "pbkdf2", + "percent-encoding", + "rand 0.9.2", + "rustc_version_runtime", + "rustls", + "serde", + "serde_bytes", + "serde_with", + "sha1", + "sha2", + "socket2 0.6.3", + "stringprep", + "strsim", + "take_mut", + "thiserror 2.0.18", + "tokio", + "tokio-rustls", + "tokio-util", + "typed-builder", + "uuid", + "webpki-roots 1.0.7", +] + +[[package]] +name = "mongodb-internal-macros" +version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" +checksum = "99ceb1a9a1018e470077ec94cf3a8c2d0e6da542b2c05ea95a59a0a627147375" dependencies = [ - "libc", - "wasi 0.11.1+wasi-snapshot-preview1", - "windows-sys 0.61.2", + "macro_magic", + "proc-macro2", + "quote", + "syn 2.0.117", ] [[package]] @@ -2546,7 +3378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3f42e7bbe13d351b6bead8286a43aac9534b82bd3cc43e47037f012ebfd62d4" dependencies = [ "bitflags 2.11.0", - "jni-sys", + "jni-sys 0.3.0", "log", "ndk-sys", "num_enum", @@ -2554,13 +3386,19 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "ndk-context" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" + [[package]] name = "ndk-sys" version = "0.6.0+11769913" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" dependencies = [ - "jni-sys", + "jni-sys 0.3.0", ] [[package]] @@ -2575,6 +3413,16 @@ version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72ef4a56884ca558e5ddb05a1d1e7e1bfd9a68d9ed024c21704cc98872dae1bb" +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint-dig" version = "0.8.6" @@ -2591,6 +3439,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.0" @@ -2878,6 +3735,10 @@ name = "once_cell" version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +dependencies = [ + "critical-section", + "portable-atomic", +] [[package]] name = "option-ext" @@ -2948,6 +3809,15 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest", +] + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -3249,6 +4119,12 @@ dependencies = [ "bstr", ] +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + [[package]] name = "postgres-protocol" version = "0.6.10" @@ -3258,7 +4134,7 @@ dependencies = [ "base64 0.22.1", "byteorder", "bytes", - "fallible-iterator", + "fallible-iterator 0.2.0", "hmac", "md-5", "memchr", @@ -3274,7 +4150,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54b858f82211e84682fecd373f68e1ceae642d8d751a1ebd13f33de6257b3e20" dependencies = [ "bytes", - "fallible-iterator", + "fallible-iterator 0.2.0", "postgres-protocol", ] @@ -3308,6 +4184,17 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c" +[[package]] +name = "prefix-trie" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cf6e3177f0684016a5c209b00882e15f8bdd3f3bb48f0491df10cd102d0c6e7" +dependencies = [ + "either", + "ipnet", + "num-traits", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -3446,7 +4333,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2", + "socket2 0.6.3", "thiserror 2.0.18", "tokio", "tracing", @@ -3483,7 +4370,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.3", "tracing", "windows-sys 0.60.2", ] @@ -3550,6 +4437,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.2.2" @@ -3607,6 +4505,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_hc" version = "0.2.0" @@ -3631,6 +4535,30 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" +[[package]] +name = "redis" +version = "0.25.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e46922bd01fefcfdcf58d9cd626da082bb2cde27211920dacfde6b2ecf9a35b" +dependencies = [ + "arc-swap", + "async-trait", + "bytes", + "combine", + "futures", + "futures-util", + "itoa", + "percent-encoding", + "pin-project-lite", + "ryu", + "sha1_smol", + "socket2 0.5.10", + "tokio", + "tokio-retry", + "tokio-util", + "url", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3726,6 +4654,7 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", + "futures-channel", "futures-core", "futures-util", "http", @@ -3793,6 +4722,12 @@ dependencies = [ "web-sys", ] +[[package]] +name = "resolv-conf" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e061d1b48cb8d38042de4ae0a7a6401009d6143dc80d2e2d6f31f0bdd6470c7" + [[package]] name = "resvg" version = "0.43.0" @@ -3943,12 +4878,49 @@ dependencies = [ "semver", ] +[[package]] +name = "rustc_version_runtime" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dd18cd2bae1820af0b6ad5e54f4a51d0f3fcc53b05f845675074efcc7af071d" +dependencies = [ + "rustc_version", + "semver", +] + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags 2.11.0", + "errno", + "libc", + "linux-raw-sys 0.4.15", + "windows-sys 0.59.0", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags 2.11.0", + "errno", + "libc", + "linux-raw-sys 0.12.1", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.38" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69f9466fb2c14ea04357e91413efb882e2a6d4a406e625449bc0a5d360d53a21" dependencies = [ + "log", "once_cell", "ring", "rustls-pki-types", @@ -4158,6 +5130,16 @@ dependencies = [ "typeid", ] +[[package]] +name = "serde_bytes" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" +dependencies = [ + "serde", + "serde_core", +] + [[package]] name = "serde_core" version = "1.0.228" @@ -4195,6 +5177,7 @@ version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ + "indexmap 2.13.0", "itoa", "memchr", "serde", @@ -4322,10 +5305,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -4333,7 +5322,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -4369,6 +5358,16 @@ version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" +[[package]] +name = "simd_cesu8" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f90157bb87cddf702797c5dadfa0be7d266cdf49e22da2fcaa32eff75b2c33" +dependencies = [ + "rustc_version", + "simdutf8", +] + [[package]] name = "simdutf8" version = "0.1.5" @@ -4420,6 +5419,16 @@ dependencies = [ "serde", ] +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.3" @@ -4772,6 +5781,27 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "subtle" version = "2.6.1" @@ -4841,6 +5871,27 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags 2.11.0", + "core-foundation 0.9.4", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "system-deps" version = "6.2.2" @@ -4854,6 +5905,18 @@ dependencies = [ "version-compare", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + +[[package]] +name = "take_mut" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f764005d11ee5f36500a149ace24e00e3da98b0158b3e2d53a7495660d3f4d60" + [[package]] name = "tao" version = "0.35.0" @@ -4862,7 +5925,7 @@ checksum = "1cf65722394c2ac443e80120064987f8914ee1d4e4e36e63cdf10f2990f01159" dependencies = [ "bitflags 2.11.0", "block2", - "core-foundation", + "core-foundation 0.10.1", "core-graphics", "crossbeam-channel", "dbus", @@ -4872,7 +5935,7 @@ dependencies = [ "gdkwayland-sys", "gdkx11-sys", "gtk", - "jni", + "jni 0.21.1", "libc", "log", "ndk", @@ -4911,6 +5974,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" +[[package]] +name = "tar" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840" +dependencies = [ + "filetime", + "libc", + "xattr", +] + [[package]] name = "target-lexicon" version = "0.12.16" @@ -4934,7 +6008,7 @@ dependencies = [ "gtk", "heck 0.5.0", "http", - "jni", + "jni 0.21.1", "libc", "log", "mime", @@ -5137,7 +6211,7 @@ dependencies = [ "dpi", "gtk", "http", - "jni", + "jni 0.21.1", "objc2", "objc2-ui-kit", "objc2-web-kit", @@ -5160,7 +6234,7 @@ checksum = "2cadb13dad0c681e1e0a2c49ae488f0e2906ded3d57e7a0017f4aaf46e387117" dependencies = [ "gtk", "http", - "jni", + "jni 0.21.1", "log", "objc2", "objc2-app-kit", @@ -5323,6 +6397,15 @@ dependencies = [ "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tiny-skia" version = "0.11.4" @@ -5407,7 +6490,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.3", "tokio-macros", "windows-sys 0.61.2", ] @@ -5432,7 +6515,7 @@ dependencies = [ "async-trait", "byteorder", "bytes", - "fallible-iterator", + "fallible-iterator 0.2.0", "futures-channel", "futures-util", "log", @@ -5443,7 +6526,7 @@ dependencies = [ "postgres-protocol", "postgres-types", "rand 0.9.2", - "socket2", + "socket2 0.6.3", "tokio", "tokio-util", "whoami 2.1.1", @@ -5464,6 +6547,17 @@ dependencies = [ "x509-cert", ] +[[package]] +name = "tokio-retry" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a129d95275ebf4c493ec53bf0f8cd95f5ac161bc4f381700809a54f595d4470" +dependencies = [ + "pin-project-lite", + "rand 0.10.1", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -5493,7 +6587,9 @@ checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", + "futures-io", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -5723,6 +6819,26 @@ dependencies = [ "core_maths", ] +[[package]] +name = "typed-builder" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "398a3a3c918c96de527dc11e6e846cd549d4508030b8a33e1da12789c856b81a" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e48cea23f68d1f78eb7bc092881b6bb88d3d6b5b7e6234f6f9c911da1ffb221" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "typeid" version = "1.0.3" @@ -5833,6 +6949,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1d386ff53b415b7fe27b50bb44679e2cc4660272694b7b6f3326d8480823a94" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode-xid" version = "0.2.6" @@ -5950,19 +7072,24 @@ name = "veloxdb" version = "0.1.0-8" dependencies = [ "base64 0.22.1", + "bson", "chrono", "csv", "deadpool-postgres", + "duckdb", "futures-util", "hex", "keyring", "log", + "mongodb", "printpdf", "rand 0.8.5", + "redis", "reqwest 0.12.28", "resvg", "rustls", "rustls-pemfile", + "semver", "serde", "serde_json", "sqlx", @@ -5971,6 +7098,7 @@ dependencies = [ "tauri-plugin-dialog", "tauri-plugin-log", "tauri-plugin-store", + "thiserror 2.0.18", "tokio", "tokio-postgres", "tokio-postgres-rustls", @@ -6363,6 +7491,12 @@ dependencies = [ "web-sys", ] +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + [[package]] name = "winapi" version = "0.3.9" @@ -6512,6 +7646,17 @@ dependencies = [ "windows-link 0.1.3", ] +[[package]] +name = "windows-registry" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +dependencies = [ + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + [[package]] name = "windows-result" version = "0.3.4" @@ -7003,7 +8148,7 @@ dependencies = [ "gtk", "http", "javascriptcore-rs", - "jni", + "jni 0.21.1", "libc", "ndk", "objc2", @@ -7071,6 +8216,16 @@ dependencies = [ "tls_codec", ] +[[package]] +name = "xattr" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" +dependencies = [ + "libc", + "rustix 1.1.4", +] + [[package]] name = "xmlwriter" version = "0.1.0" @@ -7194,12 +8349,44 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "zip" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb2a05c7c36fde6c09b08576c9f7fb4cda705990f73b58fe011abf7dfb24168b" +dependencies = [ + "arbitrary", + "crc32fast", + "flate2", + "indexmap 2.13.0", + "memchr", + "zopfli", +] + +[[package]] +name = "zlib-rs" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "977347db8caa080403f6b6b7c1cda9479a8e869316f7e13a59b19076a40f94e3" + [[package]] name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + [[package]] name = "zune-core" version = "0.4.12" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 5e16056..3822f28 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -45,4 +45,10 @@ csv = "1.3" printpdf = "0.7" reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] } futures-util = "0.3" +thiserror = "2" chrono = "0.4.45" +mongodb = "3" +bson = "2" +duckdb = { version = "1", features = ["bundled"] } +redis = { version = "0.25", features = ["tokio-comp", "connection-manager"] } +semver = "1" diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs deleted file mode 100644 index 6bd109f..0000000 --- a/src-tauri/src/commands.rs +++ /dev/null @@ -1,4082 +0,0 @@ -use std::collections::{BTreeMap, HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; -use std::time::Instant; - -use futures_util::StreamExt; -use serde_json::Value; -use tauri::{AppHandle, Emitter, State}; -use sqlx::{Column, Decode, Row, Type}; -use sqlx::mysql::{MySql, MySqlRow}; -use sqlx::sqlite::{Sqlite, SqliteRow}; -use tokio_postgres::SimpleQueryMessage; -use uuid::Uuid; - -use crate::db::{ - build_mysql_pool, build_mysql_pool_custom, build_pool, build_pool_custom, build_sqlite_pool, - disconnect_connection, drop_pool, get_or_create_mysql_pool, get_or_create_sqlite_pool, - list_connections, load_connection, persist_connection_with_password, quote_identifier, - refresh_connection_pools, require_safe_identifier, resolve_connection_engine, - with_pool_client_retry, AppState, - DEFAULT_MYSQL_PORT, MAX_QUERY_ROWS, -}; -use crate::credentials; -use crate::pg_error::{error_line_column, map_pg_err}; -use crate::sql_split::split_sql_statements; -use crate::models::{ - AskVeloxyChatRequest, AskVeloxyChatResponse, AskVeloxyConversationMessage, - AskVeloxyConversationResponse, AskVeloxyDbContextCache, AskVeloxyRequest, AskVeloxyResponse, - AskVeloxyTableRef, AskVeloxyTokenStats, ColumnInfo, ColumnProperties, ConnectionInput, - ConnectionSummary, DatabaseInfo, DatabaseEngine, DdlBatchRequest, DdlStatementRequest, - ForeignKeyEdge, IndexInfo, LintSqlRequest, LintSqlResult, QueryEditorColumn, - QueryEditorFunction, QueryEditorMetadata, QueryEditorTable, QueryRequest, QueryResult, - SchemaRequest, SqlDiagnostic, StoredConnection, SwitchDatabaseRequest, TableIndexesResult, - TableInfo, TablePropertiesApplyRequest, VeloxyStreamChunk, -}; -use crate::export::{ - DiagramExportRequest, ExportQueryRequest, - export_diagram_to_png, export_results_csv, export_results_json, -}; -use crate::ssh_tunnel::SshTunnel; - -/// Cap FK rows returned to the UI to keep IPC payloads bounded. -const MAX_FOREIGN_KEY_ROWS: i64 = 5000; - -/// Cap index rows per table (fetch limit + 1 to detect truncation). -const MAX_TABLE_INDEX_ROWS: i64 = 500; -const MAX_EDITOR_TABLES: i64 = 150; -const MAX_EDITOR_COLUMNS_PER_TABLE: i64 = 60; -const MAX_EDITOR_FUNCTIONS: i64 = 200; -const MAX_LINT_SQL_BYTES: usize = 65_536; -const ASK_VELOXY_MAX_CONTEXT_TABLES: usize = 8; -const ASK_VELOXY_MAX_CONTEXT_COLUMNS: usize = 18; -const ASK_VELOXY_MAX_CONTEXT_RELATIONSHIPS: usize = 36; -const ASK_VELOXY_SCHEMA_CHAR_BUDGET: usize = 6_000; -const ASK_VELOXY_PROMPT_CHAR_BUDGET: usize = 12_000; -const ASK_VELOXY_MAX_HISTORY_MESSAGES: usize = 30; -const ASK_VELOXY_MAX_CHAT_TOKENS: u32 = 10_000; - -fn mysql_decode_error(context: &str, column_name: &str, index: Option, detail: &str) -> String { - match index { - Some(idx) => format!( - "MySQL decode error in {} at column '{}' (index {}): {}", - context, column_name, idx, detail - ), - None => format!( - "MySQL decode error in {} at column '{}': {}", - context, column_name, detail - ), - } -} - -fn sqlite_decode_error(context: &str, column_name: &str, index: Option, detail: &str) -> String { - match index { - Some(idx) => format!( - "SQLite decode error in {} at column '{}' (index {}): {}", - context, column_name, idx, detail - ), - None => format!( - "SQLite decode error in {} at column '{}': {}", - context, column_name, detail - ), - } -} - -fn mysql_get_idx(row: &MySqlRow, index: usize, column_name: &str, context: &str) -> Result -where - for<'r> T: Decode<'r, MySql> + Type, -{ - row.try_get::(index) - .map_err(|error| mysql_decode_error(context, column_name, Some(index), &error.to_string())) -} - -fn sqlite_get_idx(row: &SqliteRow, index: usize, column_name: &str, context: &str) -> Result -where - for<'r> T: Decode<'r, Sqlite> + Type, -{ - row.try_get::(index) - .map_err(|error| sqlite_decode_error(context, column_name, Some(index), &error.to_string())) -} - -fn sqlite_get_name(row: &SqliteRow, column_name: &str, context: &str) -> Result -where - for<'r> T: Decode<'r, Sqlite> + Type, -{ - row.try_get::(column_name).map_err(|error| { - format!( - "SQLite decode error in {} at column '{}': {}", - context, column_name, error - ) - }) -} - -fn database_name_from_mysql_value( - value: Option, - context: &str, -) -> Result { - let name = value - .filter(|value| !value.is_empty()) - .ok_or_else(|| format!("{context} returned an empty database name"))?; - Ok(name) -} - -fn mysql_database_name_from_row(row: &MySqlRow, context: &str) -> Result { - let value = mysql_value_to_string(row, 0, "Database", context)?; - database_name_from_mysql_value(value, context) -} - -fn mysql_value_to_string(row: &MySqlRow, index: usize, column_name: &str, context: &str) -> Result, String> { - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::>, _>(index) { - return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::>, _>(index) { - return Ok(value.map(|v| decode_mysql_bytes_as_string(&v))); - } - Err(mysql_decode_error( - context, - column_name, - Some(index), - "unsupported value type", - )) -} - -fn decode_mysql_bytes_as_string(bytes: &[u8]) -> String { - String::from_utf8_lossy(bytes).into_owned() -} - -/// Like [`mysql_value_to_string`] but encodes raw bytes as hex (for ad-hoc query grids). -fn mysql_value_to_display_string( - row: &MySqlRow, - index: usize, - column_name: &str, - context: &str, -) -> Result, String> { - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::>, _>(index) { - return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::>, _>(index) { - return Ok(value.map(|v| format!("0x{}", hex::encode(v)))); - } - Err(mysql_decode_error( - context, - column_name, - Some(index), - "unsupported value type", - )) -} - -fn mysql_get_string(row: &MySqlRow, index: usize, column_name: &str, context: &str) -> Result { - let value = mysql_value_to_string(row, index, column_name, context)?; - value.ok_or_else(|| { - mysql_decode_error(context, column_name, Some(index), "unexpected null value") - }) -} - -fn mysql_get_optional_string( - row: &MySqlRow, - index: usize, - column_name: &str, - context: &str, -) -> Result, String> { - mysql_value_to_string(row, index, column_name, context) -} - -fn sqlite_value_to_string(row: &SqliteRow, index: usize, column_name: &str, context: &str) -> Result, String> { - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::, _>(index) { - return Ok(value.map(|v| v.to_string())); - } - if let Ok(value) = row.try_get::>, _>(index) { - return Ok(value.map(|v| format!("0x{}", hex::encode(v)))); - } - Err(sqlite_decode_error( - context, - column_name, - Some(index), - "unsupported value type", - )) -} - -fn is_row_returning_sql(sql: &str) -> bool { - let trimmed = sql.trim_start(); - let upper = trimmed.to_uppercase(); - upper.starts_with("SELECT") - || upper.starts_with("WITH") - || upper.starts_with("SHOW") - || upper.starts_with("EXPLAIN") - || upper.starts_with("DESCRIBE") - || upper.starts_with("DESC") - || upper.starts_with("PRAGMA") - || upper.starts_with("VALUES") - || upper.starts_with("TABLE ") -} - -fn map_mysql_rows( - rows: Vec, - max_query_rows: usize, -) -> Result<(Vec, Vec>>, usize, bool), String> { - let mut columns: Vec = Vec::new(); - if let Some(first) = rows.first() { - columns = first - .columns() - .iter() - .map(|column| column.name().to_string()) - .collect(); - } - let total_rows = rows.len(); - let mut mapped_rows = Vec::new(); - for row in rows.into_iter().take(max_query_rows) { - let mut mapped_row = BTreeMap::new(); - for (index, column_name) in columns.iter().enumerate() { - let value = mysql_value_to_display_string(&row, index, column_name, "run_query")?; - mapped_row.insert(column_name.clone(), value); - } - mapped_rows.push(mapped_row); - } - Ok((columns, mapped_rows, total_rows, total_rows > max_query_rows)) -} - -fn map_sqlite_rows( - rows: Vec, - max_query_rows: usize, -) -> Result<(Vec, Vec>>, usize, bool), String> { - let mut columns: Vec = Vec::new(); - if let Some(first) = rows.first() { - columns = first - .columns() - .iter() - .map(|column| column.name().to_string()) - .collect(); - } - let total_rows = rows.len(); - let mut mapped_rows = Vec::new(); - for row in rows.into_iter().take(max_query_rows) { - let mut mapped_row = BTreeMap::new(); - for (index, column_name) in columns.iter().enumerate() { - let value = sqlite_value_to_string(&row, index, column_name, "run_query")?; - mapped_row.insert(column_name.clone(), value); - } - mapped_rows.push(mapped_row); - } - Ok((columns, mapped_rows, total_rows, total_rows > max_query_rows)) -} - -async fn run_query_mysql_or_sqlite( - app: &AppHandle, - state: &AppState, - connection_id: &str, - sql: &str, - max_query_rows: usize, - engine: DatabaseEngine, -) -> Result { - let started_at = Instant::now(); - let statements = split_sql_statements(sql); - if statements.is_empty() { - return Err("Enter a SQL statement before running the query.".to_string()); - } - - let mut columns = Vec::new(); - let mut rows = Vec::new(); - let mut total_rows = 0usize; - let mut truncated = false; - let mut command_tag: Option = None; - - match engine { - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(app, state, connection_id).await?; - let mut conn = pool - .acquire() - .await - .map_err(|error| error.to_string())?; - for statement in statements { - if is_row_returning_sql(&statement) { - let fetched = sqlx::query(&statement) - .fetch_all(&mut *conn) - .await - .map_err(|error| error.to_string())?; - let mapped = map_mysql_rows(fetched, max_query_rows)?; - columns = mapped.0; - rows = mapped.1; - total_rows = mapped.2; - truncated = mapped.3; - command_tag = None; - } else { - let result = sqlx::query(&statement) - .execute(&mut *conn) - .await - .map_err(|error| error.to_string())?; - let affected = result.rows_affected(); - command_tag = Some(affected); - if rows.is_empty() { - total_rows = affected as usize; - } - } - } - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; - let mut conn = pool - .acquire() - .await - .map_err(|error| error.to_string())?; - for statement in statements { - if is_row_returning_sql(&statement) { - let fetched = sqlx::query(&statement) - .fetch_all(&mut *conn) - .await - .map_err(|error| error.to_string())?; - let mapped = map_sqlite_rows(fetched, max_query_rows)?; - columns = mapped.0; - rows = mapped.1; - total_rows = mapped.2; - truncated = mapped.3; - command_tag = None; - } else { - let result = sqlx::query(&statement) - .execute(&mut *conn) - .await - .map_err(|error| error.to_string())?; - let affected = result.rows_affected(); - command_tag = Some(affected); - if rows.is_empty() { - total_rows = affected as usize; - } - } - } - } - DatabaseEngine::Postgres => { - return Err("Internal engine routing error.".to_string()); - } - } - - Ok(QueryResult { - columns, - row_count: if rows.is_empty() { - total_rows - } else { - rows.len() - }, - rows, - execution_ms: started_at.elapsed().as_millis(), - truncated, - command_tag, - }) -} - -#[tauri::command] -pub async fn connect_db( - app: AppHandle, - state: State<'_, AppState>, - input: ConnectionInput, -) -> Result { - let connection_id = input - .id - .clone() - .unwrap_or_else(|| Uuid::new_v4().to_string()); - - match input.engine { - DatabaseEngine::Postgres => { - let pool = if let Some(ref ssh_config) = input.ssh_config { - if ssh_config.is_active() { - let tunnel = match SshTunnel::connect(ssh_config, &input.host, input.port).await { - Ok(tunnel) => tunnel, - Err(e) => return Err(format!("SSH tunnel failed: {}", e)), - }; - let local_port = tunnel.local_port; - state - .ssh_tunnels - .write() - .await - .insert(connection_id.clone(), tunnel); - build_pool_custom("127.0.0.1", local_port, &input)? - } else { - build_pool(&input)? - } - } else { - build_pool(&input)? - }; - - let client = match pool.get().await { - Ok(client) => client, - Err(e) => { - drop_pool(&state, &connection_id).await; - return Err(e.to_string()); - } - }; - - if let Err(e) = client.simple_query("select 1").await { - drop_pool(&state, &connection_id).await; - return Err(map_pg_err(e, None)); - } - - state - .pools - .write() - .await - .insert(connection_id.clone(), pool); - } - DatabaseEngine::Mysql => { - let pool = if let Some(ref ssh_config) = input.ssh_config { - if ssh_config.is_active() { - let remote_port = if input.port == 0 { DEFAULT_MYSQL_PORT } else { input.port }; - let tunnel = match SshTunnel::connect(ssh_config, &input.host, remote_port).await { - Ok(tunnel) => tunnel, - Err(e) => return Err(format!("SSH tunnel failed: {}", e)), - }; - let local_port = tunnel.local_port; - state - .ssh_tunnels - .write() - .await - .insert(connection_id.clone(), tunnel); - build_mysql_pool_custom("127.0.0.1", local_port, &input).await? - } else { - build_mysql_pool(&input).await? - } - } else { - build_mysql_pool(&input).await? - }; - - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|e| e.to_string())?; - - state - .mysql_pools - .write() - .await - .insert(connection_id.clone(), pool); - } - DatabaseEngine::Sqlite => { - let pool = build_sqlite_pool(&input).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|e| e.to_string())?; - state - .sqlite_pools - .write() - .await - .insert(connection_id.clone(), pool); - } - } - - let stored_connection = StoredConnection::from_input(connection_id.clone(), input.clone()); - persist_connection_with_password(&app, &stored_connection, &input.password)?; - - *state.active_connection_id.write().await = Some(connection_id); - - Ok(stored_connection.summary()) -} - -#[tauri::command] -pub async fn list_connections_command(app: AppHandle) -> Result, String> { - list_connections(&app) -} - -#[tauri::command] -pub async fn set_active_connection( - app: AppHandle, - state: State<'_, AppState>, - connection_id: String, -) -> Result { - let stored_connection = load_connection(&app, &connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - - match stored_connection.engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { - client - .simple_query("select 1") - .await - .map_err(|error| map_pg_err(error, None))?; - Ok(()) - }) - .await?; - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - } - } - - *state.active_connection_id.write().await = Some(connection_id); - - Ok(stored_connection.summary()) -} - -#[tauri::command] -pub async fn ping_connection( - app: AppHandle, - state: State<'_, AppState>, - connection_id: String, -) -> Result<(), String> { - let stored_connection = load_connection(&app, &connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - match stored_connection.engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { - client - .simple_query("select 1") - .await - .map_err(|error| error.to_string())?; - Ok(()) - }) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - Ok(()) - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - Ok(()) - } - } -} - -#[tauri::command] -pub async fn refresh_connection( - app: AppHandle, - state: State<'_, AppState>, - connection_id: String, -) -> Result<(), String> { - refresh_connection_pools(&app, &state, &connection_id).await -} - -#[tauri::command] -pub async fn disconnect_db( - state: State<'_, AppState>, - connection_id: String, -) -> Result<(), String> { - disconnect_connection(&state, &connection_id).await; - Ok(()) -} - -/// Renames a saved connection without affecting the active pool or SSH tunnel. -#[tauri::command] -pub async fn rename_connection( - app: AppHandle, - connection_id: String, - new_name: String, -) -> Result { - crate::db::rename_connection_in_store(&app, &connection_id, &new_name) -} - -#[tauri::command] -pub async fn delete_connection( - app: AppHandle, - state: State<'_, AppState>, - connection_id: String, -) -> Result<(), String> { - disconnect_connection(&state, &connection_id).await; - if let Err(e) = credentials::delete_password(&connection_id) { - log::warn!("Failed to delete keychain entry for {}: {}", connection_id, e); - } - crate::db::delete_connection_from_store(&app, &connection_id)?; - Ok(()) -} - -#[tauri::command] -pub async fn list_databases( - app: AppHandle, - state: State<'_, AppState>, - connection_id: Option, -) -> Result, String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { - let rows = client - .query( - "select datname from pg_database where datistemplate = false and has_database_privilege(datname, 'CONNECT') order by datname", - &[], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - Ok(rows - .into_iter() - .map(|row| { - let name: String = row.get(0); - DatabaseInfo { name } - }) - .collect()) - }) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let rows = sqlx::query("show databases") - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut databases = Vec::with_capacity(rows.len()); - for row in rows { - let name = mysql_database_name_from_row(&row, "list_databases")?; - databases.push(DatabaseInfo { name }); - } - Ok(databases) - } - DatabaseEngine::Sqlite => Ok(vec![DatabaseInfo { - name: "main".to_string(), - }]), - } -} - -#[tauri::command] -pub async fn switch_database( - app: AppHandle, - state: State<'_, AppState>, - input: SwitchDatabaseRequest, -) -> Result { - let mut stored_connection = load_connection(&app, &input.connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - - if stored_connection.engine == DatabaseEngine::Sqlite { - return Err("Switch database is not supported for SQLite connections.".to_string()); - } - - drop_pool(&state, &input.connection_id).await; - - stored_connection.database = input.database.clone(); - stored_connection.connected_at = crate::models::timestamp_string(); - persist_connection_with_password(&app, &stored_connection, &stored_connection.password.clone().unwrap_or_default())?; - - let connection_input = stored_connection.to_input(); - - match connection_input.engine { - DatabaseEngine::Postgres => { - let pool = if let Some(ref ssh_config) = connection_input.ssh_config { - if ssh_config.is_active() { - let tunnel = match SshTunnel::connect(ssh_config, &connection_input.host, connection_input.port).await { - Ok(tunnel) => tunnel, - Err(e) => return Err(format!("SSH tunnel failed: {}", e)), - }; - let local_port = tunnel.local_port; - state - .ssh_tunnels - .write() - .await - .insert(input.connection_id.clone(), tunnel); - build_pool_custom("127.0.0.1", local_port, &connection_input)? - } else { - build_pool(&connection_input)? - } - } else { - build_pool(&connection_input)? - }; - - let client = match pool.get().await { - Ok(client) => client, - Err(e) => { - drop_pool(&state, &input.connection_id).await; - return Err(e.to_string()); - } - }; - - if let Err(e) = client.simple_query("select 1").await { - drop_pool(&state, &input.connection_id).await; - return Err(map_pg_err(e, None)); - } - - state - .pools - .write() - .await - .insert(input.connection_id.clone(), pool); - } - DatabaseEngine::Mysql => { - let pool = if let Some(ref ssh_config) = connection_input.ssh_config { - if ssh_config.is_active() { - let remote_port = if connection_input.port == 0 { - DEFAULT_MYSQL_PORT - } else { - connection_input.port - }; - let tunnel = match SshTunnel::connect(ssh_config, &connection_input.host, remote_port).await { - Ok(tunnel) => tunnel, - Err(e) => return Err(format!("SSH tunnel failed: {}", e)), - }; - let local_port = tunnel.local_port; - state - .ssh_tunnels - .write() - .await - .insert(input.connection_id.clone(), tunnel); - build_mysql_pool_custom("127.0.0.1", local_port, &connection_input).await? - } else { - build_mysql_pool(&connection_input).await? - } - } else { - build_mysql_pool(&connection_input).await? - }; - - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|e| e.to_string())?; - - state - .mysql_pools - .write() - .await - .insert(input.connection_id.clone(), pool); - } - DatabaseEngine::Sqlite => {} - } - - *state.active_connection_id.write().await = Some(input.connection_id); - - Ok(stored_connection.summary()) -} - -#[tauri::command] -pub async fn run_query( - app: AppHandle, - state: State<'_, AppState>, - input: QueryRequest, -) -> Result { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id).await?; - let sql = input.sql.trim().to_string(); - - if sql.is_empty() { - return Err("Enter a SQL statement before running the query.".to_string()); - } - - if !input.allow_write.unwrap_or(false) && !is_read_only_sql(&sql) { - return Err( - "This statement modifies data or schema. Confirm execution in the editor, \ - or use the model/DDL workflow for schema changes." - .to_string(), - ); - } - - let max_query_rows = input.max_rows.unwrap_or(MAX_QUERY_ROWS); - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, sql, |client, sql| async move { - let started_at = Instant::now(); - let messages = client - .simple_query(&sql) - .await - .map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - - let mut columns = Vec::new(); - let mut rows = Vec::new(); - let mut total_rows = 0usize; - let mut command_tag = None; - - for message in messages { - match message { - SimpleQueryMessage::RowDescription(description) => { - if columns.is_empty() { - columns = description - .iter() - .map(|column| column.name().to_string()) - .collect(); - } - } - SimpleQueryMessage::Row(row) => { - total_rows += 1; - - if columns.is_empty() { - columns = row - .columns() - .iter() - .map(|column| column.name().to_string()) - .collect(); - } - - if rows.len() >= max_query_rows { - continue; - } - - let mut mapped_row = BTreeMap::new(); - for (index, column_name) in columns.iter().enumerate() { - mapped_row.insert(column_name.clone(), row.get(index).map(str::to_owned)); - } - rows.push(mapped_row); - } - SimpleQueryMessage::CommandComplete(count) => { - command_tag = Some(count); - } - _ => {} - } - } - - Ok(QueryResult { - columns, - row_count: rows.len(), - rows, - execution_ms: started_at.elapsed().as_millis(), - truncated: total_rows > max_query_rows, - command_tag, - }) - }) - .await - } - DatabaseEngine::Mysql | DatabaseEngine::Sqlite => { - run_query_mysql_or_sqlite(&app, &state, &connection_id, &sql, max_query_rows, engine).await - } - } -} - -async fn fetch_query_editor_metadata_for_connection( - app: &AppHandle, - state: &AppState, - connection_id: &str, - engine: DatabaseEngine, -) -> Result { - if engine == DatabaseEngine::Mysql { - let pool = get_or_create_mysql_pool(app, state, connection_id).await?; - let database = load_connection(app, connection_id)? - .map(|connection| connection.database) - .unwrap_or_default(); - let table_rows = sqlx::query( - " - select table_schema, table_name - from information_schema.tables - where table_type = 'BASE TABLE' - and table_schema = ? - and table_schema not in ('information_schema', 'mysql', 'performance_schema', 'sys') - order by table_schema, table_name - limit ? - ", - ) - .bind(&database) - .bind(MAX_EDITOR_TABLES + 1) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - - let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; - let mut tables = Vec::new(); - let mut truncated_columns = false; - - for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { - let schema: String = mysql_get_string(&row, 0, "table_schema", "get_query_editor_metadata")?; - let name: String = mysql_get_string(&row, 1, "table_name", "get_query_editor_metadata")?; - let column_rows = sqlx::query( - " - select column_name, data_type - from information_schema.columns - where table_schema = ? and table_name = ? - order by ordinal_position - limit ? - ", - ) - .bind(&schema) - .bind(&name) - .bind(MAX_EDITOR_COLUMNS_PER_TABLE + 1) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { - truncated_columns = true; - } - let mut columns = Vec::new(); - for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { - columns.push(QueryEditorColumn { - name: mysql_get_string(&column, 0, "column_name", "get_query_editor_metadata")?, - data_type: mysql_get_string(&column, 1, "data_type", "get_query_editor_metadata")?, - }); - } - tables.push(QueryEditorTable { schema, name, columns }); - } - - return Ok(QueryEditorMetadata { - tables, - functions: Vec::new(), - truncated_tables, - truncated_columns, - truncated_functions: false, - }); - } - - if engine == DatabaseEngine::Sqlite { - let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; - let table_rows = sqlx::query( - " - select name - from sqlite_master - where type = 'table' - and name not like 'sqlite_%' - order by name - limit ? - ", - ) - .bind(MAX_EDITOR_TABLES + 1) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; - let mut tables = Vec::new(); - let mut truncated_columns = false; - for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { - let name: String = sqlite_get_idx(&row, 0, "name", "get_query_editor_metadata")?; - require_safe_identifier(&name, "table name")?; - let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(&name)); - let column_rows = sqlx::query(&pragma_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { - truncated_columns = true; - } - let mut columns = Vec::new(); - for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { - columns.push(QueryEditorColumn { - name: sqlite_get_name(&column, "name", "get_query_editor_metadata")?, - data_type: sqlite_get_name(&column, "type", "get_query_editor_metadata")?, - }); - } - tables.push(QueryEditorTable { - schema: "main".to_string(), - name, - columns, - }); - } - return Ok(QueryEditorMetadata { - tables, - functions: Vec::new(), - truncated_tables, - truncated_columns, - truncated_functions: false, - }); - } - - with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { - let table_rows = client - .query( - " - select n.nspname::text as schema_name, c.relname::text as table_name - from pg_class c - join pg_namespace n on n.oid = c.relnamespace - where c.relkind in ('r', 'p', 'v', 'm', 'f') - and n.nspname not in ('pg_catalog', 'information_schema') - order by n.nspname, c.relname - limit $1 - ", - &[&(MAX_EDITOR_TABLES + 1)], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; - let table_rows = if truncated_tables { - table_rows - .into_iter() - .take(MAX_EDITOR_TABLES as usize) - .collect::>() - } else { - table_rows - }; - - let mut tables = Vec::with_capacity(table_rows.len()); - let mut truncated_columns = false; - - for row in table_rows { - let schema: String = row.get(0); - let name: String = row.get(1); - let column_rows = client - .query( - " - select a.attname::text as column_name, - format_type(a.atttypid, a.atttypmod)::text as data_type - from pg_attribute a - join pg_class c on c.oid = a.attrelid - join pg_namespace n on n.oid = c.relnamespace - where n.nspname = $1 - and c.relname = $2 - and a.attnum > 0 - and not a.attisdropped - order by a.attnum - limit $3 - ", - &[&schema, &name, &(MAX_EDITOR_COLUMNS_PER_TABLE + 1)], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let columns_exceeded = column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE; - if columns_exceeded { - truncated_columns = true; - } - let columns = column_rows - .into_iter() - .take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) - .map(|column| QueryEditorColumn { - name: column.get(0), - data_type: column.get(1), - }) - .collect(); - - tables.push(QueryEditorTable { - schema, - name, - columns, - }); - } - - let function_rows = client - .query( - " - select - n.nspname::text as schema_name, - p.proname::text as function_name, - coalesce(pg_get_function_identity_arguments(p.oid), '')::text as args, - pg_get_function_result(p.oid)::text as return_type - from pg_proc p - join pg_namespace n on n.oid = p.pronamespace - where n.nspname not in ('pg_catalog', 'information_schema') - order by n.nspname, p.proname - limit $1 - ", - &[&(MAX_EDITOR_FUNCTIONS + 1)], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let truncated_functions = function_rows.len() as i64 > MAX_EDITOR_FUNCTIONS; - let functions = function_rows - .into_iter() - .take(MAX_EDITOR_FUNCTIONS as usize) - .map(|row| { - let args_raw: String = row.get(2); - QueryEditorFunction { - schema: row.get(0), - name: row.get(1), - arg_types: if args_raw.trim().is_empty() { - Vec::new() - } else { - args_raw - .split(',') - .map(|value| value.trim().to_string()) - .collect() - }, - return_type: row.get(3), - } - }) - .collect(); - - Ok(QueryEditorMetadata { - tables, - functions, - truncated_tables, - truncated_columns, - truncated_functions, - }) - }) - .await -} - -async fn fetch_foreign_keys_for_connection( - app: &AppHandle, - state: &AppState, - connection_id: &str, - engine: DatabaseEngine, -) -> Result, String> { - if engine == DatabaseEngine::Mysql { - let pool = get_or_create_mysql_pool(app, state, connection_id).await?; - let rows = sqlx::query( - " - select - kcu.table_schema as from_schema, - kcu.table_name as from_table, - kcu.column_name as from_column, - kcu.referenced_table_schema as to_schema, - kcu.referenced_table_name as to_table, - kcu.referenced_column_name as to_column - from information_schema.key_column_usage kcu - where kcu.referenced_table_name is not null - order by kcu.table_schema, kcu.table_name, kcu.ordinal_position - limit ? - ", - ) - .bind(MAX_FOREIGN_KEY_ROWS) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut edges = Vec::new(); - for row in rows { - edges.push(ForeignKeyEdge { - from_schema: mysql_get_string(&row, 0, "from_schema", "get_foreign_keys")?, - from_table: mysql_get_string(&row, 1, "from_table", "get_foreign_keys")?, - from_column: mysql_get_string(&row, 2, "from_column", "get_foreign_keys")?, - to_schema: mysql_get_string(&row, 3, "to_schema", "get_foreign_keys")?, - to_table: mysql_get_string(&row, 4, "to_table", "get_foreign_keys")?, - to_column: mysql_get_string(&row, 5, "to_column", "get_foreign_keys")?, - }); - } - return Ok(edges); - } - - if engine == DatabaseEngine::Sqlite { - let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; - let tables = sqlx::query( - " - select name - from sqlite_master - where type = 'table' - and name not like 'sqlite_%' - ", - ) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut edges = Vec::new(); - for table in tables { - let table_name: String = sqlite_get_idx(&table, 0, "name", "get_foreign_keys")?; - require_safe_identifier(&table_name, "table name")?; - let fk_sql = format!("PRAGMA foreign_key_list(\"{}\");", quote_identifier(&table_name)); - let fk_rows = sqlx::query(&fk_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - for row in fk_rows { - edges.push(ForeignKeyEdge { - from_schema: "main".to_string(), - from_table: table_name.clone(), - from_column: sqlite_get_name(&row, "from", "get_foreign_keys")?, - to_schema: "main".to_string(), - to_table: sqlite_get_name(&row, "table", "get_foreign_keys")?, - to_column: sqlite_get_name(&row, "to", "get_foreign_keys")?, - }); - if edges.len() >= MAX_FOREIGN_KEY_ROWS as usize { - return Ok(edges); - } - } - } - return Ok(edges); - } - - with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { - let rows = client - .query( - " - select - src_ns.nspname::text as from_schema, - src_cls.relname::text as from_table, - src_att.attname::text as from_column, - tgt_ns.nspname::text as to_schema, - tgt_cls.relname::text as to_table, - tgt_att.attname::text as to_column - from pg_constraint c - join pg_class src_cls on src_cls.oid = c.conrelid - join pg_namespace src_ns on src_ns.oid = src_cls.relnamespace - join pg_class tgt_cls on tgt_cls.oid = c.confrelid - join pg_namespace tgt_ns on tgt_ns.oid = tgt_cls.relnamespace - cross join lateral unnest(c.conkey, c.confkey) as u(attnum, confattnum) - join pg_attribute src_att - on src_att.attrelid = c.conrelid - and src_att.attnum = u.attnum - and not src_att.attisdropped - join pg_attribute tgt_att - on tgt_att.attrelid = c.confrelid - and tgt_att.attnum = u.confattnum - and not tgt_att.attisdropped - where c.contype = 'f' - and src_ns.nspname not in ('pg_catalog', 'information_schema') - order by src_ns.nspname, src_cls.relname, c.conname, u.attnum - limit $1 - ", - &[&MAX_FOREIGN_KEY_ROWS], - ) - .await - .map_err(|error| error.to_string())?; - - Ok(rows - .into_iter() - .map(|row| ForeignKeyEdge { - from_schema: row.get(0), - from_table: row.get(1), - from_column: row.get(2), - to_schema: row.get(3), - to_table: row.get(4), - to_column: row.get(5), - }) - .collect()) - }) - .await -} - -fn ask_veloxy_context_cache_key(connection_id: &str, database_name: &str) -> String { - format!("{}::{}", connection_id, database_name) -} - -fn ask_veloxy_conversation_key(connection_id: &str, database_name: &str) -> String { - format!("{}::{}", connection_id, database_name) -} - -fn now_epoch_seconds() -> u64 { - use std::time::{SystemTime, UNIX_EPOCH}; - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - -async fn get_or_build_ask_veloxy_db_context( - app: &AppHandle, - state: &AppState, - connection_id: &str, - engine: DatabaseEngine, -) -> Result { - let stored_connection = load_connection(app, connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - let cache_key = ask_veloxy_context_cache_key(connection_id, &stored_connection.database); - if let Some(cached) = state - .ask_veloxy_db_context_cache - .read() - .await - .get(&cache_key) - .cloned() - { - return Ok(cached); - } - - let metadata = fetch_query_editor_metadata_for_connection(app, state, connection_id, engine).await?; - let foreign_keys = fetch_foreign_keys_for_connection(app, state, connection_id, engine).await?; - let cache = AskVeloxyDbContextCache { - database_name: stored_connection.database, - engine, - metadata, - foreign_keys, - }; - state - .ask_veloxy_db_context_cache - .write() - .await - .insert(cache_key, cache.clone()); - Ok(cache) -} - -fn extract_sql_draft_from_text(message: &str) -> Option { - let lowered = message.to_lowercase(); - let markers = ["select ", "with ", "insert ", "update ", "delete ", "explain "]; - let start = markers - .iter() - .filter_map(|marker| lowered.find(marker)) - .min()?; - let mut sql = message[start..].trim().to_string(); - if let Some(idx) = sql.find("```") { - sql.truncate(idx); - } - if sql.ends_with('.') { - sql.pop(); - } - if sql.is_empty() { - None - } else { - Some(sql) - } -} - -fn parse_bool_field(value: &Value, field: &str, default: bool) -> bool { - value.get(field).and_then(Value::as_bool).unwrap_or(default) -} - -#[tauri::command] -pub async fn get_query_editor_metadata( - app: AppHandle, - state: State<'_, AppState>, - connection_id: Option, -) -> Result { - let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; - - if engine == DatabaseEngine::Mysql { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let database = load_connection(&app, &connection_id)? - .map(|connection| connection.database) - .unwrap_or_default(); - let table_rows = sqlx::query( - " - select table_schema, table_name - from information_schema.tables - where table_type = 'BASE TABLE' - and table_schema = ? - and table_schema not in ('information_schema', 'mysql', 'performance_schema', 'sys') - order by table_schema, table_name - limit ? - ", - ) - .bind(&database) - .bind(MAX_EDITOR_TABLES + 1) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - - let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; - let mut tables = Vec::new(); - let mut truncated_columns = false; - - for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { - let schema: String = mysql_get_string(&row, 0, "table_schema", "get_query_editor_metadata")?; - let name: String = mysql_get_string(&row, 1, "table_name", "get_query_editor_metadata")?; - let column_rows = sqlx::query( - " - select column_name, data_type - from information_schema.columns - where table_schema = ? and table_name = ? - order by ordinal_position - limit ? - ", - ) - .bind(&schema) - .bind(&name) - .bind(MAX_EDITOR_COLUMNS_PER_TABLE + 1) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { - truncated_columns = true; - } - let mut columns = Vec::new(); - for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { - columns.push(QueryEditorColumn { - name: mysql_get_string(&column, 0, "column_name", "get_query_editor_metadata")?, - data_type: mysql_get_string(&column, 1, "data_type", "get_query_editor_metadata")?, - }); - } - tables.push(QueryEditorTable { schema, name, columns }); - } - - return Ok(QueryEditorMetadata { - tables, - functions: Vec::new(), - truncated_tables, - truncated_columns, - truncated_functions: false, - }); - } - - if engine == DatabaseEngine::Sqlite { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - let table_rows = sqlx::query( - " - select name - from sqlite_master - where type = 'table' - and name not like 'sqlite_%' - order by name - limit ? - ", - ) - .bind(MAX_EDITOR_TABLES + 1) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; - let mut tables = Vec::new(); - let mut truncated_columns = false; - for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { - let name: String = sqlite_get_idx(&row, 0, "name", "get_query_editor_metadata")?; - require_safe_identifier(&name, "table name")?; - let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(&name)); - let column_rows = sqlx::query(&pragma_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { - truncated_columns = true; - } - let mut columns = Vec::new(); - for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { - columns.push(QueryEditorColumn { - name: sqlite_get_name(&column, "name", "get_query_editor_metadata")?, - data_type: sqlite_get_name(&column, "type", "get_query_editor_metadata")?, - }); - } - tables.push(QueryEditorTable { - schema: "main".to_string(), - name, - columns, - }); - } - return Ok(QueryEditorMetadata { - tables, - functions: Vec::new(), - truncated_tables, - truncated_columns, - truncated_functions: false, - }); - } - - with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { - let table_rows = client - .query( - " - select n.nspname::text as schema_name, c.relname::text as table_name - from pg_class c - join pg_namespace n on n.oid = c.relnamespace - where c.relkind in ('r', 'p', 'v', 'm', 'f') - and n.nspname not in ('pg_catalog', 'information_schema') - order by n.nspname, c.relname - limit $1 - ", - &[&(MAX_EDITOR_TABLES + 1)], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; - let table_rows = if truncated_tables { - table_rows - .into_iter() - .take(MAX_EDITOR_TABLES as usize) - .collect::>() - } else { - table_rows - }; - - let mut tables = Vec::with_capacity(table_rows.len()); - let mut truncated_columns = false; - - for row in table_rows { - let schema: String = row.get(0); - let name: String = row.get(1); - let column_rows = client - .query( - " - select a.attname::text as column_name, - format_type(a.atttypid, a.atttypmod)::text as data_type - from pg_attribute a - join pg_class c on c.oid = a.attrelid - join pg_namespace n on n.oid = c.relnamespace - where n.nspname = $1 - and c.relname = $2 - and a.attnum > 0 - and not a.attisdropped - order by a.attnum - limit $3 - ", - &[&schema, &name, &(MAX_EDITOR_COLUMNS_PER_TABLE + 1)], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let columns_exceeded = column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE; - if columns_exceeded { - truncated_columns = true; - } - let columns = column_rows - .into_iter() - .take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) - .map(|column| QueryEditorColumn { - name: column.get(0), - data_type: column.get(1), - }) - .collect(); - - tables.push(QueryEditorTable { - schema, - name, - columns, - }); - } - - let function_rows = client - .query( - " - select - n.nspname::text as schema_name, - p.proname::text as function_name, - coalesce(pg_get_function_identity_arguments(p.oid), '')::text as args, - pg_get_function_result(p.oid)::text as return_type - from pg_proc p - join pg_namespace n on n.oid = p.pronamespace - where n.nspname not in ('pg_catalog', 'information_schema') - order by n.nspname, p.proname - limit $1 - ", - &[&(MAX_EDITOR_FUNCTIONS + 1)], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let truncated_functions = function_rows.len() as i64 > MAX_EDITOR_FUNCTIONS; - let functions = function_rows - .into_iter() - .take(MAX_EDITOR_FUNCTIONS as usize) - .map(|row| { - let args_raw: String = row.get(2); - QueryEditorFunction { - schema: row.get(0), - name: row.get(1), - arg_types: if args_raw.trim().is_empty() { - Vec::new() - } else { - args_raw - .split(',') - .map(|value| value.trim().to_string()) - .collect() - }, - return_type: row.get(3), - } - }) - .collect(); - - Ok(QueryEditorMetadata { - tables, - functions, - truncated_tables, - truncated_columns, - truncated_functions, - }) - }) - .await -} - -fn estimate_tokens(chars: usize) -> usize { - // Lightweight estimate good enough for budget telemetry. - (chars / 4).max(1) -} - -fn normalize_openrouter_base(base: Option<&str>) -> String { - let trimmed = base.unwrap_or("https://openrouter.ai/api/v1").trim(); - let value = if trimmed.is_empty() { - "https://openrouter.ai/api/v1" - } else { - trimmed - }; - value.trim_end_matches('/').to_string() -} - -fn truncate_on_char_boundary(value: &mut String, max_bytes: usize) { - if value.len() <= max_bytes { - return; - } - let mut truncate_at = max_bytes; - while !value.is_char_boundary(truncate_at) && truncate_at > 0 { - truncate_at -= 1; - } - value.truncate(truncate_at); -} - -fn table_matches_target(table: &QueryEditorTable, target: Option<&AskVeloxyTableRef>) -> bool { - let Some(target) = target else { - return false; - }; - table.schema.eq_ignore_ascii_case(&target.schema) && table.name.eq_ignore_ascii_case(&target.name) -} - -fn table_relevance_score(table: &QueryEditorTable, prompt_lower: &str) -> usize { - let mut score = 0usize; - let full_name = format!("{}.{}", table.schema.to_lowercase(), table.name.to_lowercase()); - if prompt_lower.contains(&table.name.to_lowercase()) { - score += 3; - } - if prompt_lower.contains(&table.schema.to_lowercase()) { - score += 2; - } - if prompt_lower.contains(&full_name) { - score += 4; - } - score -} - -fn relationship_relevance_score(edge: &ForeignKeyEdge, prompt_lower: &str) -> usize { - let from_name = format!("{}.{}", edge.from_schema.to_lowercase(), edge.from_table.to_lowercase()); - let to_name = format!("{}.{}", edge.to_schema.to_lowercase(), edge.to_table.to_lowercase()); - let mut score = 0usize; - if prompt_lower.contains(&edge.from_table.to_lowercase()) || prompt_lower.contains(&from_name) { - score += 2; - } - if prompt_lower.contains(&edge.to_table.to_lowercase()) || prompt_lower.contains(&to_name) { - score += 2; - } - score -} - -fn build_schema_context( - db_context: &AskVeloxyDbContextCache, - prompt: &str, - target_table: Option<&AskVeloxyTableRef>, -) -> String { - let prompt_lower = prompt.to_lowercase(); - let mut ranked: Vec<(&QueryEditorTable, usize, bool)> = db_context - .metadata - .tables - .iter() - .map(|table| { - ( - table, - table_relevance_score(table, &prompt_lower), - table_matches_target(table, target_table), - ) - }) - .collect(); - - ranked.sort_by(|a, b| b.2.cmp(&a.2).then_with(|| b.1.cmp(&a.1))); - - let mut schema_context = String::new(); - schema_context.push_str(&format!( - "database {} engine {:?}\n", - db_context.database_name, db_context.engine - )); - for (table, _score, _is_target) in ranked.into_iter().take(ASK_VELOXY_MAX_CONTEXT_TABLES) { - let columns = table - .columns - .iter() - .take(ASK_VELOXY_MAX_CONTEXT_COLUMNS) - .map(|column| format!("{}:{}", column.name, column.data_type)) - .collect::>() - .join(", "); - schema_context.push_str(&format!( - "table {}.{} columns [{}]\n", - table.schema, table.name, columns - )); - if schema_context.len() >= ASK_VELOXY_SCHEMA_CHAR_BUDGET { - truncate_on_char_boundary(&mut schema_context, ASK_VELOXY_SCHEMA_CHAR_BUDGET); - break; - } - } - - let mut ranked_relationships = db_context - .foreign_keys - .iter() - .map(|edge| (edge, relationship_relevance_score(edge, &prompt_lower))) - .collect::>(); - ranked_relationships.sort_by(|a, b| b.1.cmp(&a.1)); - for (edge, _score) in ranked_relationships - .into_iter() - .take(ASK_VELOXY_MAX_CONTEXT_RELATIONSHIPS) - { - schema_context.push_str(&format!( - "relationship {}.{}({}) -> {}.{}({})\n", - edge.from_schema, - edge.from_table, - edge.from_column, - edge.to_schema, - edge.to_table, - edge.to_column - )); - if schema_context.len() >= ASK_VELOXY_SCHEMA_CHAR_BUDGET { - truncate_on_char_boundary(&mut schema_context, ASK_VELOXY_SCHEMA_CHAR_BUDGET); - break; - } - } - schema_context -} - -fn classify_sql_intent(sql: &str) -> String { - let normalized = sql.trim_start().to_ascii_lowercase(); - if normalized.starts_with("select") || normalized.starts_with("with") { - return "select".to_string(); - } - if normalized.starts_with("insert") { - return "insert".to_string(); - } - if normalized.starts_with("update") { - return "update".to_string(); - } - if normalized.starts_with("delete") { - return "delete".to_string(); - } - if normalized.starts_with("explain") { - return "explain".to_string(); - } - "unknown".to_string() -} - -/// Whether every statement in `sql` is read-only. Transaction-control keywords -/// (`begin`/`commit`/`rollback`/`start`/`savepoint`) are ignored so a wrapped -/// `BEGIN; SELECT ...; COMMIT;` still counts as read-only, while any write -/// statement makes the whole batch non-read-only. -fn is_read_only_sql(sql: &str) -> bool { - let mut saw_statement = false; - for statement in sql.split(';').map(str::trim).filter(|s| !s.is_empty()) { - let normalized = statement.to_ascii_lowercase(); - let is_transaction_control = ["begin", "commit", "rollback", "start", "savepoint", "release"] - .iter() - .any(|kw| normalized.starts_with(kw)); - if is_transaction_control { - continue; - } - saw_statement = true; - match classify_sql_intent(statement).as_str() { - "select" | "explain" => {} - _ => return false, - } - } - saw_statement -} - -fn has_multiple_statements(sql: &str) -> bool { - let statements = sql - .split(';') - .map(str::trim) - .filter(|segment| !segment.is_empty()) - .count(); - statements > 1 -} - -fn validate_generated_sql(sql: &str) -> Result<(), String> { - let trimmed = sql.trim(); - if trimmed.is_empty() { - return Err("Ask Veloxy returned an empty SQL statement.".to_string()); - } - if has_multiple_statements(trimmed) { - return Err("Ask Veloxy generated multiple SQL statements. Please ask for a single statement.".to_string()); - } - Ok(()) -} - -fn extract_openrouter_message_content(payload: &Value) -> Result { - let content_value = payload - .get("choices") - .and_then(|choices| choices.get(0)) - .and_then(|choice| choice.get("message")) - .and_then(|message| message.get("content")) - .ok_or_else(|| "OpenRouter response missing choices[0].message.content".to_string())?; - - if let Some(content) = content_value.as_str() { - return Ok(content.to_string()); - } - - if let Some(items) = content_value.as_array() { - let mut merged = String::new(); - for item in items { - if let Some(text) = item.get("text").and_then(Value::as_str) { - merged.push_str(text); - } - } - if !merged.trim().is_empty() { - return Ok(merged); - } - } - - Err("OpenRouter returned an unsupported message format.".to_string()) -} - -fn parse_ask_veloxy_json(content: &str) -> Result { - if let Ok(value) = serde_json::from_str::(content) { - return Ok(value); - } - let start = content.find('{'); - let end = content.rfind('}'); - match (start, end) { - (Some(start_idx), Some(end_idx)) if end_idx > start_idx => { - serde_json::from_str::(&content[start_idx..=end_idx]) - .map_err(|error| format!("Ask Veloxy response was not valid JSON: {}", error)) - } - _ => Err("Ask Veloxy response did not contain JSON.".to_string()), - } -} - -fn parse_ask_veloxy_suggestions(generated: &Value) -> Vec { - generated - .get("suggestions") - .and_then(Value::as_array) - .map(|items| { - items - .iter() - .filter_map(Value::as_str) - .map(str::trim) - .filter(|item| !item.is_empty()) - .take(5) - .map(|item| { - let mut value = item.to_string(); - truncate_on_char_boundary(&mut value, 200); - value - }) - .collect::>() - }) - .unwrap_or_default() -} - -fn parse_ask_veloxy_chat_json(content: &str) -> Result { - if let Ok(value) = serde_json::from_str::(content) { - return Ok(value); - } - let start = content.find('{'); - let end = content.rfind('}'); - match (start, end) { - (Some(start_idx), Some(end_idx)) if end_idx > start_idx => { - serde_json::from_str::(&content[start_idx..=end_idx]) - .map_err(|error| format!("Ask Veloxy chat JSON was invalid: {}", error)) - } - _ => Err("Ask Veloxy chat response did not contain JSON.".to_string()), - } -} - -fn decode_json_quoted_string(value: &str) -> Option { - serde_json::from_str::(&format!("\"{}\"", value)).ok() -} - -fn unescape_json_string_fragment(raw: &str) -> String { - let mut out = String::with_capacity(raw.len()); - let mut chars = raw.chars().peekable(); - while let Some(ch) = chars.next() { - if ch == '\\' { - match chars.next() { - Some('n') => out.push('\n'), - Some('t') => out.push('\t'), - Some('r') => out.push('\r'), - Some('"') => out.push('"'), - Some('\\') => out.push('\\'), - Some(other) => { - out.push('\\'); - out.push(other); - } - None => out.push('\\'), - } - } else { - out.push(ch); - } - } - out -} - -fn extract_json_string_field(content: &str, key: &str, allow_partial: bool) -> Option { - let marker = format!("\"{}\"", key); - let marker_idx = content.find(&marker)?; - let mut idx = marker_idx + marker.len(); - let bytes = content.as_bytes(); - - while idx < bytes.len() && bytes[idx].is_ascii_whitespace() { - idx += 1; - } - if idx >= bytes.len() || bytes[idx] != b':' { - return None; - } - idx += 1; - while idx < bytes.len() && bytes[idx].is_ascii_whitespace() { - idx += 1; - } - if idx >= bytes.len() || bytes[idx] != b'"' { - return None; - } - idx += 1; - let start = idx; - let mut escaped = false; - while idx < bytes.len() { - let byte = bytes[idx]; - if escaped { - escaped = false; - idx += 1; - continue; - } - if byte == b'\\' { - escaped = true; - idx += 1; - continue; - } - if byte == b'"' { - let raw = &content[start..idx]; - return decode_json_quoted_string(raw) - .or_else(|| Some(unescape_json_string_fragment(raw))) - .map(|text| text.trim().to_string()) - .filter(|text| !text.is_empty()); - } - idx += 1; - } - - if allow_partial && start < bytes.len() { - let raw = &content[start..]; - let text = unescape_json_string_fragment(raw).trim().to_string(); - if !text.is_empty() { - return Some(text); - } - } - None -} - -fn extract_message_from_loose_json(content: &str) -> Option { - let trimmed = content.trim(); - if trimmed.is_empty() { - return None; - } - let unwrapped = trimmed - .strip_prefix("```json") - .or_else(|| trimmed.strip_prefix("```JSON")) - .map(str::trim_start) - .unwrap_or(trimmed); - let unwrapped = unwrapped.strip_suffix("```").unwrap_or(unwrapped).trim(); - - ["message", "reply", "content"] - .iter() - .find_map(|key| extract_json_string_field(unwrapped, key, false)) - .or_else(|| { - ["message", "reply", "content"] - .iter() - .find_map(|key| extract_json_string_field(unwrapped, key, true)) - }) -} - -fn looks_like_json_response(content: &str) -> bool { - let trimmed = content.trim_start(); - trimmed.starts_with('{') || trimmed.starts_with("```") -} - -fn streaming_display_text(accumulated: &str) -> String { - let trimmed = accumulated.trim(); - if trimmed.is_empty() { - return String::new(); - } - if let Some(text) = extract_message_from_loose_json(trimmed) { - return text; - } - if !looks_like_json_response(trimmed) { - return trimmed.to_string(); - } - String::new() -} - -fn parse_chat_message(value: &Value) -> Option { - if let Some(text) = value - .as_str() - .map(str::trim) - .filter(|text| !text.is_empty()) - .map(str::to_string) - { - return Some(text); - } - value - .get("message") - .and_then(Value::as_str) - .or_else(|| value.get("reply").and_then(Value::as_str)) - .map(str::trim) - .filter(|text| !text.is_empty()) - .map(str::to_string) -} - -type ParsedAskVeloxyChat = ( - String, - Vec, - Vec, - Option, - bool, - bool, -); - -fn parse_ask_veloxy_chat_content(message_content: &str) -> ParsedAskVeloxyChat { - match parse_ask_veloxy_chat_json(message_content) { - Ok(value) => { - let message = - parse_chat_message(&value).unwrap_or_else(|| message_content.trim().to_string()); - let mut draft = value - .get("sqlDraft") - .and_then(Value::as_str) - .or_else(|| value.get("sql_draft").and_then(Value::as_str)) - .map(str::trim) - .filter(|text| !text.is_empty()) - .map(str::to_string); - if draft.is_none() { - draft = extract_sql_draft_from_text(&message); - } - let suggestions = value - .get("suggestions") - .and_then(Value::as_array) - .map(|items| { - items - .iter() - .filter_map(Value::as_str) - .map(str::trim) - .filter(|text| !text.is_empty()) - .take(5) - .map(str::to_string) - .collect::>() - }) - .unwrap_or_default(); - let warnings = value - .get("warnings") - .and_then(Value::as_array) - .map(|items| { - items - .iter() - .filter_map(Value::as_str) - .map(str::to_string) - .collect::>() - }) - .unwrap_or_default(); - let needs_sql_generation = parse_bool_field(&value, "needsSqlGeneration", draft.is_some()); - let needs_clarification = parse_bool_field(&value, "needsClarification", false); - ( - message, - suggestions, - warnings, - draft, - needs_sql_generation, - needs_clarification, - ) - } - Err(_) => { - let normalized_message = extract_message_from_loose_json(message_content) - .unwrap_or_else(|| { - if looks_like_json_response(message_content) { - String::new() - } else { - message_content.trim().to_string() - } - }); - let mut warnings = vec!["Model returned non-JSON chat output. Parsed in tolerant mode.".to_string()]; - if normalized_message.is_empty() && looks_like_json_response(message_content) { - warnings.push("Response JSON could not be parsed. Try asking again.".to_string()); - } - let draft = extract_sql_draft_from_text(&normalized_message); - let needs_sql_generation = draft.is_some(); - ( - normalized_message, - Vec::new(), - warnings, - draft, - needs_sql_generation, - false, - ) - } - } -} - -fn extract_openrouter_stream_delta(data: &str) -> Option { - let payload: Value = serde_json::from_str(data).ok()?; - payload - .get("choices") - .and_then(|choices| choices.get(0)) - .and_then(|choice| choice.get("delta")) - .and_then(|delta| delta.get("content")) - .and_then(Value::as_str) - .filter(|text| !text.is_empty()) - .map(str::to_string) -} - -fn emit_veloxy_stream_chunk(app: &AppHandle, chunk: VeloxyStreamChunk) { - let _ = app.emit("veloxy-stream-chunk", chunk); -} - -fn extract_openrouter_finish_reason(data: &str) -> Option { - let payload: Value = serde_json::from_str(data).ok()?; - payload - .get("choices") - .and_then(|choices| choices.get(0)) - .and_then(|choice| choice.get("finish_reason")) - .and_then(Value::as_str) - .map(str::to_string) -} - -async fn stream_openrouter_chat_completion( - app: &AppHandle, - client: &reqwest::Client, - endpoint: &str, - api_key: &str, - model: &str, - system_prompt: &str, - user_prompt: &str, - request_id: &str, - cancel: Arc, -) -> Result<(String, bool), String> { - let response = client - .post(endpoint) - .header("Authorization", format!("Bearer {}", api_key)) - .header("Content-Type", "application/json") - .json(&serde_json::json!({ - "model": model, - "temperature": 0.2, - "max_tokens": ASK_VELOXY_MAX_CHAT_TOKENS, - "stream": true, - "messages": [ - { "role": "system", "content": system_prompt }, - { "role": "user", "content": user_prompt } - ] - })) - .send() - .await - .map_err(|error| format!("OpenRouter request failed: {}", error))?; - - let status = response.status(); - if !status.is_success() { - let body = response - .text() - .await - .unwrap_or_else(|_| "Unknown OpenRouter error".to_string()); - if let Ok(payload) = serde_json::from_str::(&body) { - let message = payload - .get("error") - .and_then(|error| error.get("message")) - .and_then(Value::as_str) - .unwrap_or("Unknown OpenRouter error"); - return Err(format!("OpenRouter error ({}): {}", status.as_u16(), message)); - } - return Err(format!("OpenRouter error ({}): {}", status.as_u16(), body)); - } - - let mut stream = response.bytes_stream(); - let mut buffer = String::new(); - let mut accumulated = String::new(); - let mut last_display_len = 0usize; - let mut hit_token_limit = false; - - while let Some(chunk) = stream.next().await { - if cancel.load(Ordering::Relaxed) { - return Ok((accumulated, hit_token_limit)); - } - let bytes = chunk.map_err(|error| format!("OpenRouter stream read failed: {}", error))?; - buffer.push_str(&String::from_utf8_lossy(&bytes)); - - while let Some(line_end) = buffer.find('\n') { - let line = buffer[..line_end].trim_end_matches('\r').to_string(); - buffer.drain(..=line_end); - - if !line.starts_with("data: ") { - continue; - } - let data = line["data: ".len()..].trim(); - if data == "[DONE]" { - continue; - } - if extract_openrouter_finish_reason(data).as_deref() == Some("length") { - hit_token_limit = true; - } - if let Some(delta) = extract_openrouter_stream_delta(data) { - accumulated.push_str(&delta); - let display = streaming_display_text(&accumulated); - let display_delta = if display.len() > last_display_len { - display[last_display_len..].to_string() - } else { - String::new() - }; - last_display_len = display.len(); - if !display_delta.is_empty() { - emit_veloxy_stream_chunk( - app, - VeloxyStreamChunk { - request_id: request_id.to_string(), - delta: display_delta, - done: false, - message: None, - suggestions: Vec::new(), - warnings: Vec::new(), - sql_draft: None, - needs_sql_generation: false, - needs_clarification: false, - }, - ); - } - } - } - } - - Ok((accumulated, hit_token_limit)) -} - -#[tauri::command] -pub async fn cancel_veloxy_request(state: State<'_, AppState>) -> Result<(), String> { - if let Some(cancel) = state.veloxy_cancel.read().await.as_ref() { - cancel.store(true, Ordering::Relaxed); - } - Ok(()) -} - -#[tauri::command] -pub async fn chat_with_db( - app: AppHandle, - state: State<'_, AppState>, - input: AskVeloxyChatRequest, -) -> Result { - let natural_prompt = input.natural_prompt.trim(); - if natural_prompt.is_empty() { - return Err("Ask Veloxy prompt cannot be empty.".to_string()); - } - if input.provider_config.api_key.trim().is_empty() { - return Err("OpenRouter API key is required.".to_string()); - } - if input.provider_config.model.trim().is_empty() { - return Err("OpenRouter model is required.".to_string()); - } - - let (connection_id, engine) = - resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let stored_connection = load_connection(&app, &connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - let db_context = get_or_build_ask_veloxy_db_context(&app, &state, &connection_id, engine).await?; - let schema_context = build_schema_context(&db_context, natural_prompt, input.target_table.as_ref()); - let conversation_key = ask_veloxy_conversation_key(&connection_id, &stored_connection.database); - let history = state - .ask_veloxy_conversations - .read() - .await - .get(&conversation_key) - .cloned() - .unwrap_or_default(); - - let history_block = history - .iter() - .rev() - .take(8) - .rev() - .map(|message| format!("{}: {}", message.role, message.text)) - .collect::>() - .join("\n"); - - let mut user_prompt = format!( - "Engine: {:?}\nDatabase: {}\nTask: {}\nMaxRows: {}\nRecentConversation:\n{}\nSchemaContext:\n{}\n", - db_context.engine, - db_context.database_name, - natural_prompt, - input.max_rows.unwrap_or(MAX_QUERY_ROWS), - history_block, - schema_context - ); - truncate_on_char_boundary(&mut user_prompt, ASK_VELOXY_PROMPT_CHAR_BUDGET); - - let system_prompt = "You are Ask Veloxy chat mode. Return JSON when possible with keys: message (string), suggestions (array of strings), sqlDraft (string optional), needsSqlGeneration (boolean), needsClarification (boolean), warnings (array of strings). If JSON is not possible, return helpful plain text."; - let base_url = normalize_openrouter_base(input.provider_config.base_url.as_deref()); - let endpoint = format!("{}/chat/completions", base_url); - let client = state.openrouter_client.get_or_init(reqwest::Client::new); - let request_id = input - .request_id - .clone() - .filter(|value| !value.trim().is_empty()) - .unwrap_or_else(|| format!("req-{}", uuid::Uuid::new_v4())); - - let cancel = Arc::new(AtomicBool::new(false)); - { - let mut guard = state.veloxy_cancel.write().await; - *guard = Some(cancel.clone()); - } - - let (message_content, hit_token_limit) = stream_openrouter_chat_completion( - &app, - client, - &endpoint, - input.provider_config.api_key.trim(), - input.provider_config.model.trim(), - system_prompt, - &user_prompt, - &request_id, - cancel.clone(), - ) - .await?; - - { - let mut guard = state.veloxy_cancel.write().await; - *guard = None; - } - - let (message, suggestions, mut warnings, sql_draft, needs_sql_generation, needs_clarification) = - parse_ask_veloxy_chat_content(&message_content); - - if cancel.load(Ordering::Relaxed) { - warnings.push("Stopped early.".to_string()); - } - if hit_token_limit { - warnings.push(format!( - "Response may be truncated (model output limit of {} tokens).", - ASK_VELOXY_MAX_CHAT_TOKENS - )); - } - - emit_veloxy_stream_chunk( - &app, - VeloxyStreamChunk { - request_id: request_id.clone(), - delta: String::new(), - done: true, - message: Some(message.clone()), - suggestions: suggestions.clone(), - warnings: warnings.clone(), - sql_draft: sql_draft.clone(), - needs_sql_generation, - needs_clarification, - }, - ); - - { - let mut conversations = state.ask_veloxy_conversations.write().await; - let bucket = conversations.entry(conversation_key).or_default(); - bucket.push(AskVeloxyConversationMessage { - id: format!("msg-{}", uuid::Uuid::new_v4()), - role: "user".to_string(), - mode: "chat".to_string(), - text: natural_prompt.to_string(), - created_at: now_epoch_seconds(), - sql_draft: None, - }); - bucket.push(AskVeloxyConversationMessage { - id: format!("msg-{}", uuid::Uuid::new_v4()), - role: "assistant".to_string(), - mode: "chat".to_string(), - text: message.clone(), - created_at: now_epoch_seconds(), - sql_draft: sql_draft.clone(), - }); - if bucket.len() > ASK_VELOXY_MAX_HISTORY_MESSAGES { - let remove_count = bucket.len() - ASK_VELOXY_MAX_HISTORY_MESSAGES; - bucket.drain(0..remove_count); - } - } - - Ok(AskVeloxyChatResponse { - message, - suggestions, - warnings, - sql_draft, - needs_sql_generation, - needs_clarification, - }) -} - -#[tauri::command] -pub async fn load_veloxy_conversation( - app: AppHandle, - state: State<'_, AppState>, - connection_id: Option, -) -> Result { - let (resolved_connection_id, _) = resolve_connection_engine(&app, &state, connection_id).await?; - let stored_connection = load_connection(&app, &resolved_connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - let key = ask_veloxy_conversation_key(&resolved_connection_id, &stored_connection.database); - let messages = state - .ask_veloxy_conversations - .read() - .await - .get(&key) - .cloned() - .unwrap_or_default(); - Ok(AskVeloxyConversationResponse { messages }) -} - -#[tauri::command] -pub async fn clear_veloxy_conversation( - app: AppHandle, - state: State<'_, AppState>, - connection_id: Option, -) -> Result<(), String> { - let (resolved_connection_id, _) = resolve_connection_engine(&app, &state, connection_id).await?; - let stored_connection = load_connection(&app, &resolved_connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - let key = ask_veloxy_conversation_key(&resolved_connection_id, &stored_connection.database); - state.ask_veloxy_conversations.write().await.remove(&key); - Ok(()) -} - -#[tauri::command] -pub async fn generate_sql_from_nl( - app: AppHandle, - state: State<'_, AppState>, - input: AskVeloxyRequest, -) -> Result { - let natural_prompt = input.natural_prompt.trim(); - if natural_prompt.is_empty() { - return Err("Ask Veloxy prompt cannot be empty.".to_string()); - } - if input.provider_config.api_key.trim().is_empty() { - return Err("OpenRouter API key is required.".to_string()); - } - if input.provider_config.model.trim().is_empty() { - return Err("OpenRouter model is required.".to_string()); - } - - let (connection_id, engine) = - resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let db_context = get_or_build_ask_veloxy_db_context(&app, &state, &connection_id, engine).await?; - let schema_context = build_schema_context(&db_context, natural_prompt, input.target_table.as_ref()); - - let mut user_prompt = format!( - "Engine: {:?}\nDatabase: {}\nTask: {}\nMaxRows: {}\nSchemaContext:\n{}\n", - db_context.engine, - db_context.database_name, - natural_prompt, - input.max_rows.unwrap_or(MAX_QUERY_ROWS), - schema_context - ); - truncate_on_char_boundary(&mut user_prompt, ASK_VELOXY_PROMPT_CHAR_BUDGET); - - let system_prompt = "You are Ask Veloxy. Return JSON only with keys: sql (string), intent (string), confidence (number 0..1), explanation (string), suggestions (array of short strings), warnings (array of strings). Generate exactly one SQL statement, keep explanation concise, and never include markdown."; - let base_url = normalize_openrouter_base(input.provider_config.base_url.as_deref()); - let endpoint = format!("{}/chat/completions", base_url); - - let client = state.openrouter_client.get_or_init(reqwest::Client::new); - let response = client - .post(&endpoint) - .header("Authorization", format!("Bearer {}", input.provider_config.api_key.trim())) - .header("Content-Type", "application/json") - .json(&serde_json::json!({ - "model": input.provider_config.model.trim(), - "temperature": 0.1, - "max_tokens": 500, - "messages": [ - { "role": "system", "content": system_prompt }, - { "role": "user", "content": user_prompt } - ] - })) - .send() - .await - .map_err(|error| format!("OpenRouter request failed: {}", error))?; - - let status = response.status(); - let payload = response - .json::() - .await - .map_err(|error| format!("Invalid OpenRouter JSON response: {}", error))?; - if !status.is_success() { - let message = payload - .get("error") - .and_then(|error| error.get("message")) - .and_then(Value::as_str) - .unwrap_or("Unknown OpenRouter error"); - return Err(format!("OpenRouter error ({}): {}", status.as_u16(), message)); - } - - let message_content = extract_openrouter_message_content(&payload)?; - let generated = parse_ask_veloxy_json(&message_content)?; - let sql = generated - .get("sql") - .and_then(Value::as_str) - .unwrap_or_default() - .trim() - .to_string(); - validate_generated_sql(&sql)?; - - let mut warnings = generated - .get("warnings") - .and_then(Value::as_array) - .map(|items| { - items - .iter() - .filter_map(Value::as_str) - .map(str::to_string) - .collect::>() - }) - .unwrap_or_default(); - - let intent = generated - .get("intent") - .and_then(Value::as_str) - .map(str::to_string) - .unwrap_or_else(|| classify_sql_intent(&sql)); - let confidence = generated - .get("confidence") - .and_then(Value::as_f64) - .unwrap_or(0.6) - .clamp(0.0, 1.0); - let explanation = generated - .get("explanation") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(|value| { - let mut truncated = value.to_string(); - truncate_on_char_boundary(&mut truncated, 350); - truncated - }); - let suggestions = parse_ask_veloxy_suggestions(&generated); - - if intent != "select" { - warnings.push("Generated SQL is not read-only. Review before execution.".to_string()); - } - if confidence < 0.5 { - warnings.push("Low confidence result. Review SQL carefully.".to_string()); - } - - let token_stats = AskVeloxyTokenStats { - schema_chars: schema_context.len(), - schema_tokens_estimate: estimate_tokens(schema_context.len()), - prompt_chars: user_prompt.len() + system_prompt.len(), - prompt_tokens_estimate: estimate_tokens(user_prompt.len() + system_prompt.len()), - }; - - Ok(AskVeloxyResponse { - sql, - intent, - confidence, - explanation, - suggestions, - warnings, - token_stats, - }) -} - -#[tauri::command] -pub async fn lint_sql( - app: AppHandle, - state: State<'_, AppState>, - input: LintSqlRequest, -) -> Result { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let sql = input.sql.trim().to_string(); - if sql.is_empty() { - return Ok(LintSqlResult { - diagnostics: Vec::new(), - }); - } - if sql.len() > MAX_LINT_SQL_BYTES { - return Err("SQL is too large to lint in the editor.".to_string()); - } - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, sql, |client, sql| async move { - let lint_sql = format!("EXPLAIN {}", sql); - let diagnostics = match client.simple_query(&lint_sql).await { - Ok(_) => Vec::new(), - Err(error) => { - let (line, column) = error_line_column(&error, &sql) - .map(|(l, c)| (Some(l), Some(c))) - .unwrap_or((None, None)); - vec![SqlDiagnostic { - message: map_pg_err(error, Some(sql.as_str())), - severity: "error".to_string(), - line, - column, - end_line: line, - end_column: column.map(|value| value + 1), - }] - } - }; - Ok(LintSqlResult { diagnostics }) - }) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let lint_sql = format!("EXPLAIN {}", sql); - let diagnostics = match sqlx::query(&lint_sql).execute(&pool).await { - Ok(_) => Vec::new(), - Err(error) => vec![SqlDiagnostic { - message: error.to_string(), - severity: "error".to_string(), - line: None, - column: None, - end_line: None, - end_column: None, - }], - }; - Ok(LintSqlResult { diagnostics }) - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - let lint_sql = format!("EXPLAIN QUERY PLAN {}", sql); - let diagnostics = match sqlx::query(&lint_sql).execute(&pool).await { - Ok(_) => Vec::new(), - Err(error) => vec![SqlDiagnostic { - message: error.to_string(), - severity: "error".to_string(), - line: None, - column: None, - end_line: None, - end_column: None, - }], - }; - Ok(LintSqlResult { diagnostics }) - } - } -} - -#[tauri::command] -pub async fn get_tables( - app: AppHandle, - state: State<'_, AppState>, - connection_id: Option, -) -> Result, String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { - let rows = client - .query( - " - select table_schema, table_name - from information_schema.tables - where table_type = 'BASE TABLE' - and table_schema not in ('pg_catalog', 'information_schema') - order by table_schema, table_name - ", - &[], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - Ok(rows - .into_iter() - .map(|row| { - let schema: String = row.get(0); - let name: String = row.get(1); - let preview_query = format!( - "select * from \"{}\".\"{}\" limit 100;", - quote_identifier(&schema), - quote_identifier(&name) - ); - - TableInfo { - schema, - name, - preview_query, - } - }) - .collect()) - }) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let database = load_connection(&app, &connection_id)? - .map(|connection| connection.database) - .unwrap_or_default(); - let rows = sqlx::query( - " - select table_schema, table_name - from information_schema.tables - where table_type = 'BASE TABLE' - and table_schema = ? - and table_schema not in ('information_schema', 'mysql', 'performance_schema', 'sys') - order by table_schema, table_name - ", - ) - .bind(&database) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut tables = Vec::new(); - for row in rows { - let schema: String = mysql_get_string(&row, 0, "table_schema", "get_tables")?; - let name: String = mysql_get_string(&row, 1, "table_name", "get_tables")?; - tables.push(TableInfo { - preview_query: format!("select * from `{}`.`{}` limit 100;", schema, name), - schema, - name, - }); - } - Ok(tables) - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - let rows = sqlx::query( - " - select name - from sqlite_master - where type = 'table' - and name not like 'sqlite_%' - order by name - ", - ) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut tables = Vec::new(); - for row in rows { - let name: String = sqlite_get_idx(&row, 0, "name", "get_tables")?; - require_safe_identifier(&name, "table name")?; - tables.push(TableInfo { - schema: "main".to_string(), - preview_query: format!("select * from \"{}\" limit 100;", quote_identifier(&name)), - name, - }); - } - Ok(tables) - } - } -} - -#[tauri::command] -pub async fn get_schema( - app: AppHandle, - state: State<'_, AppState>, - input: SchemaRequest, -) -> Result, String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let schema_request = input.clone(); - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry( - &app, - &state, - &connection_id, - schema_request, - |client, input| async move { - let rows = client - .query( - " - select table_schema, table_name, column_name, data_type, is_nullable - from information_schema.columns - where table_schema = $1 and table_name = $2 - order by ordinal_position - ", - &[&input.table_schema, &input.table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - Ok(rows - .into_iter() - .map(|row| ColumnInfo { - table_schema: row.get(0), - table_name: row.get(1), - column_name: row.get(2), - data_type: row.get(3), - is_nullable: row.get::<_, String>(4) == "YES", - }) - .collect()) - }, - ) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let rows = sqlx::query( - " - select table_schema, table_name, column_name, data_type, is_nullable - from information_schema.columns - where table_schema = ? and table_name = ? - order by ordinal_position - ", - ) - .bind(&schema_request.table_schema) - .bind(&schema_request.table_name) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut columns = Vec::new(); - for row in rows { - columns.push(ColumnInfo { - table_schema: mysql_get_string(&row, 0, "table_schema", "get_schema")?, - table_name: mysql_get_string(&row, 1, "table_name", "get_schema")?, - column_name: mysql_get_string(&row, 2, "column_name", "get_schema")?, - data_type: mysql_get_string(&row, 3, "data_type", "get_schema")?, - is_nullable: mysql_get_string(&row, 4, "is_nullable", "get_schema")? == "YES", - }); - } - Ok(columns) - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - require_safe_identifier(&schema_request.table_name, "table name")?; - let pragma_sql = format!( - "PRAGMA table_info(\"{}\");", - quote_identifier(&schema_request.table_name) - ); - let rows = sqlx::query(&pragma_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut columns = Vec::new(); - for row in rows { - let col_name: String = sqlite_get_name(&row, "name", "get_schema")?; - let col_type: String = sqlite_get_name(&row, "type", "get_schema")?; - let notnull: i64 = sqlite_get_name(&row, "notnull", "get_schema")?; - columns.push(ColumnInfo { - table_schema: "main".to_string(), - table_name: schema_request.table_name.clone(), - column_name: col_name, - data_type: col_type, - is_nullable: notnull == 0, - }); - } - Ok(columns) - } - } -} - -fn veloxdb_unique_constraint_name(table_name: &str, column_name: &str) -> String { - // Postgres constraint names are limited to 63 bytes. - // Keep this deterministic so we can drop the exact constraint later. - let suffix = "_uniq"; - let max_base_len = 63usize.saturating_sub(suffix.len()); - - let mut base = format!("veloxdb_{}_{}", table_name, column_name); - base.truncate(max_base_len); - - format!("{}{}", base, suffix) -} - -#[tauri::command] -pub async fn get_table_properties( - app: AppHandle, - state: State<'_, AppState>, - input: SchemaRequest, -) -> Result, String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let ctx = input.clone(); - - if engine == DatabaseEngine::Mysql { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let rows = sqlx::query( - " - select - c.table_schema, - c.table_name, - c.column_name, - c.data_type, - c.is_nullable, - c.column_default, - c.extra - from information_schema.columns c - where c.table_schema = ? and c.table_name = ? - order by c.ordinal_position - ", - ) - .bind(&ctx.table_schema) - .bind(&ctx.table_name) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - - let pk_rows = sqlx::query( - " - select column_name - from information_schema.key_column_usage - where table_schema = ? and table_name = ? and constraint_name = 'PRIMARY' - ", - ) - .bind(&ctx.table_schema) - .bind(&ctx.table_name) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let pk_cols: HashSet = pk_rows - .into_iter() - .map(|row| mysql_get_string(&row, 0, "column_name", "get_table_properties")) - .collect::, _>>()?; - - let unique_rows = sqlx::query( - " - select index_name, column_name, seq_in_index - from information_schema.statistics - where table_schema = ? - and table_name = ? - and non_unique = 0 - order by index_name, seq_in_index - ", - ) - .bind(&ctx.table_schema) - .bind(&ctx.table_name) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut unique_by_index: HashMap> = HashMap::new(); - for row in unique_rows { - let index_name: String = mysql_get_string(&row, 0, "index_name", "get_table_properties")?; - if index_name == "PRIMARY" { - continue; - } - let column_name: String = mysql_get_string(&row, 1, "column_name", "get_table_properties")?; - unique_by_index.entry(index_name).or_default().push(column_name); - } - let mut unique_cols: HashSet = HashSet::new(); - let mut composite_unique_cols: HashSet = HashSet::new(); - for cols in unique_by_index.values() { - for col in cols { - unique_cols.insert(col.clone()); - } - if cols.len() > 1 { - for col in cols { - composite_unique_cols.insert(col.clone()); - } - } - } - - let mut properties = Vec::new(); - for row in rows { - let column_name: String = mysql_get_string(&row, 2, "column_name", "get_table_properties")?; - let is_primary_key = pk_cols.contains(&column_name); - let is_unique = is_primary_key || unique_cols.contains(&column_name); - let is_part_of_composite_unique = composite_unique_cols.contains(&column_name); - let extra: String = mysql_get_string(&row, 6, "extra", "get_table_properties")?; - let lower_extra = extra.to_lowercase(); - properties.push(ColumnProperties { - table_schema: mysql_get_string(&row, 0, "table_schema", "get_table_properties")?, - table_name: mysql_get_string(&row, 1, "table_name", "get_table_properties")?, - column_name, - data_type: mysql_get_string(&row, 3, "data_type", "get_table_properties")?, - is_nullable: mysql_get_string(&row, 4, "is_nullable", "get_table_properties")? == "YES", - is_primary_key, - is_unique, - is_part_of_composite_unique, - column_default: mysql_get_optional_string(&row, 5, "column_default", "get_table_properties")?, - is_identity: lower_extra.contains("auto_increment"), - identity_generation: if lower_extra.contains("auto_increment") { - Some("BY DEFAULT".to_string()) - } else { - None - }, - is_generated: if lower_extra.contains("generated") { - Some("ALWAYS".to_string()) - } else { - None - }, - }); - } - return Ok(properties); - } - - if engine == DatabaseEngine::Sqlite { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - require_safe_identifier(&ctx.table_name, "table name")?; - let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(&ctx.table_name)); - let rows = sqlx::query(&pragma_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let index_list_sql = format!("PRAGMA index_list(\"{}\");", quote_identifier(&ctx.table_name)); - let index_rows = sqlx::query(&index_list_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut unique_cols: HashSet = HashSet::new(); - let mut composite_unique_cols: HashSet = HashSet::new(); - for index in index_rows { - let is_unique = sqlite_get_name::(&index, "unique", "get_table_properties")? == 1; - if !is_unique { - continue; - } - let origin = sqlite_get_name::(&index, "origin", "get_table_properties")?; - if origin == "pk" { - continue; - } - let index_name = sqlite_get_name::(&index, "name", "get_table_properties")?; - require_safe_identifier(&index_name, "index name")?; - let info_sql = format!("PRAGMA index_info(\"{}\");", quote_identifier(&index_name)); - let info_rows = sqlx::query(&info_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut cols: Vec = Vec::new(); - for info in info_rows { - if let Ok(name) = sqlite_get_name::(&info, "name", "get_table_properties") { - cols.push(name); - } - } - for col in &cols { - unique_cols.insert(col.clone()); - } - if cols.len() > 1 { - for col in cols { - composite_unique_cols.insert(col); - } - } - } - let mut properties = Vec::new(); - for row in rows { - let column_name: String = sqlite_get_name(&row, "name", "get_table_properties")?; - let is_primary_key = sqlite_get_name::(&row, "pk", "get_table_properties")? == 1; - let is_unique = is_primary_key || unique_cols.contains(&column_name); - let is_part_of_composite_unique = composite_unique_cols.contains(&column_name); - properties.push(ColumnProperties { - table_schema: "main".to_string(), - table_name: ctx.table_name.clone(), - column_name, - data_type: sqlite_get_name(&row, "type", "get_table_properties")?, - is_nullable: sqlite_get_name::(&row, "notnull", "get_table_properties")? == 0, - is_primary_key, - is_unique, - is_part_of_composite_unique, - column_default: sqlite_get_name::>(&row, "dflt_value", "get_table_properties")?, - is_identity: false, - identity_generation: None, - is_generated: None, - }); - } - return Ok(properties); - } - - with_pool_client_retry(&app, &state, &connection_id, ctx, |client, input| async move { - let columns = client - .query( - " - select - c.table_schema, - c.table_name, - c.column_name, - c.data_type, - c.is_nullable, - c.column_default, - c.is_identity, - c.identity_generation, - c.is_generated - from information_schema.columns c - where c.table_schema = $1 and c.table_name = $2 - order by c.ordinal_position - ", - &[&input.table_schema, &input.table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let primary_keys = client - .query( - " - select kcu.column_name - from information_schema.table_constraints tc - join information_schema.key_column_usage kcu - on tc.constraint_name = kcu.constraint_name - and tc.table_schema = kcu.table_schema - where tc.table_schema = $1 - and tc.table_name = $2 - and tc.constraint_type = 'PRIMARY KEY' - order by kcu.ordinal_position - ", - &[&input.table_schema, &input.table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let primary_key_columns: HashSet = primary_keys - .into_iter() - .filter_map(|row| Some(row.get::<_, String>(0))) - .collect(); - - let unique_constraints = client - .query( - " - select tc.constraint_name, kcu.column_name, kcu.ordinal_position - from information_schema.table_constraints tc - join information_schema.key_column_usage kcu - on tc.constraint_name = kcu.constraint_name - and tc.table_schema = kcu.table_schema - where tc.table_schema = $1 - and tc.table_name = $2 - and tc.constraint_type = 'UNIQUE' - order by tc.constraint_name, kcu.ordinal_position - ", - &[&input.table_schema, &input.table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let mut unique_by_name: HashMap> = HashMap::new(); - for row in unique_constraints { - let constraint_name: String = row.get(0); - let column_name: String = row.get(1); - unique_by_name.entry(constraint_name).or_default().push(column_name); - } - - let mut unique_columns: HashSet = HashSet::new(); - let mut composite_unique_columns: HashSet = HashSet::new(); - - for (_constraint_name, cols) in unique_by_name { - for c in &cols { - unique_columns.insert(c.clone()); - } - if cols.len() > 1 { - for c in &cols { - composite_unique_columns.insert(c.clone()); - } - } - } - - Ok(columns - .into_iter() - .map(|row| { - let table_schema: String = row.get(0); - let table_name: String = row.get(1); - let column_name: String = row.get(2); - let data_type: String = row.get(3); - let is_nullable = row.get::<_, String>(4) == "YES"; - let column_default: Option = row.get(5); - let is_identity = row.get::<_, Option>(6).as_deref() == Some("YES"); - let identity_generation: Option = row.get(7); - let is_generated: Option = row.get(8); - - let is_primary_key = primary_key_columns.contains(&column_name); - let is_unique = is_primary_key || unique_columns.contains(&column_name); - let is_part_of_composite_unique = composite_unique_columns.contains(&column_name); - - ColumnProperties { - table_schema, - table_name, - column_name, - data_type, - is_nullable, - is_primary_key, - is_unique, - is_part_of_composite_unique, - column_default, - is_identity, - identity_generation, - is_generated, - } - }) - .collect()) - }) - .await -} - -#[tauri::command] -pub async fn apply_table_properties( - app: AppHandle, - state: State<'_, AppState>, - input: TablePropertiesApplyRequest, -) -> Result<(), String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - if engine != DatabaseEngine::Postgres { - return Err(format!( - "Table property editing is not supported for {} connections yet.", - match engine { - DatabaseEngine::Postgres => "PostgreSQL", - DatabaseEngine::Mysql => "MySQL", - DatabaseEngine::Sqlite => "SQLite", - } - )); - } - - with_pool_client_retry(&app, &state, &connection_id, input, |mut client, input| async move { - let table_schema = input.table_schema; - let table_name = input.table_name; - let columns = input.columns; - - require_safe_identifier(&table_schema, "schema name")?; - require_safe_identifier(&table_name, "table name")?; - - let current_columns = client - .query( - " - select column_name, is_nullable - from information_schema.columns - where table_schema = $1 and table_name = $2 - ", - &[&table_schema, &table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let mut current_nullable: HashMap = HashMap::new(); - for row in current_columns { - let column_name: String = row.get(0); - let is_nullable = row.get::<_, String>(1) == "YES"; - current_nullable.insert(column_name, is_nullable); - } - - let primary_keys = client - .query( - " - select kcu.column_name - from information_schema.table_constraints tc - join information_schema.key_column_usage kcu - on tc.constraint_name = kcu.constraint_name - and tc.table_schema = kcu.table_schema - where tc.table_schema = $1 - and tc.table_name = $2 - and tc.constraint_type = 'PRIMARY KEY' - ", - &[&table_schema, &table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let primary_key_columns: HashSet = primary_keys - .into_iter() - .filter_map(|row| Some(row.get::<_, String>(0))) - .collect(); - - let unique_constraints = client - .query( - " - select tc.constraint_name, kcu.column_name, kcu.ordinal_position - from information_schema.table_constraints tc - join information_schema.key_column_usage kcu - on tc.constraint_name = kcu.constraint_name - and tc.table_schema = kcu.table_schema - where tc.table_schema = $1 - and tc.table_name = $2 - and tc.constraint_type = 'UNIQUE' - order by tc.constraint_name, kcu.ordinal_position - ", - &[&table_schema, &table_name], - ) - .await - .map_err(|error| map_pg_err(error, None))?; - - let mut unique_by_name: HashMap> = HashMap::new(); - for row in unique_constraints { - let constraint_name: String = row.get(0); - let column_name: String = row.get(1); - unique_by_name.entry(constraint_name).or_default().push(column_name); - } - - let mut composite_unique_columns: HashSet = HashSet::new(); - let mut single_unique_constraint_names_by_column: HashMap> = HashMap::new(); - - for (constraint_name, cols) in &unique_by_name { - if cols.len() > 1 { - for c in cols { - composite_unique_columns.insert(c.clone()); - } - } else if cols.len() == 1 { - let c = &cols[0]; - single_unique_constraint_names_by_column - .entry(c.clone()) - .or_default() - .push(constraint_name.clone()); - } - } - - let mut desired_by_column: HashMap = HashMap::new(); - for update in columns { - desired_by_column.insert(update.column_name, (update.is_nullable, update.is_unique)); - } - - let txn = client.transaction().await.map_err(|error| map_pg_err(error, None))?; - - // 1) Nullable changes - for (column_name, (desired_is_nullable, _desired_is_unique)) in &desired_by_column { - let current_is_nullable = current_nullable - .get(column_name) - .ok_or_else(|| format!("Unknown column: {}", column_name))?; - - if *current_is_nullable == *desired_is_nullable { - continue; - } - - let qualified_table = format!( - "\"{}\".\"{}\"", - quote_identifier(&table_schema), - quote_identifier(&table_name) - ); - - require_safe_identifier(column_name, "column name")?; - let qualified_column = format!("\"{}\"", quote_identifier(column_name)); - - if *desired_is_nullable { - let sql = format!( - "ALTER TABLE {} ALTER COLUMN {} DROP NOT NULL", - qualified_table, qualified_column - ); - txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - } else { - let sql = format!( - "ALTER TABLE {} ALTER COLUMN {} SET NOT NULL", - qualified_table, qualified_column - ); - txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - } - } - - // 2) UNIQUE changes (v1: only support single-column UNIQUE constraints) - for (column_name, (_desired_is_nullable, desired_is_unique)) in &desired_by_column { - let is_primary_key = primary_key_columns.contains(column_name); - let is_part_of_composite_unique = composite_unique_columns.contains(column_name); - - if !*desired_is_unique { - if is_primary_key { - return Err(format!( - "Cannot disable UNIQUE for primary key column: {}", - column_name - )); - } - - if is_part_of_composite_unique { - return Err(format!( - "Cannot disable UNIQUE for column in a composite UNIQUE constraint: {}", - column_name - )); - } - } - - // Compute current uniqueness: - let has_single_unique = single_unique_constraint_names_by_column - .get(column_name) - .map(|names| !names.is_empty()) - .unwrap_or(false); - - let current_is_unique = is_primary_key || has_single_unique || is_part_of_composite_unique; - - if *desired_is_unique == current_is_unique { - continue; - } - - let qualified_table = format!( - "\"{}\".\"{}\"", - quote_identifier(&table_schema), - quote_identifier(&table_name) - ); - require_safe_identifier(column_name, "column name")?; - let qualified_column = format!("\"{}\"", quote_identifier(column_name)); - - if *desired_is_unique { - // Add a new single-column UNIQUE constraint. - if current_is_unique { - continue; - } - - let generated_name = veloxdb_unique_constraint_name(&table_name, column_name); - - // If a constraint with that name exists and doesn't match our target column, fail fast. - if let Some(existing_cols) = unique_by_name.get(&generated_name) { - if existing_cols.len() != 1 || existing_cols[0] != *column_name { - return Err(format!( - "Cannot create UNIQUE constraint due to name collision ({}). Rename the existing constraint.", - generated_name - )); - } - } - - let sql = format!( - "ALTER TABLE {} ADD CONSTRAINT \"{}\" UNIQUE ({})", - qualified_table, - quote_identifier(&generated_name), - qualified_column - ); - txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - } else { - // Drop the existing single-column UNIQUE constraints for this column. - let constraint_names = single_unique_constraint_names_by_column - .get(column_name) - .cloned() - .unwrap_or_default(); - - for constraint_name in constraint_names { - require_safe_identifier(&constraint_name, "constraint name")?; - let sql = format!( - "ALTER TABLE {} DROP CONSTRAINT \"{}\"", - qualified_table, - quote_identifier(&constraint_name) - ); - txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - } - } - } - - txn.commit().await.map_err(|error| map_pg_err(error, None))?; - Ok(()) - }) - .await -} - -#[tauri::command] -pub async fn get_foreign_keys( - app: AppHandle, - state: State<'_, AppState>, - connection_id: Option, -) -> Result, String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; - - if engine == DatabaseEngine::Mysql { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let rows = sqlx::query( - " - select - kcu.table_schema as from_schema, - kcu.table_name as from_table, - kcu.column_name as from_column, - kcu.referenced_table_schema as to_schema, - kcu.referenced_table_name as to_table, - kcu.referenced_column_name as to_column - from information_schema.key_column_usage kcu - where kcu.referenced_table_name is not null - order by kcu.table_schema, kcu.table_name, kcu.ordinal_position - limit ? - ", - ) - .bind(MAX_FOREIGN_KEY_ROWS) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut edges = Vec::new(); - for row in rows { - edges.push(ForeignKeyEdge { - from_schema: mysql_get_string(&row, 0, "from_schema", "get_foreign_keys")?, - from_table: mysql_get_string(&row, 1, "from_table", "get_foreign_keys")?, - from_column: mysql_get_string(&row, 2, "from_column", "get_foreign_keys")?, - to_schema: mysql_get_string(&row, 3, "to_schema", "get_foreign_keys")?, - to_table: mysql_get_string(&row, 4, "to_table", "get_foreign_keys")?, - to_column: mysql_get_string(&row, 5, "to_column", "get_foreign_keys")?, - }); - } - return Ok(edges); - } - - if engine == DatabaseEngine::Sqlite { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - let tables = sqlx::query( - " - select name - from sqlite_master - where type = 'table' - and name not like 'sqlite_%' - ", - ) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let mut edges = Vec::new(); - for table in tables { - let table_name: String = sqlite_get_idx(&table, 0, "name", "get_foreign_keys")?; - require_safe_identifier(&table_name, "table name")?; - let fk_sql = format!("PRAGMA foreign_key_list(\"{}\");", quote_identifier(&table_name)); - let fk_rows = sqlx::query(&fk_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - for row in fk_rows { - edges.push(ForeignKeyEdge { - from_schema: "main".to_string(), - from_table: table_name.clone(), - from_column: sqlite_get_name(&row, "from", "get_foreign_keys")?, - to_schema: "main".to_string(), - to_table: sqlite_get_name(&row, "table", "get_foreign_keys")?, - to_column: sqlite_get_name(&row, "to", "get_foreign_keys")?, - }); - if edges.len() >= MAX_FOREIGN_KEY_ROWS as usize { - return Ok(edges); - } - } - } - return Ok(edges); - } - - with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { - let rows = client - .query( - " - select - src_ns.nspname::text as from_schema, - src_cls.relname::text as from_table, - src_att.attname::text as from_column, - tgt_ns.nspname::text as to_schema, - tgt_cls.relname::text as to_table, - tgt_att.attname::text as to_column - from pg_constraint c - join pg_class src_cls on src_cls.oid = c.conrelid - join pg_namespace src_ns on src_ns.oid = src_cls.relnamespace - join pg_class tgt_cls on tgt_cls.oid = c.confrelid - join pg_namespace tgt_ns on tgt_ns.oid = tgt_cls.relnamespace - cross join lateral unnest(c.conkey, c.confkey) as u(attnum, confattnum) - join pg_attribute src_att - on src_att.attrelid = c.conrelid - and src_att.attnum = u.attnum - and not src_att.attisdropped - join pg_attribute tgt_att - on tgt_att.attrelid = c.confrelid - and tgt_att.attnum = u.confattnum - and not tgt_att.attisdropped - where c.contype = 'f' - and src_ns.nspname not in ('pg_catalog', 'information_schema') - order by src_ns.nspname, src_cls.relname, c.conname, u.attnum - limit $1 - ", - &[&MAX_FOREIGN_KEY_ROWS], - ) - .await - .map_err(|error| error.to_string())?; - - Ok(rows - .into_iter() - .map(|row| ForeignKeyEdge { - from_schema: row.get(0), - from_table: row.get(1), - from_column: row.get(2), - to_schema: row.get(3), - to_table: row.get(4), - to_column: row.get(5), - }) - .collect()) - }) - .await -} - -#[tauri::command] -pub async fn get_table_indexes( - app: AppHandle, - state: State<'_, AppState>, - input: SchemaRequest, -) -> Result { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let ctx = input.clone(); - - if engine == DatabaseEngine::Mysql { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let fetch_limit = MAX_TABLE_INDEX_ROWS + 1; - let rows = sqlx::query( - " - select - table_schema as index_schema, - index_name, - table_schema, - table_name, - non_unique = 0 as is_unique, - index_name = 'PRIMARY' as is_primary, - true as is_valid, - false as is_partial, - concat(index_name, ' (', group_concat(column_name order by seq_in_index separator ', '), ')') as definition, - 0 as index_bytes, - 0 as idx_scan, - 0 as idx_tup_read, - 0 as idx_tup_fetch - from information_schema.statistics - where table_schema = ? - and table_name = ? - group by table_schema, table_name, index_name, non_unique - order by index_name - limit ? - ", - ) - .bind(&ctx.table_schema) - .bind(&ctx.table_name) - .bind(fetch_limit) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let truncated = rows.len() as i64 > MAX_TABLE_INDEX_ROWS; - let mut indexes = Vec::new(); - for row in rows.into_iter().take(MAX_TABLE_INDEX_ROWS as usize) { - indexes.push(IndexInfo { - index_schema: mysql_get_string(&row, 0, "index_schema", "get_table_indexes")?, - index_name: mysql_get_string(&row, 1, "index_name", "get_table_indexes")?, - table_schema: mysql_get_string(&row, 2, "table_schema", "get_table_indexes")?, - table_name: mysql_get_string(&row, 3, "table_name", "get_table_indexes")?, - is_unique: mysql_get_idx(&row, 4, "is_unique", "get_table_indexes")?, - is_primary: mysql_get_idx(&row, 5, "is_primary", "get_table_indexes")?, - is_valid: true, - is_partial: false, - definition: mysql_get_string(&row, 8, "definition", "get_table_indexes")?, - index_bytes: 0, - idx_scan: 0, - idx_tup_read: 0, - idx_tup_fetch: 0, - }); - } - return Ok(TableIndexesResult { indexes, truncated }); - } - - if engine == DatabaseEngine::Sqlite { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - require_safe_identifier(&ctx.table_name, "table name")?; - let pragma_sql = format!( - "PRAGMA index_list(\"{}\");", - quote_identifier(&ctx.table_name) - ); - let rows = sqlx::query(&pragma_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let truncated = rows.len() as i64 > MAX_TABLE_INDEX_ROWS; - let mut indexes = Vec::new(); - for row in rows.into_iter().take(MAX_TABLE_INDEX_ROWS as usize) { - let index_name: String = sqlite_get_name(&row, "name", "get_table_indexes")?; - require_safe_identifier(&index_name, "index name")?; - let index_info_sql = format!("PRAGMA index_info(\"{}\");", quote_identifier(&index_name)); - let index_info_rows = sqlx::query(&index_info_sql) - .fetch_all(&pool) - .await - .map_err(|error| error.to_string())?; - let index_columns = index_info_rows - .into_iter() - .filter_map(|idx| sqlite_get_name::(&idx, "name", "get_table_indexes").ok()) - .collect::>(); - indexes.push(IndexInfo { - index_schema: "main".to_string(), - index_name: index_name.clone(), - table_schema: "main".to_string(), - table_name: ctx.table_name.clone(), - is_unique: sqlite_get_name::(&row, "unique", "get_table_indexes")? == 1, - is_primary: sqlite_get_name::(&row, "origin", "get_table_indexes")? == "pk", - is_valid: true, - is_partial: sqlite_get_name::(&row, "partial", "get_table_indexes")? == 1, - definition: if index_columns.is_empty() { - format!("index {}", index_name) - } else { - format!("index {} ({})", index_name, index_columns.join(", ")) - }, - index_bytes: 0, - idx_scan: 0, - idx_tup_read: 0, - idx_tup_fetch: 0, - }); - } - return Ok(TableIndexesResult { indexes, truncated }); - } - - with_pool_client_retry(&app, &state, &connection_id, ctx, |client, input| async move { - let table_schema = input.table_schema; - let table_name = input.table_name; - let fetch_limit = MAX_TABLE_INDEX_ROWS + 1; - - let rows = client - .query( - " - select - ins.nspname::text as index_schema, - ic.relname::text as index_name, - tn.nspname::text as table_schema, - tc.relname::text as table_name, - i.indisunique as is_unique, - i.indisprimary as is_primary, - i.indisvalid as is_valid, - (i.indpred is not null) as is_partial, - pg_get_indexdef(i.indexrelid) as definition, - coalesce(pg_relation_size(i.indexrelid::regclass), 0)::bigint as index_bytes, - coalesce(s.idx_scan, 0)::bigint as idx_scan, - coalesce(s.idx_tup_read, 0)::bigint as idx_tup_read, - coalesce(s.idx_tup_fetch, 0)::bigint as idx_tup_fetch - from pg_index i - join pg_class ic on ic.oid = i.indexrelid - join pg_namespace ins on ins.oid = ic.relnamespace - join pg_class tc on tc.oid = i.indrelid - join pg_namespace tn on tn.oid = tc.relnamespace - left join pg_stat_user_indexes s on s.indexrelid = i.indexrelid - where tn.nspname = $1 - and tc.relname = $2 - and ins.nspname not in ('pg_catalog', 'information_schema') - order by ic.relname - limit $3 - ", - &[&table_schema, &table_name, &fetch_limit], - ) - .await - .map_err(|error| error.to_string())?; - - let truncated = rows.len() as i64 > MAX_TABLE_INDEX_ROWS; - let take = if truncated { - MAX_TABLE_INDEX_ROWS as usize - } else { - rows.len() - }; - - let mut indexes = Vec::with_capacity(take); - for row in rows.into_iter().take(take) { - indexes.push(IndexInfo { - index_schema: row.get(0), - index_name: row.get(1), - table_schema: row.get(2), - table_name: row.get(3), - is_unique: row.get(4), - is_primary: row.get(5), - is_valid: row.get(6), - is_partial: row.get(7), - definition: row.get(8), - index_bytes: row.get(9), - idx_scan: row.get(10), - idx_tup_read: row.get(11), - idx_tup_fetch: row.get(12), - }); - } - - Ok(TableIndexesResult { indexes, truncated }) - }) - .await -} - -#[tauri::command] -pub async fn execute_ddl_transaction( - app: AppHandle, - state: State<'_, AppState>, - input: DdlBatchRequest, -) -> Result<(), String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, input, |mut client, input| async move { - let stmts: Vec = input - .statements - .into_iter() - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) - .collect(); - - if stmts.is_empty() { - return Err("No SQL statements to execute.".to_string()); - } - - let txn = client.transaction().await.map_err(|error| map_pg_err(error, None))?; - for sql in &stmts { - txn.execute(sql.as_str(), &[]) - .await - .map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - } - txn.commit().await.map_err(|error| map_pg_err(error, None))?; - Ok(()) - }) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - let mut tx = pool.begin().await.map_err(|error| error.to_string())?; - for sql in input - .statements - .iter() - .map(|s| s.trim()) - .filter(|s| !s.is_empty()) - { - sqlx::query(sql) - .execute(&mut *tx) - .await - .map_err(|error| error.to_string())?; - } - tx.commit().await.map_err(|error| error.to_string()) - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - let mut tx = pool.begin().await.map_err(|error| error.to_string())?; - for sql in input - .statements - .iter() - .map(|s| s.trim()) - .filter(|s| !s.is_empty()) - { - sqlx::query(sql) - .execute(&mut *tx) - .await - .map_err(|error| error.to_string())?; - } - tx.commit().await.map_err(|error| error.to_string()) - } - } -} - -/// Run a single DDL statement outside an explicit transaction (required for `CREATE INDEX CONCURRENTLY`). -#[tauri::command] -pub async fn execute_ddl_statement( - app: AppHandle, - state: State<'_, AppState>, - input: DdlStatementRequest, -) -> Result<(), String> { - let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; - let sql = input.statement.trim().to_string(); - if sql.is_empty() { - return Err("No SQL statement to execute.".to_string()); - } - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(&app, &state, &connection_id, sql, |client, sql| async move { - client - .execute(sql.as_str(), &[]) - .await - .map_err(|error| map_pg_err(error, Some(sql.as_str())))?; - Ok(()) - }) - .await - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; - sqlx::query(&sql) - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - Ok(()) - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; - sqlx::query(&sql) - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - Ok(()) - } - } -} - -#[tauri::command] -pub async fn export_diagram_png( - input: DiagramExportRequest, - output_path: String, -) -> Result<(), String> { - let path = std::path::PathBuf::from(&output_path); - tokio::task::spawn_blocking(move || export_diagram_to_png(&input, &path)) - .await - .map_err(|e| e.to_string())? -} - -#[tauri::command] -pub async fn export_results_csv_command( - app: AppHandle, - state: State<'_, AppState>, - input: ExportQueryRequest, -) -> Result<(), String> { - export_results_csv(&app, &state, &input).await -} - -#[tauri::command] -pub async fn export_results_json_command( - app: AppHandle, - state: State<'_, AppState>, - input: ExportQueryRequest, -) -> Result<(), String> { - export_results_json(&app, &state, &input).await -} - -#[tauri::command] -pub async fn save_base64_png(data: String, output_path: String) -> Result<(), String> { - use base64::Engine; - let bytes = base64::engine::general_purpose::STANDARD - .decode(data.strip_prefix("data:image/png;base64,").unwrap_or(&data)) - .map_err(|e| e.to_string())?; - std::fs::write(&output_path, bytes).map_err(|e| e.to_string()) -} - -#[tauri::command] -pub async fn save_text_file(content: String, output_path: String) -> Result<(), String> { - std::fs::write(&output_path, content).map_err(|e| e.to_string()) -} - -#[tauri::command] -pub async fn store_openrouter_api_key(api_key: String) -> Result<(), String> { - if api_key.trim().is_empty() { - return credentials::delete_openrouter_api_key(); - } - credentials::store_openrouter_api_key(&api_key) -} - -#[tauri::command] -pub async fn get_openrouter_api_key() -> Result, String> { - credentials::get_openrouter_api_key() -} - -#[tauri::command] -pub async fn delete_openrouter_api_key() -> Result<(), String> { - credentials::delete_openrouter_api_key() -} - -#[cfg(test)] -mod tests { - use super::{ - build_schema_context, classify_sql_intent, database_name_from_mysql_value, - decode_mysql_bytes_as_string, extract_openrouter_stream_delta, mysql_decode_error, - parse_ask_veloxy_json, sqlite_decode_error, streaming_display_text, - validate_generated_sql, - }; - use crate::models::{ - AskVeloxyDbContextCache, DatabaseEngine, QueryEditorColumn, QueryEditorMetadata, - QueryEditorTable, - }; - - #[test] - fn streaming_display_text_extracts_partial_json_message() { - let partial = r#"{ "message": "The messages table has relationships with:\n- delivery_reports"#; - let display = streaming_display_text(partial); - assert!(display.contains("messages table")); - assert!(display.contains("delivery_reports")); - } - - #[test] - fn streaming_display_text_returns_plain_text_directly() { - assert_eq!( - streaming_display_text("Hello from Veloxy"), - "Hello from Veloxy" - ); - } - - #[test] - fn extract_openrouter_stream_delta_reads_content() { - let data = r#"{"choices":[{"delta":{"content":"Hello"}}]}"#; - assert_eq!( - extract_openrouter_stream_delta(data).as_deref(), - Some("Hello") - ); - } - - #[test] - fn database_name_from_mysql_value_rejects_empty() { - assert!(database_name_from_mysql_value(None, "list_databases").is_err()); - assert!(database_name_from_mysql_value(Some(String::new()), "list_databases").is_err()); - } - - #[test] - fn database_name_from_mysql_value_accepts_non_empty() { - let name = - database_name_from_mysql_value(Some("my_app".to_string()), "list_databases").expect("name"); - assert_eq!(name, "my_app"); - } - - #[test] - fn decode_mysql_bytes_as_string_uses_utf8_text() { - assert_eq!( - decode_mysql_bytes_as_string(b"my_schema"), - "my_schema" - ); - } - - #[test] - fn mysql_decode_error_is_explicit() { - let message = mysql_decode_error("get_tables", "table_schema", Some(0), "mismatched types"); - assert!(message.contains("MySQL decode error")); - assert!(message.contains("get_tables")); - assert!(message.contains("table_schema")); - } - - #[test] - fn sqlite_decode_error_is_explicit() { - let message = sqlite_decode_error("get_schema", "name", Some(0), "unsupported value type"); - assert!(message.contains("SQLite decode error")); - assert!(message.contains("get_schema")); - assert!(message.contains("name")); - } - - #[test] - fn schema_context_is_bounded() { - let columns = (0..40) - .map(|idx| QueryEditorColumn { - name: format!("column_{}", idx), - data_type: "text".to_string(), - }) - .collect::>(); - let tables = (0..20) - .map(|idx| QueryEditorTable { - schema: "public".to_string(), - name: format!("events_{}", idx), - columns: columns.clone(), - }) - .collect::>(); - let metadata = QueryEditorMetadata { - tables, - functions: Vec::new(), - truncated_tables: false, - truncated_columns: false, - truncated_functions: false, - }; - let db_context = AskVeloxyDbContextCache { - database_name: "test".to_string(), - engine: DatabaseEngine::Postgres, - metadata, - foreign_keys: Vec::new(), - }; - - let context = build_schema_context(&db_context, "show events", None); - assert!(!context.is_empty()); - assert!(context.len() <= super::ASK_VELOXY_SCHEMA_CHAR_BUDGET); - } - - #[test] - fn ask_veloxy_json_parser_handles_embedded_block() { - let content = "Here is the output {\"sql\":\"select 1\",\"intent\":\"select\",\"confidence\":0.9,\"warnings\":[]}"; - let parsed = parse_ask_veloxy_json(content).expect("json should parse"); - assert_eq!(parsed.get("sql").and_then(|v| v.as_str()), Some("select 1")); - } - - #[test] - fn sql_validation_rejects_multi_statement() { - let multi = "select 1; select 2;"; - assert!(validate_generated_sql(multi).is_err()); - } - - #[test] - fn sql_intent_classifier_recognizes_update() { - assert_eq!(classify_sql_intent("UPDATE foo SET bar = 1"), "update"); - } - - #[test] - fn read_only_check_allows_selects_and_explain() { - assert!(super::is_read_only_sql("SELECT 1")); - assert!(super::is_read_only_sql("EXPLAIN ANALYZE SELECT * FROM t")); - assert!(super::is_read_only_sql("WITH x AS (SELECT 1) SELECT * FROM x")); - assert!(super::is_read_only_sql("BEGIN; SELECT 1; COMMIT;")); - } - - #[test] - fn read_only_check_blocks_writes() { - assert!(!super::is_read_only_sql("DELETE FROM t")); - assert!(!super::is_read_only_sql("DROP TABLE t")); - assert!(!super::is_read_only_sql("BEGIN; UPDATE t SET a = 1; COMMIT;")); - assert!(!super::is_read_only_sql("SELECT 1; DELETE FROM t")); - assert!(!super::is_read_only_sql("")); - } - - #[test] - fn mysql_timestamp_formats_as_datetime_string() { - let dt = chrono::DateTime::parse_from_rfc3339("2024-03-15T10:30:45Z") - .unwrap() - .with_timezone(&chrono::Utc); - assert_eq!(dt.format("%Y-%m-%d %H:%M:%S").to_string(), "2024-03-15 10:30:45"); - } - - #[test] - fn mysql_datetime_formats_as_naive_datetime_string() { - let dt = chrono::NaiveDateTime::parse_from_str("2024-03-15 10:30:45", "%Y-%m-%d %H:%M:%S").unwrap(); - assert_eq!(dt.format("%Y-%m-%d %H:%M:%S").to_string(), "2024-03-15 10:30:45"); - } - - #[test] - fn mysql_date_formats_as_iso_date() { - let d = chrono::NaiveDate::from_ymd_opt(2024, 3, 15).unwrap(); - assert_eq!(d.to_string(), "2024-03-15"); - } - - #[test] - fn mysql_time_formats_as_iso_time() { - let t = chrono::NaiveTime::from_hms_opt(10, 30, 45).unwrap(); - assert_eq!(t.to_string(), "10:30:45"); - } -} diff --git a/src-tauri/src/commands/connections.rs b/src-tauri/src/commands/connections.rs new file mode 100644 index 0000000..2416fb3 --- /dev/null +++ b/src-tauri/src/commands/connections.rs @@ -0,0 +1,218 @@ +use uuid::Uuid; +use tauri::{AppHandle, State}; + +use crate::db::{ + build_mysql_pool, build_mysql_pool_custom, + build_pool, build_pool_custom, disconnect_connection, + drop_pool, get_or_create_duckdb_connection, get_or_create_mongo_client, + get_or_create_redis_client, + load_connection, persist_connection_with_password, refresh_connection_pools, + resolve_connection_engine, AppState, DEFAULT_MYSQL_PORT, +}; +use mongodb::bson::doc; +use crate::credentials; +use crate::engines; +use crate::engines::DatabaseEngineOps; +use crate::models::{ + ConnectionInput, ConnectionSummary, DatabaseEngine, DatabaseInfo, + StoredConnection, SwitchDatabaseRequest, +}; +use crate::pg_error::map_pg_err; +use crate::ssh_tunnel::SshTunnel; + +#[tauri::command] +pub async fn connect_db( + app: AppHandle, + state: State<'_, AppState>, + input: ConnectionInput, +) -> Result { + let connection_id = input.id.clone().unwrap_or_else(|| Uuid::new_v4().to_string()); + + let ops = engines::get_engine(input.engine); + ops.connect(&app, &state, &input, &connection_id).await.map_err(String::from)?; + + let stored_connection = StoredConnection::from_input(connection_id.clone(), input.clone()); + persist_connection_with_password(&app, &stored_connection, &input.password)?; + + *state.active_connection_id.write().await = Some(connection_id); + + Ok(stored_connection.summary()) +} + +#[tauri::command] +pub async fn list_connections_command(app: AppHandle) -> Result, String> { + crate::db::list_connections(&app).map_err(String::from) +} + +#[tauri::command] +pub async fn set_active_connection( + app: AppHandle, + state: State<'_, AppState>, + connection_id: String, +) -> Result { + let stored_connection = load_connection(&app, &connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + + let ops = engines::get_engine(stored_connection.engine); + ops.ping(&app, &state, &connection_id).await.map_err(String::from)?; + + *state.active_connection_id.write().await = Some(connection_id); + Ok(stored_connection.summary()) +} + +#[tauri::command] +pub async fn ping_connection( + app: AppHandle, + state: State<'_, AppState>, + connection_id: String, +) -> Result<(), String> { + let stored_connection = load_connection(&app, &connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let ops = engines::get_engine(stored_connection.engine); + ops.ping(&app, &state, &connection_id).await.map_err(String::from) +} + +#[tauri::command] +pub async fn refresh_connection( + app: AppHandle, + state: State<'_, AppState>, + connection_id: String, +) -> Result<(), String> { + refresh_connection_pools(&app, &state, &connection_id).await.map_err(String::from) +} + +#[tauri::command] +pub async fn disconnect_db( + state: State<'_, AppState>, + connection_id: String, +) -> Result<(), String> { + disconnect_connection(&state, &connection_id).await; + Ok(()) +} + +#[tauri::command] +pub async fn rename_connection( + app: AppHandle, + connection_id: String, + new_name: String, +) -> Result { + crate::db::rename_connection_in_store(&app, &connection_id, &new_name).map_err(String::from) +} + +#[tauri::command] +pub async fn delete_connection( + app: AppHandle, + state: State<'_, AppState>, + connection_id: String, +) -> Result<(), String> { + disconnect_connection(&state, &connection_id).await; + if let Err(e) = credentials::delete_password(&connection_id) { + log::warn!("Failed to delete keychain entry for {}: {}", connection_id, e); + } + crate::db::delete_connection_from_store(&app, &connection_id)?; + Ok(()) +} + +#[tauri::command] +pub async fn list_databases( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; + + let ops = engines::get_engine(engine); + ops.list_databases(&app, &state, &connection_id).await.map_err(String::from) +} + +#[tauri::command] +pub async fn switch_database( + app: AppHandle, + state: State<'_, AppState>, + input: SwitchDatabaseRequest, +) -> Result { + let mut stored_connection = load_connection(&app, &input.connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + + if stored_connection.engine == DatabaseEngine::Sqlite { + return Err("Switch database is not supported for SQLite connections.".to_string()); + } + + drop_pool(&state, &input.connection_id).await; + + stored_connection.database = input.database.clone(); + stored_connection.connected_at = crate::models::timestamp_string(); + persist_connection_with_password(&app, &stored_connection, &stored_connection.password.clone().unwrap_or_default())?; + + let connection_input = stored_connection.to_input(); + + match connection_input.engine { + DatabaseEngine::Postgres => { + let pool = if let Some(ref ssh_config) = connection_input.ssh_config { + if ssh_config.is_active() { + let tunnel = match SshTunnel::connect(ssh_config, &connection_input.host, connection_input.port).await { + Ok(tunnel) => tunnel, + Err(e) => return Err(format!("SSH tunnel failed: {}", e)), + }; + let local_port = tunnel.local_port; + state.ssh_tunnels.write().await.insert(input.connection_id.clone(), tunnel); + build_pool_custom("127.0.0.1", local_port, &connection_input)? + } else { + build_pool(&connection_input)? + } + } else { + build_pool(&connection_input)? + }; + + let client = match pool.get().await { + Ok(client) => client, + Err(e) => { drop_pool(&state, &input.connection_id).await; return Err(e.to_string()); } + }; + + if let Err(e) = client.simple_query("select 1").await { + drop_pool(&state, &input.connection_id).await; + return Err(map_pg_err(e, None)); + } + + state.pools.write().await.insert(input.connection_id.clone(), pool); + } + DatabaseEngine::Mysql => { + let pool = if let Some(ref ssh_config) = connection_input.ssh_config { + if ssh_config.is_active() { + let remote_port = if connection_input.port == 0 { DEFAULT_MYSQL_PORT } else { connection_input.port }; + let tunnel = match SshTunnel::connect(ssh_config, &connection_input.host, remote_port).await { + Ok(tunnel) => tunnel, + Err(e) => return Err(format!("SSH tunnel failed: {}", e)), + }; + let local_port = tunnel.local_port; + state.ssh_tunnels.write().await.insert(input.connection_id.clone(), tunnel); + build_mysql_pool_custom("127.0.0.1", local_port, &connection_input).await? + } else { + build_mysql_pool(&connection_input).await? + } + } else { + build_mysql_pool(&connection_input).await? + }; + + sqlx::query("select 1").execute(&pool).await.map_err(|e| e.to_string())?; + state.mysql_pools.write().await.insert(input.connection_id.clone(), pool); + } + DatabaseEngine::Sqlite => {} + DatabaseEngine::Mongo => { + let client = get_or_create_mongo_client(&app, &state, &input.connection_id).await?; + client.database(&input.database).run_command(doc! { "ping": 1 }).await + .map_err(|e| format!("MongoDB ping failed: {}", e))?; + } + DatabaseEngine::Duckdb => { + get_or_create_duckdb_connection(&app, &state, &input.connection_id).await?; + } + DatabaseEngine::Redis => { + let mut client = get_or_create_redis_client(&app, &state, &input.connection_id).await?; + redis::cmd("PING").query_async::<_, String>(&mut client).await + .map_err(|e| format!("Redis ping failed: {}", e))?; + } + } + + *state.active_connection_id.write().await = Some(input.connection_id); + Ok(stored_connection.summary()) +} diff --git a/src-tauri/src/commands/ddl.rs b/src-tauri/src/commands/ddl.rs new file mode 100644 index 0000000..2b5adbb --- /dev/null +++ b/src-tauri/src/commands/ddl.rs @@ -0,0 +1,113 @@ +use tauri::{AppHandle, State}; + +use crate::db::{ + get_or_create_duckdb_connection, get_or_create_mysql_pool, get_or_create_sqlite_pool, + resolve_connection_engine, with_pool_client_retry, AppState, +}; +use crate::models::{DatabaseEngine, DdlBatchRequest, DdlStatementRequest}; +use crate::pg_error::map_pg_err; + +#[tauri::command] +pub async fn execute_ddl_transaction( + app: AppHandle, + state: State<'_, AppState>, + input: DdlBatchRequest, +) -> Result<(), String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + + match engine { + DatabaseEngine::Postgres => { + with_pool_client_retry(&app, &state, &connection_id, input, |mut client, input| async move { + let stmts: Vec = input.statements.into_iter() + .map(|s| s.trim().to_string()).filter(|s| !s.is_empty()).collect(); + + if stmts.is_empty() { + return Err("No SQL statements to execute.".to_string()); + } + + let txn = client.transaction().await.map_err(|error| map_pg_err(error, None))?; + for sql in &stmts { + txn.execute(sql.as_str(), &[]).await + .map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + } + txn.commit().await.map_err(|error| map_pg_err(error, None))?; + Ok(()) + }).await.map_err(String::from) + } + DatabaseEngine::Mysql => { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + let mut tx = pool.begin().await.map_err(|error| error.to_string())?; + for sql in input.statements.iter().map(|s| s.trim()).filter(|s| !s.is_empty()) { + sqlx::query(sql).execute(&mut *tx).await.map_err(|error| error.to_string())?; + } + tx.commit().await.map_err(|error| error.to_string()) + } + DatabaseEngine::Sqlite => { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + let mut tx = pool.begin().await.map_err(|error| error.to_string())?; + for sql in input.statements.iter().map(|s| s.trim()).filter(|s| !s.is_empty()) { + sqlx::query(sql).execute(&mut *tx).await.map_err(|error| error.to_string())?; + } + tx.commit().await.map_err(|error| error.to_string()) + } + DatabaseEngine::Mongo => { + Err("MongoDB does not support DDL transactions.".to_string()) + } + DatabaseEngine::Duckdb => { + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(&connection_id).ok_or("DuckDB connection not found")?; + let conn = conn_mutex.lock().await; + for sql in input.statements.iter().map(|s| s.trim()).filter(|s| !s.is_empty()) { + conn.execute(sql, []).map_err(|e| format!("DuckDB DDL failed: {}", e))?; + } + Ok(()) + } + DatabaseEngine::Redis => Err("Not supported for Redis.".to_string()), + } +} + +#[tauri::command] +pub async fn execute_ddl_statement( + app: AppHandle, + state: State<'_, AppState>, + input: DdlStatementRequest, +) -> Result<(), String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let sql = input.statement.trim().to_string(); + if sql.is_empty() { + return Err("No SQL statement to execute.".to_string()); + } + + match engine { + DatabaseEngine::Postgres => { + with_pool_client_retry(&app, &state, &connection_id, sql, |client, sql| async move { + client.execute(sql.as_str(), &[]).await + .map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + Ok(()) + }).await.map_err(String::from) + } + DatabaseEngine::Mysql => { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + sqlx::query(&sql).execute(&pool).await.map_err(|error| error.to_string())?; + Ok(()) + } + DatabaseEngine::Sqlite => { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + sqlx::query(&sql).execute(&pool).await.map_err(|error| error.to_string())?; + Ok(()) + } + DatabaseEngine::Mongo => { + Err("MongoDB does not support DDL statements.".to_string()) + } + DatabaseEngine::Duckdb => { + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(&connection_id).ok_or("DuckDB connection not found")?; + let conn = conn_mutex.lock().await; + conn.execute(&sql, []).map_err(|e| format!("DuckDB DDL failed: {}", e))?; + Ok(()) + } + DatabaseEngine::Redis => Err("Not supported for Redis.".to_string()), + } +} diff --git a/src-tauri/src/commands/duckdb.rs b/src-tauri/src/commands/duckdb.rs new file mode 100644 index 0000000..ef35db9 --- /dev/null +++ b/src-tauri/src/commands/duckdb.rs @@ -0,0 +1,520 @@ +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::time::Instant; + +use tauri::{AppHandle, State}; + +use crate::db::{get_or_create_duckdb_connection, quote_identifier, resolve_connection_engine, AppState, MAX_QUERY_ROWS}; +use crate::models::{ColumnInfo, ColumnProperties, ForeignKeyEdge, IndexInfo, QueryRequest, QueryResult, SchemaRequest, TableIndexesResult, TableInfo, TablePropertiesApplyRequest}; + +/// Execute a SQL query against a DuckDB connection. +#[tauri::command] +pub async fn duckdb_run_query( + app: AppHandle, + state: State<'_, AppState>, + input: QueryRequest, +) -> Result { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let sql = input.sql.trim().to_string(); + if sql.is_empty() { + return Err("Enter a SQL statement before running the query.".to_string()); + } + + let started_at = Instant::now(); + let max_rows = input.max_rows.unwrap_or(MAX_QUERY_ROWS); + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + let stmt = conn + .prepare(&sql) + .map_err(|e| format!("DuckDB query preparation failed: {}", e))?; + + let columns: Vec = stmt + .column_names() + .iter() + .map(|c| c.to_string()) + .collect(); + + let col_count = columns.len(); + + if col_count > 0 { + // Build a wrapper query that casts all columns to VARCHAR for safe string extraction + let cast_cols: Vec = columns + .iter() + .map(|c| { + let quoted = quote_identifier(c); + format!("CAST(\"{quoted}\" AS VARCHAR) as \"{quoted}\"") + }) + .collect(); + let wrapper_sql = format!("SELECT {} FROM ({})", cast_cols.join(", "), sql); + + drop(stmt); + + let mut stmt2 = conn + .prepare(&wrapper_sql) + .map_err(|e| format!("DuckDB query preparation failed: {}", e))?; + + let mut rows: Vec>> = Vec::new(); + let mut total = 0usize; + + let row_iter = stmt2 + .query_map([], |row| { + let mut map = BTreeMap::new(); + for (i, col) in columns.iter().enumerate() { + let val: Option = row.get(i).ok().flatten(); + map.insert(col.clone(), val); + } + Ok(map) + }) + .map_err(|e| format!("DuckDB query execution failed: {}", e))?; + + for row_result in row_iter { + match row_result { + Ok(row) => { + if rows.len() < max_rows { + rows.push(row); + } + total += 1; + } + Err(e) => return Err(format!("DuckDB row error: {}", e)), + } + } + + Ok(QueryResult { + columns, + row_count: rows.len(), + rows, + execution_ms: started_at.elapsed().as_millis(), + truncated: total > max_rows, + command_tag: None, + }) + } else { + // Non-row-returning statement (INSERT, UPDATE, DDL) + drop(stmt); + let affected = conn + .execute(&sql, []) + .map_err(|e| format!("DuckDB execution failed: {}", e))?; + + Ok(QueryResult { + columns: Vec::new(), + rows: Vec::new(), + row_count: affected, + execution_ms: started_at.elapsed().as_millis(), + truncated: false, + command_tag: Some(affected as u64), + }) + } +} + +/// List all tables in the DuckDB database. +#[tauri::command] +pub async fn duckdb_get_tables( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + let mut stmt = conn + .prepare( + "SELECT table_name FROM information_schema.tables \ + WHERE table_schema = 'main' AND table_type = 'BASE TABLE' \ + ORDER BY table_name", + ) + .map_err(|e| format!("DuckDB table listing failed: {}", e))?; + + let tables: Vec = stmt + .query_map([], |row| { + let name: String = row.get(0)?; + Ok(TableInfo { + schema: "main".to_string(), + name: name.clone(), + preview_query: format!( + "SELECT * FROM \"{}\" LIMIT 100;", + quote_identifier(&name) + ), + }) + }) + .map_err(|e| format!("DuckDB table listing failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + Ok(tables) +} + +/// Get the column schema for a DuckDB table. +#[tauri::command] +pub async fn duckdb_get_schema( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, + table_schema: String, + table_name: String, +) -> Result, String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + let mut stmt = conn + .prepare( + "SELECT column_name, data_type, is_nullable \ + FROM information_schema.columns \ + WHERE table_schema = ? AND table_name = ? \ + ORDER BY ordinal_position", + ) + .map_err(|e| format!("DuckDB schema query failed: {}", e))?; + + let columns: Vec = stmt + .query_map( + duckdb::params![table_schema, table_name], + |row| { + let col_name: String = row.get(0)?; + let data_type: String = row.get(1)?; + let is_nullable: String = row.get(2)?; + Ok(ColumnInfo { + table_schema: table_schema.clone(), + table_name: table_name.clone(), + column_name: col_name, + data_type, + is_nullable: is_nullable == "YES", + }) + }, + ) + .map_err(|e| format!("DuckDB schema query failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + Ok(columns) +} + +/// List foreign key relationships in the DuckDB database. +#[tauri::command] +pub async fn duckdb_get_foreign_keys( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + let mut stmt = conn + .prepare( + "SELECT + kcu.table_schema AS from_schema, + kcu.table_name AS from_table, + kcu.column_name AS from_column, + kcu.referenced_table_schema AS to_schema, + kcu.referenced_table_name AS to_table, + kcu.referenced_column_name AS to_column + FROM information_schema.key_column_usage kcu + WHERE kcu.referenced_table_name IS NOT NULL + ORDER BY kcu.table_schema, kcu.table_name", + ) + .map_err(|e| format!("DuckDB FK query failed: {}", e))?; + + let edges: Vec = stmt + .query_map([], |row| { + Ok(ForeignKeyEdge { + from_schema: row.get(0)?, + from_table: row.get(1)?, + from_column: row.get(2)?, + to_schema: row.get(3)?, + to_table: row.get(4)?, + to_column: row.get(5)?, + }) + }) + .map_err(|e| format!("DuckDB FK query failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + Ok(edges) +} + +/// List indexes on a DuckDB table. +#[tauri::command] +pub async fn duckdb_get_table_indexes( + app: AppHandle, + state: State<'_, AppState>, + input: SchemaRequest, +) -> Result { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + let mut stmt = conn + .prepare( + "SELECT + index_name, + index_schema, + table_name, + is_unique, + is_primary, + sql + FROM duckdb_indexes() + WHERE table_name = ? AND schema_name = ? + ORDER BY index_name", + ) + .map_err(|e| format!("DuckDB index query failed: {}", e))?; + + let indexes: Vec = stmt + .query_map( + duckdb::params![input.table_name, input.table_schema], + |row| { + let idx_name: String = row.get(0)?; + let idx_schema: String = row.get(1)?; + let tbl_name: String = row.get(2)?; + let is_unique: bool = row.get(3)?; + let is_primary: bool = row.get(4)?; + let definition: String = row.get(5)?; + Ok(IndexInfo { + index_schema: idx_schema, + index_name: idx_name, + table_schema: input.table_schema.clone(), + table_name: tbl_name, + is_unique, + is_primary, + is_valid: true, + is_partial: false, + definition, + index_bytes: 0, + idx_scan: 0, + idx_tup_read: 0, + idx_tup_fetch: 0, + }) + }, + ) + .map_err(|e| format!("DuckDB index query failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + let truncated = indexes.len() > 500; + Ok(TableIndexesResult { + indexes: if truncated { indexes.into_iter().take(500).collect() } else { indexes }, + truncated, + }) +} + +/// Get column properties (nullable, primary key, unique, default) for a DuckDB table. +#[tauri::command] +pub async fn duckdb_get_table_properties( + app: AppHandle, + state: State<'_, AppState>, + input: SchemaRequest, +) -> Result, String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + // Query columns + let mut col_stmt = conn + .prepare( + "SELECT table_schema, table_name, column_name, data_type, is_nullable, column_default + FROM information_schema.columns + WHERE table_schema = ? AND table_name = ? + ORDER BY ordinal_position", + ) + .map_err(|e| format!("DuckDB properties query failed: {}", e))?; + + let columns: Vec<(String, String, String, String, String, Option)> = col_stmt + .query_map( + duckdb::params![input.table_schema, input.table_name], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?, row.get(4)?, row.get(5)?)), + ) + .map_err(|e| format!("DuckDB properties query failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + // Query primary keys + let mut pk_stmt = conn + .prepare( + "SELECT kcu.column_name + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + WHERE tc.table_schema = ? AND tc.table_name = ? + AND tc.constraint_type = 'PRIMARY KEY' + ORDER BY kcu.ordinal_position", + ) + .map_err(|e| format!("DuckDB PK query failed: {}", e))?; + + let pk_cols: std::collections::HashSet = pk_stmt + .query_map(duckdb::params![input.table_schema, input.table_name], |row| row.get(0)) + .map_err(|e| format!("DuckDB PK query failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + // Query unique constraints + let mut uniq_stmt = conn + .prepare( + "SELECT tc.constraint_name, kcu.column_name + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + WHERE tc.table_schema = ? AND tc.table_name = ? + AND tc.constraint_type = 'UNIQUE' + ORDER BY tc.constraint_name, kcu.ordinal_position", + ) + .map_err(|e| format!("DuckDB unique query failed: {}", e))?; + + let mut unique_by_name: HashMap> = HashMap::new(); + let unique_rows: Vec<(String, String)> = uniq_stmt + .query_map(duckdb::params![input.table_schema, input.table_name], |row| Ok((row.get(0)?, row.get(1)?))) + .map_err(|e| format!("DuckDB unique query failed: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + for (constraint_name, column_name) in unique_rows { + unique_by_name.entry(constraint_name).or_default().push(column_name); + } + + let mut unique_columns: HashSet = HashSet::new(); + let mut composite_unique_columns: HashSet = HashSet::new(); + for cols in unique_by_name.values() { + for c in cols { unique_columns.insert(c.clone()); } + if cols.len() > 1 { for c in cols { composite_unique_columns.insert(c.clone()); } } + } + + Ok(columns.into_iter().map(|(table_schema, table_name, column_name, data_type, is_nullable, column_default)| { + let is_pk = pk_cols.contains(&column_name); + ColumnProperties { + table_schema, + table_name, + column_name: column_name.clone(), + data_type, + is_nullable: is_nullable == "YES", + is_primary_key: is_pk, + is_unique: is_pk || unique_columns.contains(&column_name), + is_part_of_composite_unique: composite_unique_columns.contains(&column_name), + column_default, + is_identity: false, + identity_generation: None, + is_generated: None, + } + }).collect()) +} + +/// Apply nullable/unique changes to a DuckDB table via ALTER TABLE. +#[tauri::command] +pub async fn duckdb_apply_table_properties( + app: AppHandle, + state: State<'_, AppState>, + input: TablePropertiesApplyRequest, +) -> Result<(), String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + + get_or_create_duckdb_connection(&app, &state, &connection_id).await?; + + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or_else(|| "DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + + let table_schema = &input.table_schema; + let table_name = &input.table_name; + + // Get current nullability + let mut col_stmt = conn + .prepare( + "SELECT column_name, is_nullable FROM information_schema.columns + WHERE table_schema = ? AND table_name = ?", + ) + .map_err(|e| format!("DuckDB: {}", e))?; + let current_nullable: HashMap = col_stmt + .query_map(duckdb::params![table_schema, table_name], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)? == "YES")) + }) + .map_err(|e| format!("DuckDB: {}", e))? + .filter_map(|r| r.ok()) + .collect(); + + let qualified_table = format!( + "\"{}\".\"{}\"", + quote_identifier(table_schema), + quote_identifier(table_name) + ); + + for update in &input.columns { + let current = current_nullable.get(&update.column_name) + .ok_or_else(|| format!("Unknown column: {}", update.column_name))?; + if *current == update.is_nullable { continue; } + let qualified_col = format!("\"{}\"", quote_identifier(&update.column_name)); + if update.is_nullable { + conn.execute( + &format!("ALTER TABLE {} ALTER {} DROP NOT NULL", qualified_table, qualified_col), + [], + ).map_err(|e| format!("DuckDB ALTER failed: {}", e))?; + } else { + conn.execute( + &format!("ALTER TABLE {} ALTER {} SET NOT NULL", qualified_table, qualified_col), + [], + ).map_err(|e| format!("DuckDB ALTER failed: {}", e))?; + } + } + + // Unique constraints + for update in &input.columns { + if !update.is_unique { continue; } + let qualified_col = format!("\"{}\"", quote_identifier(&update.column_name)); + let constraint_name = format!("veloxdb_unq_{}", update.column_name); + // Add if not exists — best-effort (DuckDB may error on duplicate, which is fine) + conn.execute( + &format!( + "ALTER TABLE {} ADD CONSTRAINT \"{}\" UNIQUE ({})", + qualified_table, + quote_identifier(&constraint_name), + qualified_col + ), + [], + ).ok(); // Ignore if constraint already exists + } + + Ok(()) +} diff --git a/src-tauri/src/commands/editor_meta.rs b/src-tauri/src/commands/editor_meta.rs new file mode 100644 index 0000000..8c770c0 --- /dev/null +++ b/src-tauri/src/commands/editor_meta.rs @@ -0,0 +1,426 @@ +use tauri::{AppHandle, State}; + +use crate::db::{ + get_or_create_mysql_pool, get_or_create_sqlite_pool, load_connection, quote_identifier, + require_safe_identifier, resolve_connection_engine, with_pool_client_retry, AppState, +}; +use crate::models::{ + DatabaseEngine, ForeignKeyEdge, QueryEditorColumn, QueryEditorFunction, + QueryEditorMetadata, QueryEditorTable, +}; +use crate::pg_error::map_pg_err; + +use super::{ + MAX_EDITOR_TABLES, MAX_EDITOR_COLUMNS_PER_TABLE, MAX_EDITOR_FUNCTIONS, MAX_FOREIGN_KEY_ROWS, + mysql_get_string, sqlite_get_idx, sqlite_get_name, +}; + +pub(crate) async fn fetch_query_editor_metadata_for_connection( + app: &AppHandle, + state: &AppState, + connection_id: &str, + engine: DatabaseEngine, +) -> Result { + if engine == DatabaseEngine::Mysql { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + let database = load_connection(app, connection_id)? + .map(|connection| connection.database).unwrap_or_default(); + let table_rows = sqlx::query( + "select table_schema, table_name \ + from information_schema.tables \ + where table_type = 'BASE TABLE' \ + and table_schema = ? \ + and table_schema not in ('information_schema', 'mysql', 'performance_schema', 'sys') \ + order by table_schema, table_name \ + limit ?", + ).bind(&database).bind(MAX_EDITOR_TABLES + 1) + .fetch_all(&pool).await.map_err(|error| error.to_string())?; + + let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; + let mut tables = Vec::new(); + let mut truncated_columns = false; + + for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { + let schema: String = mysql_get_string(&row, 0, "table_schema", "get_query_editor_metadata")?; + let name: String = mysql_get_string(&row, 1, "table_name", "get_query_editor_metadata")?; + let column_rows = sqlx::query( + "select column_name, data_type \ + from information_schema.columns \ + where table_schema = ? and table_name = ? \ + order by ordinal_position \ + limit ?", + ).bind(&schema).bind(&name).bind(MAX_EDITOR_COLUMNS_PER_TABLE + 1) + .fetch_all(&pool).await.map_err(|error| error.to_string())?; + if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { truncated_columns = true; } + let mut columns = Vec::new(); + for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { + columns.push(QueryEditorColumn { + name: mysql_get_string(&column, 0, "column_name", "get_query_editor_metadata")?, + data_type: mysql_get_string(&column, 1, "data_type", "get_query_editor_metadata")?, + }); + } + tables.push(QueryEditorTable { schema, name, columns }); + } + + return Ok(QueryEditorMetadata { + tables, functions: Vec::new(), + truncated_tables, truncated_columns, truncated_functions: false, + }); + } + + if engine == DatabaseEngine::Sqlite { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + let table_rows = sqlx::query( + "select name from sqlite_master \ + where type = 'table' and name not like 'sqlite_%' \ + order by name limit ?", + ).bind(MAX_EDITOR_TABLES + 1).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; + let mut tables = Vec::new(); + let mut truncated_columns = false; + for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { + let name: String = sqlite_get_idx(&row, 0, "name", "get_query_editor_metadata")?; + require_safe_identifier(&name, "table name")?; + let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(&name)); + let column_rows = sqlx::query(&pragma_sql).fetch_all(&pool).await + .map_err(|error| error.to_string())?; + if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { truncated_columns = true; } + let mut columns = Vec::new(); + for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { + columns.push(QueryEditorColumn { + name: sqlite_get_name(&column, "name", "get_query_editor_metadata")?, + data_type: sqlite_get_name(&column, "type", "get_query_editor_metadata")?, + }); + } + tables.push(QueryEditorTable { schema: "main".to_string(), name, columns }); + } + return Ok(QueryEditorMetadata { + tables, functions: Vec::new(), + truncated_tables, truncated_columns, truncated_functions: false, + }); + } + + with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { + let table_rows = client.query( + "select n.nspname::text as schema_name, c.relname::text as table_name \ + from pg_class c \ + join pg_namespace n on n.oid = c.relnamespace \ + where c.relkind in ('r', 'p', 'v', 'm', 'f') \ + and n.nspname not in ('pg_catalog', 'information_schema') \ + order by n.nspname, c.relname \ + limit $1", + &[&(MAX_EDITOR_TABLES + 1)], + ).await.map_err(|error| map_pg_err(error, None))?; + + let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; + let table_rows = if truncated_tables { + table_rows.into_iter().take(MAX_EDITOR_TABLES as usize).collect::>() + } else { table_rows }; + + let mut tables = Vec::with_capacity(table_rows.len()); + let mut truncated_columns = false; + + for row in table_rows { + let schema: String = row.get(0); + let name: String = row.get(1); + let column_rows = client.query( + "select a.attname::text as column_name, format_type(a.atttypid, a.atttypmod)::text as data_type \ + from pg_attribute a \ + join pg_class c on c.oid = a.attrelid \ + join pg_namespace n on n.oid = c.relnamespace \ + where n.nspname = $1 and c.relname = $2 \ + and a.attnum > 0 and not a.attisdropped \ + order by a.attnum limit $3", + &[&schema, &name, &(MAX_EDITOR_COLUMNS_PER_TABLE + 1)], + ).await.map_err(|error| map_pg_err(error, None))?; + + if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { truncated_columns = true; } + let columns = column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) + .map(|column| QueryEditorColumn { name: column.get(0), data_type: column.get(1) }) + .collect(); + + tables.push(QueryEditorTable { schema, name, columns }); + } + + let function_rows = client.query( + "select n.nspname::text as schema_name, p.proname::text as function_name, \ + coalesce(pg_get_function_identity_arguments(p.oid), '')::text as args, \ + pg_get_function_result(p.oid)::text as return_type \ + from pg_proc p \ + join pg_namespace n on n.oid = p.pronamespace \ + where n.nspname not in ('pg_catalog', 'information_schema') \ + order by n.nspname, p.proname limit $1", + &[&(MAX_EDITOR_FUNCTIONS + 1)], + ).await.map_err(|error| map_pg_err(error, None))?; + + let truncated_functions = function_rows.len() as i64 > MAX_EDITOR_FUNCTIONS; + let functions = function_rows.into_iter().take(MAX_EDITOR_FUNCTIONS as usize) + .map(|row| { + let args_raw: String = row.get(2); + QueryEditorFunction { + schema: row.get(0), + name: row.get(1), + arg_types: if args_raw.trim().is_empty() { Vec::new() } + else { args_raw.split(',').map(|value| value.trim().to_string()).collect() }, + return_type: row.get(3), + } + }).collect(); + + Ok(QueryEditorMetadata { + tables, functions, + truncated_tables, truncated_columns, truncated_functions, + }) + }).await.map_err(String::from) +} + +pub(crate) async fn fetch_foreign_keys_for_connection( + app: &AppHandle, + state: &AppState, + connection_id: &str, + engine: DatabaseEngine, +) -> Result, String> { + if engine == DatabaseEngine::Mysql { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + let rows = sqlx::query( + "select kcu.table_schema as from_schema, kcu.table_name as from_table, \ + kcu.column_name as from_column, kcu.referenced_table_schema as to_schema, \ + kcu.referenced_table_name as to_table, kcu.referenced_column_name as to_column \ + from information_schema.key_column_usage kcu \ + where kcu.referenced_table_name is not null \ + order by kcu.table_schema, kcu.table_name, kcu.ordinal_position \ + limit ?", + ).bind(MAX_FOREIGN_KEY_ROWS).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let mut edges = Vec::new(); + for row in rows { + edges.push(ForeignKeyEdge { + from_schema: mysql_get_string(&row, 0, "from_schema", "get_foreign_keys")?, + from_table: mysql_get_string(&row, 1, "from_table", "get_foreign_keys")?, + from_column: mysql_get_string(&row, 2, "from_column", "get_foreign_keys")?, + to_schema: mysql_get_string(&row, 3, "to_schema", "get_foreign_keys")?, + to_table: mysql_get_string(&row, 4, "to_table", "get_foreign_keys")?, + to_column: mysql_get_string(&row, 5, "to_column", "get_foreign_keys")?, + }); + } + return Ok(edges); + } + + if engine == DatabaseEngine::Sqlite { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + let tables = sqlx::query( + "select name from sqlite_master \ + where type = 'table' and name not like 'sqlite_%'", + ).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let mut edges = Vec::new(); + for table in tables { + let table_name: String = sqlite_get_idx(&table, 0, "name", "get_foreign_keys")?; + require_safe_identifier(&table_name, "table name")?; + let fk_sql = format!("PRAGMA foreign_key_list(\"{}\");", quote_identifier(&table_name)); + let fk_rows = sqlx::query(&fk_sql).fetch_all(&pool).await.map_err(|error| error.to_string())?; + for row in fk_rows { + edges.push(ForeignKeyEdge { + from_schema: "main".to_string(), + from_table: table_name.clone(), + from_column: sqlite_get_name(&row, "from", "get_foreign_keys")?, + to_schema: "main".to_string(), + to_table: sqlite_get_name(&row, "table", "get_foreign_keys")?, + to_column: sqlite_get_name(&row, "to", "get_foreign_keys")?, + }); + if edges.len() >= MAX_FOREIGN_KEY_ROWS as usize { + return Ok(edges); + } + } + } + return Ok(edges); + } + + with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { + let rows = client.query( + "select src_ns.nspname::text as from_schema, src_cls.relname::text as from_table, \ + src_att.attname::text as from_column, tgt_ns.nspname::text as to_schema, \ + tgt_cls.relname::text as to_table, tgt_att.attname::text as to_column \ + from pg_constraint c \ + join pg_class src_cls on src_cls.oid = c.conrelid \ + join pg_namespace src_ns on src_ns.oid = src_cls.relnamespace \ + join pg_class tgt_cls on tgt_cls.oid = c.confrelid \ + join pg_namespace tgt_ns on tgt_ns.oid = tgt_cls.relnamespace \ + cross join lateral unnest(c.conkey, c.confkey) as u(attnum, confattnum) \ + join pg_attribute src_att on src_att.attrelid = c.conrelid \ + and src_att.attnum = u.attnum and not src_att.attisdropped \ + join pg_attribute tgt_att on tgt_att.attrelid = c.confrelid \ + and tgt_att.attnum = u.confattnum and not tgt_att.attisdropped \ + where c.contype = 'f' \ + and src_ns.nspname not in ('pg_catalog', 'information_schema') \ + order by src_ns.nspname, src_cls.relname, c.conname, u.attnum \ + limit $1", + &[&MAX_FOREIGN_KEY_ROWS], + ).await.map_err(|error| error.to_string())?; + + Ok(rows.into_iter().map(|row| ForeignKeyEdge { + from_schema: row.get(0), + from_table: row.get(1), + from_column: row.get(2), + to_schema: row.get(3), + to_table: row.get(4), + to_column: row.get(5), + }).collect()) + }).await.map_err(String::from) +} + +#[tauri::command] +pub async fn get_query_editor_metadata( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result { + let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; + + if engine == DatabaseEngine::Mysql { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + let database = load_connection(&app, &connection_id)? + .map(|connection| connection.database).unwrap_or_default(); + let table_rows = sqlx::query( + "select table_schema, table_name \ + from information_schema.tables \ + where table_type = 'BASE TABLE' \ + and table_schema = ? \ + and table_schema not in ('information_schema', 'mysql', 'performance_schema', 'sys') \ + order by table_schema, table_name \ + limit ?", + ).bind(&database).bind(MAX_EDITOR_TABLES + 1) + .fetch_all(&pool).await.map_err(|error| error.to_string())?; + + let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; + let mut tables = Vec::new(); + let mut truncated_columns = false; + + for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { + let schema: String = mysql_get_string(&row, 0, "table_schema", "get_query_editor_metadata")?; + let name: String = mysql_get_string(&row, 1, "table_name", "get_query_editor_metadata")?; + let column_rows = sqlx::query( + "select column_name, data_type \ + from information_schema.columns \ + where table_schema = ? and table_name = ? \ + order by ordinal_position limit ?", + ).bind(&schema).bind(&name).bind(MAX_EDITOR_COLUMNS_PER_TABLE + 1) + .fetch_all(&pool).await.map_err(|error| error.to_string())?; + if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { truncated_columns = true; } + let mut columns = Vec::new(); + for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { + columns.push(QueryEditorColumn { + name: mysql_get_string(&column, 0, "column_name", "get_query_editor_metadata")?, + data_type: mysql_get_string(&column, 1, "data_type", "get_query_editor_metadata")?, + }); + } + tables.push(QueryEditorTable { schema, name, columns }); + } + + return Ok(QueryEditorMetadata { + tables, functions: Vec::new(), + truncated_tables, truncated_columns, truncated_functions: false, + }); + } + + if engine == DatabaseEngine::Sqlite { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + let table_rows = sqlx::query( + "select name from sqlite_master \ + where type = 'table' and name not like 'sqlite_%' \ + order by name limit ?", + ).bind(MAX_EDITOR_TABLES + 1).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; + let mut tables = Vec::new(); + let mut truncated_columns = false; + for row in table_rows.into_iter().take(MAX_EDITOR_TABLES as usize) { + let name: String = sqlite_get_idx(&row, 0, "name", "get_query_editor_metadata")?; + require_safe_identifier(&name, "table name")?; + let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(&name)); + let column_rows = sqlx::query(&pragma_sql).fetch_all(&pool).await + .map_err(|error| error.to_string())?; + if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { truncated_columns = true; } + let mut columns = Vec::new(); + for column in column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) { + columns.push(QueryEditorColumn { + name: sqlite_get_name(&column, "name", "get_query_editor_metadata")?, + data_type: sqlite_get_name(&column, "type", "get_query_editor_metadata")?, + }); + } + tables.push(QueryEditorTable { schema: "main".to_string(), name, columns }); + } + return Ok(QueryEditorMetadata { + tables, functions: Vec::new(), + truncated_tables, truncated_columns, truncated_functions: false, + }); + } + + with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { + let table_rows = client.query( + "select n.nspname::text as schema_name, c.relname::text as table_name \ + from pg_class c \ + join pg_namespace n on n.oid = c.relnamespace \ + where c.relkind in ('r', 'p', 'v', 'm', 'f') \ + and n.nspname not in ('pg_catalog', 'information_schema') \ + order by n.nspname, c.relname \ + limit $1", + &[&(MAX_EDITOR_TABLES + 1)], + ).await.map_err(|error| map_pg_err(error, None))?; + + let truncated_tables = table_rows.len() as i64 > MAX_EDITOR_TABLES; + let table_rows = if truncated_tables { + table_rows.into_iter().take(MAX_EDITOR_TABLES as usize).collect::>() + } else { table_rows }; + + let mut tables = Vec::with_capacity(table_rows.len()); + let mut truncated_columns = false; + + for row in table_rows { + let schema: String = row.get(0); + let name: String = row.get(1); + let column_rows = client.query( + "select a.attname::text as column_name, format_type(a.atttypid, a.atttypmod)::text as data_type \ + from pg_attribute a \ + join pg_class c on c.oid = a.attrelid \ + join pg_namespace n on n.oid = c.relnamespace \ + where n.nspname = $1 and c.relname = $2 \ + and a.attnum > 0 and not a.attisdropped \ + order by a.attnum limit $3", + &[&schema, &name, &(MAX_EDITOR_COLUMNS_PER_TABLE + 1)], + ).await.map_err(|error| map_pg_err(error, None))?; + + if column_rows.len() as i64 > MAX_EDITOR_COLUMNS_PER_TABLE { truncated_columns = true; } + let columns = column_rows.into_iter().take(MAX_EDITOR_COLUMNS_PER_TABLE as usize) + .map(|column| QueryEditorColumn { name: column.get(0), data_type: column.get(1) }) + .collect(); + + tables.push(QueryEditorTable { schema, name, columns }); + } + + let function_rows = client.query( + "select n.nspname::text as schema_name, p.proname::text as function_name, \ + coalesce(pg_get_function_identity_arguments(p.oid), '')::text as args, \ + pg_get_function_result(p.oid)::text as return_type \ + from pg_proc p \ + join pg_namespace n on n.oid = p.pronamespace \ + where n.nspname not in ('pg_catalog', 'information_schema') \ + order by n.nspname, p.proname limit $1", + &[&(MAX_EDITOR_FUNCTIONS + 1)], + ).await.map_err(|error| map_pg_err(error, None))?; + + let truncated_functions = function_rows.len() as i64 > MAX_EDITOR_FUNCTIONS; + let functions = function_rows.into_iter().take(MAX_EDITOR_FUNCTIONS as usize) + .map(|row| { + let args_raw: String = row.get(2); + QueryEditorFunction { + schema: row.get(0), + name: row.get(1), + arg_types: if args_raw.trim().is_empty() { Vec::new() } + else { args_raw.split(',').map(|value| value.trim().to_string()).collect() }, + return_type: row.get(3), + } + }).collect(); + + Ok(QueryEditorMetadata { + tables, functions, + truncated_tables, truncated_columns, truncated_functions, + }) + }).await.map_err(String::from) +} diff --git a/src-tauri/src/commands/export_cmds.rs b/src-tauri/src/commands/export_cmds.rs new file mode 100644 index 0000000..471871c --- /dev/null +++ b/src-tauri/src/commands/export_cmds.rs @@ -0,0 +1,65 @@ +use tauri::{AppHandle, State}; + +use crate::db::AppState; +use crate::export::{DiagramExportRequest, ExportQueryRequest, export_diagram_to_png, export_results_csv, export_results_json}; +use crate::credentials; + +#[tauri::command] +pub async fn export_diagram_png( + input: DiagramExportRequest, + output_path: String, +) -> Result<(), String> { + let path = std::path::PathBuf::from(&output_path); + tokio::task::spawn_blocking(move || export_diagram_to_png(&input, &path)) + .await.map_err(|e| e.to_string())? +} + +#[tauri::command] +pub async fn export_results_csv_command( + app: AppHandle, + state: State<'_, AppState>, + input: ExportQueryRequest, +) -> Result<(), String> { + export_results_csv(&app, &state, &input).await +} + +#[tauri::command] +pub async fn export_results_json_command( + app: AppHandle, + state: State<'_, AppState>, + input: ExportQueryRequest, +) -> Result<(), String> { + export_results_json(&app, &state, &input).await +} + +#[tauri::command] +pub async fn save_base64_png(data: String, output_path: String) -> Result<(), String> { + use base64::Engine; + let bytes = base64::engine::general_purpose::STANDARD + .decode(data.strip_prefix("data:image/png;base64,").unwrap_or(&data)) + .map_err(|e| e.to_string())?; + std::fs::write(&output_path, bytes).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn save_text_file(content: String, output_path: String) -> Result<(), String> { + std::fs::write(&output_path, content).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn store_openrouter_api_key(api_key: String) -> Result<(), String> { + if api_key.trim().is_empty() { + return credentials::delete_openrouter_api_key(); + } + credentials::store_openrouter_api_key(&api_key) +} + +#[tauri::command] +pub async fn get_openrouter_api_key() -> Result, String> { + credentials::get_openrouter_api_key() +} + +#[tauri::command] +pub async fn delete_openrouter_api_key() -> Result<(), String> { + credentials::delete_openrouter_api_key() +} diff --git a/src-tauri/src/commands/lint.rs b/src-tauri/src/commands/lint.rs new file mode 100644 index 0000000..4f48ed0 --- /dev/null +++ b/src-tauri/src/commands/lint.rs @@ -0,0 +1,102 @@ +use tauri::{AppHandle, State}; + +use crate::db::{ + get_or_create_mysql_pool, get_or_create_sqlite_pool, resolve_connection_engine, + with_pool_client_retry, AppState, +}; +use crate::models::{DatabaseEngine, LintSqlRequest, LintSqlResult, SqlDiagnostic}; +use crate::pg_error::{error_line_column, map_pg_err}; + +use super::MAX_LINT_SQL_BYTES; + +#[tauri::command] +pub async fn lint_sql( + app: AppHandle, + state: State<'_, AppState>, + input: LintSqlRequest, +) -> Result { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let sql = input.sql.trim().to_string(); + if sql.is_empty() { + return Ok(LintSqlResult { diagnostics: Vec::new() }); + } + if sql.len() > MAX_LINT_SQL_BYTES { + return Err("SQL is too large to lint in the editor.".to_string()); + } + + match engine { + DatabaseEngine::Postgres => { + with_pool_client_retry(&app, &state, &connection_id, sql, |client, sql| async move { + let lint_sql = format!("EXPLAIN {}", sql); + let diagnostics = match client.simple_query(&lint_sql).await { + Ok(_) => Vec::new(), + Err(error) => { + let (line, column) = error_line_column(&error, &sql) + .map(|(l, c)| (Some(l), Some(c))) + .unwrap_or((None, None)); + vec![SqlDiagnostic { + message: map_pg_err(error, Some(sql.as_str())), + severity: "error".to_string(), + line, + column, + end_line: line, + end_column: column.map(|value| value + 1), + }] + } + }; + Ok(LintSqlResult { diagnostics }) + }).await.map_err(String::from) + } + DatabaseEngine::Mysql => { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + let lint_sql = format!("EXPLAIN {}", sql); + let diagnostics = match sqlx::query(&lint_sql).execute(&pool).await { + Ok(_) => Vec::new(), + Err(error) => vec![SqlDiagnostic { + message: error.to_string(), + severity: "error".to_string(), + line: None, column: None, end_line: None, end_column: None, + }], + }; + Ok(LintSqlResult { diagnostics }) + } + DatabaseEngine::Sqlite => { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + let lint_sql = format!("EXPLAIN QUERY PLAN {}", sql); + let diagnostics = match sqlx::query(&lint_sql).execute(&pool).await { + Ok(_) => Vec::new(), + Err(error) => vec![SqlDiagnostic { + message: error.to_string(), + severity: "error".to_string(), + line: None, column: None, end_line: None, end_column: None, + }], + }; + Ok(LintSqlResult { diagnostics }) + } + DatabaseEngine::Mongo => { + Err("MongoDB does not support SQL linting.".to_string()) + } + DatabaseEngine::Duckdb => { + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns + .get(&connection_id) + .ok_or("DuckDB connection not found.".to_string())?; + let conn = conn_mutex.lock().await; + let lint_sql = format!("EXPLAIN {}", sql); + match conn.prepare(&lint_sql) { + Ok(_) => Ok(LintSqlResult { diagnostics: vec![] }), + Err(e) => Ok(LintSqlResult { + diagnostics: vec![SqlDiagnostic { + message: e.to_string(), + severity: "error".to_string(), + line: None, + column: None, + end_line: None, + end_column: None, + }], + }), + } + } + DatabaseEngine::Redis => Err("Not supported for Redis.".to_string()), + } +} diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs new file mode 100644 index 0000000..55b4179 --- /dev/null +++ b/src-tauri/src/commands/mod.rs @@ -0,0 +1,982 @@ +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; + +use futures_util::StreamExt; +use serde_json::Value; +use tauri::{AppHandle, Emitter}; +use sqlx::{Column, Decode, Row, Type}; +use sqlx::mysql::{MySql, MySqlRow}; +use sqlx::sqlite::{Sqlite, SqliteRow}; + +use crate::db::{ + get_or_create_mysql_pool, get_or_create_sqlite_pool, AppState, +}; +use crate::models::{ + AskVeloxyDbContextCache, AskVeloxyTableRef, DatabaseEngine, + ForeignKeyEdge, QueryEditorTable, QueryResult, VeloxyStreamChunk, +}; +use crate::sql_split::split_sql_statements; + +// --- Constants --- + +pub(crate) const MAX_FOREIGN_KEY_ROWS: i64 = 5000; +pub(crate) const MAX_TABLE_INDEX_ROWS: i64 = 500; +pub(crate) const MAX_EDITOR_TABLES: i64 = 150; +pub(crate) const MAX_EDITOR_COLUMNS_PER_TABLE: i64 = 60; +pub(crate) const MAX_EDITOR_FUNCTIONS: i64 = 200; +pub(crate) const MAX_LINT_SQL_BYTES: usize = 65_536; +pub(crate) const ASK_VELOXY_MAX_CONTEXT_TABLES: usize = 8; +pub(crate) const ASK_VELOXY_MAX_CONTEXT_COLUMNS: usize = 18; +pub(crate) const ASK_VELOXY_MAX_CONTEXT_RELATIONSHIPS: usize = 36; +pub(crate) const ASK_VELOXY_SCHEMA_CHAR_BUDGET: usize = 6_000; +pub(crate) const ASK_VELOXY_PROMPT_CHAR_BUDGET: usize = 12_000; +pub(crate) const ASK_VELOXY_MAX_HISTORY_MESSAGES: usize = 30; +pub(crate) const ASK_VELOXY_MAX_CHAT_TOKENS: u32 = 10_000; + +// --- MySQL / SQLite decode helpers --- + +pub(crate) fn mysql_decode_error(context: &str, column_name: &str, index: Option, detail: &str) -> String { + match index { + Some(idx) => format!( + "MySQL decode error in {} at column '{}' (index {}): {}", + context, column_name, idx, detail + ), + None => format!( + "MySQL decode error in {} at column '{}': {}", + context, column_name, detail + ), + } +} + +pub(crate) fn sqlite_decode_error(context: &str, column_name: &str, index: Option, detail: &str) -> String { + match index { + Some(idx) => format!( + "SQLite decode error in {} at column '{}' (index {}): {}", + context, column_name, idx, detail + ), + None => format!( + "SQLite decode error in {} at column '{}': {}", + context, column_name, detail + ), + } +} + +pub(crate) fn mysql_get_idx(row: &MySqlRow, index: usize, column_name: &str, context: &str) -> Result +where + for<'r> T: Decode<'r, MySql> + Type, +{ + row.try_get::(index) + .map_err(|error| mysql_decode_error(context, column_name, Some(index), &error.to_string())) +} + +pub(crate) fn sqlite_get_idx(row: &SqliteRow, index: usize, column_name: &str, context: &str) -> Result +where + for<'r> T: Decode<'r, Sqlite> + Type, +{ + row.try_get::(index) + .map_err(|error| sqlite_decode_error(context, column_name, Some(index), &error.to_string())) +} + +pub(crate) fn sqlite_get_name(row: &SqliteRow, column_name: &str, context: &str) -> Result +where + for<'r> T: Decode<'r, Sqlite> + Type, +{ + row.try_get::(column_name).map_err(|error| { + format!( + "SQLite decode error in {} at column '{}': {}", + context, column_name, error + ) + }) +} + +pub(crate) fn database_name_from_mysql_value(value: Option, context: &str) -> Result { + let name = value + .filter(|value| !value.is_empty()) + .ok_or_else(|| format!("{context} returned an empty database name"))?; + Ok(name) +} + +pub(crate) fn mysql_database_name_from_row(row: &MySqlRow, context: &str) -> Result { + let value = mysql_value_to_string(row, 0, "Database", context)?; + database_name_from_mysql_value(value, context) +} + +pub(crate) fn mysql_value_to_string(row: &MySqlRow, index: usize, column_name: &str, context: &str) -> Result, String> { + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::>, _>(index) { + return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::>, _>(index) { + return Ok(value.map(|v| decode_mysql_bytes_as_string(&v))); + } + Err(mysql_decode_error(context, column_name, Some(index), "unsupported value type")) +} + +pub(crate) fn decode_mysql_bytes_as_string(bytes: &[u8]) -> String { + String::from_utf8_lossy(bytes).into_owned() +} + +pub(crate) fn mysql_value_to_display_string( + row: &MySqlRow, + index: usize, + column_name: &str, + context: &str, +) -> Result, String> { + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::>, _>(index) { + return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.format("%Y-%m-%d %H:%M:%S").to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::>, _>(index) { + return Ok(value.map(|v| format!("0x{}", hex::encode(v)))); + } + Err(mysql_decode_error(context, column_name, Some(index), "unsupported value type")) +} + +pub(crate) fn mysql_get_string(row: &MySqlRow, index: usize, column_name: &str, context: &str) -> Result { + let value = mysql_value_to_string(row, index, column_name, context)?; + value.ok_or_else(|| mysql_decode_error(context, column_name, Some(index), "unexpected null value")) +} + +pub(crate) fn mysql_get_optional_string( + row: &MySqlRow, + index: usize, + column_name: &str, + context: &str, +) -> Result, String> { + mysql_value_to_string(row, index, column_name, context) +} + +pub(crate) fn sqlite_value_to_string(row: &SqliteRow, index: usize, column_name: &str, context: &str) -> Result, String> { + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::, _>(index) { + return Ok(value.map(|v| v.to_string())); + } + if let Ok(value) = row.try_get::>, _>(index) { + return Ok(value.map(|v| format!("0x{}", hex::encode(v)))); + } + Err(sqlite_decode_error(context, column_name, Some(index), "unsupported value type")) +} + +// --- SQL helpers --- + +pub(crate) fn is_row_returning_sql(sql: &str) -> bool { + let trimmed = sql.trim_start(); + let upper = trimmed.to_uppercase(); + upper.starts_with("SELECT") + || upper.starts_with("WITH") + || upper.starts_with("SHOW") + || upper.starts_with("EXPLAIN") + || upper.starts_with("DESCRIBE") + || upper.starts_with("DESC") + || upper.starts_with("PRAGMA") + || upper.starts_with("VALUES") + || upper.starts_with("TABLE ") +} + +pub(crate) fn classify_sql_intent(sql: &str) -> String { + let normalized = sql.trim_start().to_ascii_lowercase(); + if normalized.starts_with("select") || normalized.starts_with("with") { + return "select".to_string(); + } + if normalized.starts_with("insert") { + return "insert".to_string(); + } + if normalized.starts_with("update") { + return "update".to_string(); + } + if normalized.starts_with("delete") { + return "delete".to_string(); + } + if normalized.starts_with("explain") { + return "explain".to_string(); + } + "unknown".to_string() +} + +pub(crate) fn is_read_only_sql(sql: &str) -> bool { + let mut saw_statement = false; + for statement in sql.split(';').map(str::trim).filter(|s| !s.is_empty()) { + let normalized = statement.to_ascii_lowercase(); + let is_transaction_control = ["begin", "commit", "rollback", "start", "savepoint", "release"] + .iter() + .any(|kw| normalized.starts_with(kw)); + if is_transaction_control { + continue; + } + saw_statement = true; + match classify_sql_intent(statement).as_str() { + "select" | "explain" => {} + _ => return false, + } + } + saw_statement +} + +pub(crate) fn has_multiple_statements(sql: &str) -> bool { + sql.split(';') + .map(str::trim) + .filter(|segment| !segment.is_empty()) + .count() > 1 +} + +pub(crate) fn validate_generated_sql(sql: &str) -> Result<(), String> { + let trimmed = sql.trim(); + if trimmed.is_empty() { + return Err("Ask Veloxy returned an empty SQL statement.".to_string()); + } + if has_multiple_statements(trimmed) { + return Err("Ask Veloxy generated multiple SQL statements. Please ask for a single statement.".to_string()); + } + Ok(()) +} + +// --- Row mapping --- + +pub(crate) type MappedRowsResult = (Vec, Vec>>, usize, bool); + +pub(crate) fn map_mysql_rows( + rows: Vec, + max_query_rows: usize, +) -> Result { + let mut columns: Vec = Vec::new(); + if let Some(first) = rows.first() { + columns = first.columns().iter().map(|column| column.name().to_string()).collect(); + } + let total_rows = rows.len(); + let mut mapped_rows = Vec::new(); + for row in rows.into_iter().take(max_query_rows) { + let mut mapped_row = BTreeMap::new(); + for (index, column_name) in columns.iter().enumerate() { + let value = mysql_value_to_display_string(&row, index, column_name, "run_query")?; + mapped_row.insert(column_name.clone(), value); + } + mapped_rows.push(mapped_row); + } + Ok((columns, mapped_rows, total_rows, total_rows > max_query_rows)) +} + +pub(crate) fn map_sqlite_rows( + rows: Vec, + max_query_rows: usize, +) -> Result { + let mut columns: Vec = Vec::new(); + if let Some(first) = rows.first() { + columns = first.columns().iter().map(|column| column.name().to_string()).collect(); + } + let total_rows = rows.len(); + let mut mapped_rows = Vec::new(); + for row in rows.into_iter().take(max_query_rows) { + let mut mapped_row = BTreeMap::new(); + for (index, column_name) in columns.iter().enumerate() { + let value = sqlite_value_to_string(&row, index, column_name, "run_query")?; + mapped_row.insert(column_name.clone(), value); + } + mapped_rows.push(mapped_row); + } + Ok((columns, mapped_rows, total_rows, total_rows > max_query_rows)) +} + +pub(crate) async fn run_query_mysql_or_sqlite( + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_query_rows: usize, + engine: DatabaseEngine, +) -> Result { + let started_at = Instant::now(); + let statements = split_sql_statements(sql); + if statements.is_empty() { + return Err("Enter a SQL statement before running the query.".to_string()); + } + + let mut columns = Vec::new(); + let mut rows = Vec::new(); + let mut total_rows = 0usize; + let mut truncated = false; + let mut command_tag: Option = None; + + match engine { + DatabaseEngine::Mysql => { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + let mut conn = pool.acquire().await.map_err(|error| error.to_string())?; + for statement in statements { + if is_row_returning_sql(&statement) { + let fetched = sqlx::query(&statement) + .fetch_all(&mut *conn) + .await + .map_err(|error| error.to_string())?; + let mapped = map_mysql_rows(fetched, max_query_rows)?; + columns = mapped.0; + rows = mapped.1; + total_rows = mapped.2; + truncated = mapped.3; + command_tag = None; + } else { + let result = sqlx::query(&statement) + .execute(&mut *conn) + .await + .map_err(|error| error.to_string())?; + let affected = result.rows_affected(); + command_tag = Some(affected); + if rows.is_empty() { + total_rows = affected as usize; + } + } + } + } + DatabaseEngine::Sqlite => { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + let mut conn = pool.acquire().await.map_err(|error| error.to_string())?; + for statement in statements { + if is_row_returning_sql(&statement) { + let fetched = sqlx::query(&statement) + .fetch_all(&mut *conn) + .await + .map_err(|error| error.to_string())?; + let mapped = map_sqlite_rows(fetched, max_query_rows)?; + columns = mapped.0; + rows = mapped.1; + total_rows = mapped.2; + truncated = mapped.3; + command_tag = None; + } else { + let result = sqlx::query(&statement) + .execute(&mut *conn) + .await + .map_err(|error| error.to_string())?; + let affected = result.rows_affected(); + command_tag = Some(affected); + if rows.is_empty() { + total_rows = affected as usize; + } + } + } + } + DatabaseEngine::Mongo => { + return Err("Internal engine routing error (MongoDB uses its own query path).".to_string()); + } + DatabaseEngine::Duckdb => { + return Err("Internal engine routing error (DuckDB uses its own query path).".to_string()); + } + DatabaseEngine::Redis => { + return Err("Internal engine routing error (Redis uses its own command path).".to_string()); + } + DatabaseEngine::Postgres => { + return Err("Internal engine routing error.".to_string()); + } + } + + Ok(QueryResult { + columns, + row_count: if rows.is_empty() { total_rows } else { rows.len() }, + rows, + execution_ms: started_at.elapsed().as_millis(), + truncated, + command_tag, + }) +} + +// --- Veloxy helpers --- + +pub(crate) fn estimate_tokens(chars: usize) -> usize { + (chars / 4).max(1) +} + +pub(crate) fn normalize_openrouter_base(base: Option<&str>) -> String { + let trimmed = base.unwrap_or("https://openrouter.ai/api/v1").trim(); + let value = if trimmed.is_empty() { "https://openrouter.ai/api/v1" } else { trimmed }; + value.trim_end_matches('/').to_string() +} + +pub(crate) fn truncate_on_char_boundary(value: &mut String, max_bytes: usize) { + if value.len() <= max_bytes { + return; + } + let mut truncate_at = max_bytes; + while !value.is_char_boundary(truncate_at) && truncate_at > 0 { + truncate_at -= 1; + } + value.truncate(truncate_at); +} + +pub(crate) fn ask_veloxy_context_cache_key(connection_id: &str, database_name: &str) -> String { + format!("{}::{}", connection_id, database_name) +} + +pub(crate) fn ask_veloxy_conversation_key(connection_id: &str, database_name: &str) -> String { + format!("{}::{}", connection_id, database_name) +} + +pub(crate) fn now_epoch_seconds() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() +} + +pub(crate) fn table_matches_target(table: &QueryEditorTable, target: Option<&AskVeloxyTableRef>) -> bool { + let Some(target) = target else { return false }; + table.schema.eq_ignore_ascii_case(&target.schema) && table.name.eq_ignore_ascii_case(&target.name) +} + +pub(crate) fn table_relevance_score(table: &QueryEditorTable, prompt_lower: &str) -> usize { + let mut score = 0usize; + let full_name = format!("{}.{}", table.schema.to_lowercase(), table.name.to_lowercase()); + if prompt_lower.contains(&table.name.to_lowercase()) { score += 3; } + if prompt_lower.contains(&table.schema.to_lowercase()) { score += 2; } + if prompt_lower.contains(&full_name) { score += 4; } + score +} + +pub(crate) fn relationship_relevance_score(edge: &ForeignKeyEdge, prompt_lower: &str) -> usize { + let from_name = format!("{}.{}", edge.from_schema.to_lowercase(), edge.from_table.to_lowercase()); + let to_name = format!("{}.{}", edge.to_schema.to_lowercase(), edge.to_table.to_lowercase()); + let mut score = 0usize; + if prompt_lower.contains(&edge.from_table.to_lowercase()) || prompt_lower.contains(&from_name) { + score += 2; + } + if prompt_lower.contains(&edge.to_table.to_lowercase()) || prompt_lower.contains(&to_name) { + score += 2; + } + score +} + +pub(crate) fn build_schema_context( + db_context: &AskVeloxyDbContextCache, + prompt: &str, + target_table: Option<&AskVeloxyTableRef>, +) -> String { + let prompt_lower = prompt.to_lowercase(); + let mut ranked: Vec<(&QueryEditorTable, usize, bool)> = db_context + .metadata + .tables + .iter() + .map(|table| { + (table, table_relevance_score(table, &prompt_lower), table_matches_target(table, target_table)) + }) + .collect(); + + ranked.sort_by(|a, b| b.2.cmp(&a.2).then_with(|| b.1.cmp(&a.1))); + + let mut schema_context = String::new(); + schema_context.push_str(&format!( + "database {} engine {:?}\n", db_context.database_name, db_context.engine + )); + for (table, _score, _is_target) in ranked.into_iter().take(ASK_VELOXY_MAX_CONTEXT_TABLES) { + let columns = table + .columns + .iter() + .take(ASK_VELOXY_MAX_CONTEXT_COLUMNS) + .map(|column| format!("{}:{}", column.name, column.data_type)) + .collect::>() + .join(", "); + schema_context.push_str(&format!( + "table {}.{} columns [{}]\n", table.schema, table.name, columns + )); + if schema_context.len() >= ASK_VELOXY_SCHEMA_CHAR_BUDGET { + truncate_on_char_boundary(&mut schema_context, ASK_VELOXY_SCHEMA_CHAR_BUDGET); + break; + } + } + + let mut ranked_relationships = db_context + .foreign_keys + .iter() + .map(|edge| (edge, relationship_relevance_score(edge, &prompt_lower))) + .collect::>(); + ranked_relationships.sort_by(|a, b| b.1.cmp(&a.1)); + for (edge, _score) in ranked_relationships + .into_iter() + .take(ASK_VELOXY_MAX_CONTEXT_RELATIONSHIPS) + { + schema_context.push_str(&format!( + "relationship {}.{}({}) -> {}.{}({})\n", + edge.from_schema, edge.from_table, edge.from_column, + edge.to_schema, edge.to_table, edge.to_column + )); + if schema_context.len() >= ASK_VELOXY_SCHEMA_CHAR_BUDGET { + truncate_on_char_boundary(&mut schema_context, ASK_VELOXY_SCHEMA_CHAR_BUDGET); + break; + } + } + schema_context +} + +pub(crate) fn extract_sql_draft_from_text(message: &str) -> Option { + let lowered = message.to_lowercase(); + let markers = ["select ", "with ", "insert ", "update ", "delete ", "explain "]; + let start = markers.iter().filter_map(|marker| lowered.find(marker)).min()?; + let mut sql = message[start..].trim().to_string(); + if let Some(idx) = sql.find("```") { sql.truncate(idx); } + if sql.ends_with('.') { sql.pop(); } + if sql.is_empty() { None } else { Some(sql) } +} + +pub(crate) fn parse_bool_field(value: &Value, field: &str, default: bool) -> bool { + value.get(field).and_then(Value::as_bool).unwrap_or(default) +} + +pub(crate) fn extract_openrouter_message_content(payload: &Value) -> Result { + let content_value = payload + .get("choices") + .and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("message")) + .and_then(|message| message.get("content")) + .ok_or_else(|| "OpenRouter response missing choices[0].message.content".to_string())?; + + if let Some(content) = content_value.as_str() { + return Ok(content.to_string()); + } + + if let Some(items) = content_value.as_array() { + let mut merged = String::new(); + for item in items { + if let Some(text) = item.get("text").and_then(Value::as_str) { + merged.push_str(text); + } + } + if !merged.trim().is_empty() { + return Ok(merged); + } + } + + Err("OpenRouter returned an unsupported message format.".to_string()) +} + +pub(crate) fn parse_ask_veloxy_json(content: &str) -> Result { + if let Ok(value) = serde_json::from_str::(content) { + return Ok(value); + } + let start = content.find('{'); + let end = content.rfind('}'); + match (start, end) { + (Some(start_idx), Some(end_idx)) if end_idx > start_idx => { + serde_json::from_str::(&content[start_idx..=end_idx]) + .map_err(|error| format!("Ask Veloxy response was not valid JSON: {}", error)) + } + _ => Err("Ask Veloxy response did not contain JSON.".to_string()), + } +} + +pub(crate) fn parse_ask_veloxy_suggestions(generated: &Value) -> Vec { + generated + .get("suggestions") + .and_then(Value::as_array) + .map(|items| { + items.iter() + .filter_map(Value::as_str) + .map(str::trim) + .filter(|item| !item.is_empty()) + .take(5) + .map(|item| { + let mut value = item.to_string(); + truncate_on_char_boundary(&mut value, 200); + value + }) + .collect::>() + }) + .unwrap_or_default() +} + +pub(crate) fn parse_ask_veloxy_chat_json(content: &str) -> Result { + if let Ok(value) = serde_json::from_str::(content) { + return Ok(value); + } + let start = content.find('{'); + let end = content.rfind('}'); + match (start, end) { + (Some(start_idx), Some(end_idx)) if end_idx > start_idx => { + serde_json::from_str::(&content[start_idx..=end_idx]) + .map_err(|error| format!("Ask Veloxy chat JSON was invalid: {}", error)) + } + _ => Err("Ask Veloxy chat response did not contain JSON.".to_string()), + } +} + +fn decode_json_quoted_string(value: &str) -> Option { + serde_json::from_str::(&format!("\"{}\"", value)).ok() +} + +fn unescape_json_string_fragment(raw: &str) -> String { + let mut out = String::with_capacity(raw.len()); + let mut chars = raw.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == '\\' { + match chars.next() { + Some('n') => out.push('\n'), + Some('t') => out.push('\t'), + Some('r') => out.push('\r'), + Some('"') => out.push('"'), + Some('\\') => out.push('\\'), + Some(other) => { out.push('\\'); out.push(other); } + None => out.push('\\'), + } + } else { + out.push(ch); + } + } + out +} + +fn extract_json_string_field(content: &str, key: &str, allow_partial: bool) -> Option { + let marker = format!("\"{}\"", key); + let marker_idx = content.find(&marker)?; + let mut idx = marker_idx + marker.len(); + let bytes = content.as_bytes(); + + while idx < bytes.len() && bytes[idx].is_ascii_whitespace() { idx += 1; } + if idx >= bytes.len() || bytes[idx] != b':' { return None; } + idx += 1; + while idx < bytes.len() && bytes[idx].is_ascii_whitespace() { idx += 1; } + if idx >= bytes.len() || bytes[idx] != b'"' { return None; } + idx += 1; + let start = idx; + let mut escaped = false; + while idx < bytes.len() { + let byte = bytes[idx]; + if escaped { escaped = false; idx += 1; continue; } + if byte == b'\\' { escaped = true; idx += 1; continue; } + if byte == b'"' { + let raw = &content[start..idx]; + return decode_json_quoted_string(raw) + .or_else(|| Some(unescape_json_string_fragment(raw))) + .map(|text| text.trim().to_string()) + .filter(|text| !text.is_empty()); + } + idx += 1; + } + + if allow_partial && start < bytes.len() { + let raw = &content[start..]; + let text = unescape_json_string_fragment(raw).trim().to_string(); + if !text.is_empty() { return Some(text); } + } + None +} + +pub(crate) fn extract_message_from_loose_json(content: &str) -> Option { + let trimmed = content.trim(); + if trimmed.is_empty() { return None; } + let unwrapped = trimmed + .strip_prefix("```json") + .or_else(|| trimmed.strip_prefix("```JSON")) + .map(str::trim_start) + .unwrap_or(trimmed); + let unwrapped = unwrapped.strip_suffix("```").unwrap_or(unwrapped).trim(); + + ["message", "reply", "content"] + .iter() + .find_map(|key| extract_json_string_field(unwrapped, key, false)) + .or_else(|| { + ["message", "reply", "content"] + .iter() + .find_map(|key| extract_json_string_field(unwrapped, key, true)) + }) +} + +pub(crate) fn looks_like_json_response(content: &str) -> bool { + let trimmed = content.trim_start(); + trimmed.starts_with('{') || trimmed.starts_with("```") +} + +pub(crate) fn streaming_display_text(accumulated: &str) -> String { + let trimmed = accumulated.trim(); + if trimmed.is_empty() { return String::new(); } + if let Some(text) = extract_message_from_loose_json(trimmed) { return text; } + if !looks_like_json_response(trimmed) { return trimmed.to_string(); } + String::new() +} + +fn parse_chat_message(value: &Value) -> Option { + if let Some(text) = value.as_str().map(str::trim).filter(|text| !text.is_empty()).map(str::to_string) { + return Some(text); + } + value.get("message").and_then(Value::as_str) + .or_else(|| value.get("reply").and_then(Value::as_str)) + .map(str::trim).filter(|text| !text.is_empty()).map(str::to_string) +} + +pub(crate) type ParsedAskVeloxyChat = ( + String, Vec, Vec, Option, bool, bool, +); + +pub(crate) fn parse_ask_veloxy_chat_content(message_content: &str) -> ParsedAskVeloxyChat { + match parse_ask_veloxy_chat_json(message_content) { + Ok(value) => { + let message = parse_chat_message(&value).unwrap_or_else(|| message_content.trim().to_string()); + let mut draft = value.get("sqlDraft").and_then(Value::as_str) + .or_else(|| value.get("sql_draft").and_then(Value::as_str)) + .map(str::trim).filter(|text| !text.is_empty()).map(str::to_string); + if draft.is_none() { draft = extract_sql_draft_from_text(&message); } + let suggestions = value.get("suggestions").and_then(Value::as_array) + .map(|items| items.iter().filter_map(Value::as_str).map(str::trim) + .filter(|text| !text.is_empty()).take(5).map(str::to_string).collect::>()) + .unwrap_or_default(); + let warnings = value.get("warnings").and_then(Value::as_array) + .map(|items| items.iter().filter_map(Value::as_str).map(str::to_string).collect::>()) + .unwrap_or_default(); + let needs_sql_generation = parse_bool_field(&value, "needsSqlGeneration", draft.is_some()); + let needs_clarification = parse_bool_field(&value, "needsClarification", false); + (message, suggestions, warnings, draft, needs_sql_generation, needs_clarification) + } + Err(_) => { + let normalized_message = extract_message_from_loose_json(message_content).unwrap_or_else(|| { + if looks_like_json_response(message_content) { String::new() } + else { message_content.trim().to_string() } + }); + let mut warnings = vec!["Model returned non-JSON chat output. Parsed in tolerant mode.".to_string()]; + if normalized_message.is_empty() && looks_like_json_response(message_content) { + warnings.push("Response JSON could not be parsed. Try asking again.".to_string()); + } + let draft = extract_sql_draft_from_text(&normalized_message); + let needs_sql_generation = draft.is_some(); + (normalized_message, Vec::new(), warnings, draft, needs_sql_generation, false) + } + } +} + +pub(crate) fn extract_openrouter_stream_delta(data: &str) -> Option { + let payload: Value = serde_json::from_str(data).ok()?; + payload.get("choices").and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("delta")) + .and_then(|delta| delta.get("content")) + .and_then(Value::as_str).filter(|text| !text.is_empty()).map(str::to_string) +} + +pub(crate) fn extract_openrouter_finish_reason(data: &str) -> Option { + let payload: Value = serde_json::from_str(data).ok()?; + payload.get("choices").and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("finish_reason")) + .and_then(Value::as_str).map(str::to_string) +} + +pub(crate) fn emit_veloxy_stream_chunk(app: &AppHandle, chunk: VeloxyStreamChunk) { + let _ = app.emit("veloxy-stream-chunk", chunk); +} + +#[allow(clippy::too_many_arguments)] +pub(crate) async fn stream_openrouter_chat_completion( + app: &AppHandle, + client: &reqwest::Client, + endpoint: &str, + api_key: &str, + model: &str, + system_prompt: &str, + user_prompt: &str, + request_id: &str, + cancel: Arc, +) -> Result<(String, bool), String> { + let response = client + .post(endpoint) + .header("Authorization", format!("Bearer {}", api_key)) + .header("Content-Type", "application/json") + .json(&serde_json::json!({ + "model": model, + "temperature": 0.2, + "max_tokens": ASK_VELOXY_MAX_CHAT_TOKENS, + "stream": true, + "messages": [ + { "role": "system", "content": system_prompt }, + { "role": "user", "content": user_prompt } + ] + })) + .send() + .await + .map_err(|error| format!("OpenRouter request failed: {}", error))?; + + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_else(|_| "Unknown OpenRouter error".to_string()); + if let Ok(payload) = serde_json::from_str::(&body) { + let message = payload.get("error").and_then(|error| error.get("message")) + .and_then(Value::as_str).unwrap_or("Unknown OpenRouter error"); + return Err(format!("OpenRouter error ({}): {}", status.as_u16(), message)); + } + return Err(format!("OpenRouter error ({}): {}", status.as_u16(), body)); + } + + let mut stream = response.bytes_stream(); + let mut buffer = String::new(); + let mut accumulated = String::new(); + let mut last_display_len = 0usize; + let mut hit_token_limit = false; + + while let Some(chunk) = stream.next().await { + if cancel.load(Ordering::Relaxed) { + return Ok((accumulated, hit_token_limit)); + } + let bytes = chunk.map_err(|error| format!("OpenRouter stream read failed: {}", error))?; + buffer.push_str(&String::from_utf8_lossy(&bytes)); + + while let Some(line_end) = buffer.find('\n') { + let line = buffer[..line_end].trim_end_matches('\r').to_string(); + buffer.drain(..=line_end); + + if !line.starts_with("data: ") { continue; } + let data = line["data: ".len()..].trim(); + if data == "[DONE]" { continue; } + if extract_openrouter_finish_reason(data).as_deref() == Some("length") { + hit_token_limit = true; + } + if let Some(delta) = extract_openrouter_stream_delta(data) { + accumulated.push_str(&delta); + let display = streaming_display_text(&accumulated); + let display_delta = if display.len() > last_display_len { + display[last_display_len..].to_string() + } else { + String::new() + }; + last_display_len = display.len(); + if !display_delta.is_empty() { + emit_veloxy_stream_chunk(app, VeloxyStreamChunk { + request_id: request_id.to_string(), + delta: display_delta, + done: false, + message: None, + suggestions: Vec::new(), + warnings: Vec::new(), + sql_draft: None, + needs_sql_generation: false, + needs_clarification: false, + }); + } + } + } + } + + Ok((accumulated, hit_token_limit)) +} + +pub(crate) fn veloxdb_unique_constraint_name(table_name: &str, column_name: &str) -> String { + let suffix = "_uniq"; + let max_base_len = 63usize.saturating_sub(suffix.len()); + let mut base = format!("veloxdb_{}_{}", table_name, column_name); + base.truncate(max_base_len); + format!("{}{}", base, suffix) +} + +// --- Sub-modules --- + +mod connections; +mod query; +mod table_props; +mod ddl; +mod editor_meta; +mod veloxy; +mod lint; +pub(crate) mod mongo; +pub(crate) mod duckdb; +mod export_cmds; +pub(crate) mod redis; +mod updater; + +// --- Re-exports --- + +pub use connections::{ + connect_db, list_connections_command, set_active_connection, ping_connection, + refresh_connection, disconnect_db, rename_connection, delete_connection, + list_databases, switch_database, +}; +pub use query::{ + run_query, get_tables, get_schema, +}; +pub use table_props::{ + get_table_properties, apply_table_properties, get_foreign_keys, get_table_indexes, +}; +pub use ddl::{ + execute_ddl_transaction, execute_ddl_statement, +}; +pub use editor_meta::{ + get_query_editor_metadata, +}; +pub use veloxy::{ + cancel_veloxy_request, chat_with_db, clear_veloxy_conversation, generate_sql_from_nl, + load_veloxy_conversation, +}; +pub use lint::{ + lint_sql, +}; +pub use export_cmds::{ + export_diagram_png, export_results_csv_command, export_results_json_command, + save_base64_png, save_text_file, store_openrouter_api_key, get_openrouter_api_key, + delete_openrouter_api_key, +}; +pub use mongo::{ + mongo_run_query, mongo_get_collections, mongo_get_schema, +}; +pub use duckdb::{ + duckdb_run_query, duckdb_get_tables, duckdb_get_schema, +}; +pub use redis::{ + redis_run_query, redis_get_keys, redis_get_schema, +}; +pub use updater::{ + check_for_updates, +}; diff --git a/src-tauri/src/commands/mongo.rs b/src-tauri/src/commands/mongo.rs new file mode 100644 index 0000000..f5b0aee --- /dev/null +++ b/src-tauri/src/commands/mongo.rs @@ -0,0 +1,909 @@ +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::time::Instant; + +use futures_util::StreamExt; +use mongodb::bson::{doc, Document}; +use tauri::{AppHandle, State}; + +use crate::db::{get_or_create_mongo_client, resolve_connection_engine, AppState, MAX_QUERY_ROWS}; +use crate::models::{ColumnInfo, QueryRequest, QueryResult, TableIndexesResult, TableInfo}; + +/// Parse a user-supplied MongoDB query string into a filter Document. +/// +/// Supports these input forms: +/// `{"status": "active"}` — raw JSON filter +/// `db.collection.find({...})` — full shell syntax (ignores collection prefix) +/// `{ status: "active" }` — relaxed JSON (unquoted keys, single quotes) +fn parse_mongo_filter(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Ok(doc! {}); + } + + // Try strict JSON first + if let Ok(doc) = serde_json::from_str::(trimmed) { + if !doc.is_empty() { + return Ok(doc); + } + } + + // Try shell syntax: db.collection.find({...}) + if let Some(start) = trimmed.find(".find(") { + let after_find = &trimmed[start + 6..]; // skip ".find(" + let inner = extract_braced_json(after_find) + .ok_or_else(|| "Could not parse .find() arguments.".to_string())?; + return normalize_to_document(&inner); + } + + // Try relaxed JSON (single quotes, unquoted keys) + let relaxed = trimmed + .replace('\'', "\"") + .replace("ObjectId(", "\"ObjectId(") + .replace("ISODate(", "\"ISODate("); + if let Ok(doc) = serde_json::from_str::(&relaxed) { + return Ok(doc); + } + + // Last attempt: treat as a raw key-value search + normalize_to_document(trimmed) +} + +fn extract_braced_json(input: &str) -> Option { + let trimmed = input.trim(); + if !trimmed.starts_with('{') { + return None; + } + let mut depth = 0i32; + let mut end = 0usize; + for (i, ch) in trimmed.char_indices() { + match ch { + '{' => depth += 1, + '}' => { + depth -= 1; + if depth == 0 { + end = i + 1; + break; + } + } + _ => {} + } + } + if end > 0 { + Some(trimmed[..end].to_string()) + } else { + None + } +} + +/// Extract the collection name from MongoDB shell syntax like `db.users` or `db.getSiblingDB('x').users`. +fn extract_shell_collection_name(prefix: &str) -> String { + prefix + .split('.') + .last() + .unwrap_or(prefix) + .trim() + .to_string() +} + +/// Resolve `(database, collection)` for export from a MongoDB query string. +pub fn resolve_mongo_export_target( + query: &str, + default_database: &str, +) -> Result<(String, String), String> { + let trimmed = query.trim(); + if trimmed.is_empty() { + return Err("No MongoDB query to export.".to_string()); + } + + // Shell syntax: db.collection.find(...) + if let Some(dot_find) = trimmed.find(".find(") { + let collection = extract_shell_collection_name(trimmed[..dot_find].trim()); + if collection.is_empty() { + return Err("Could not determine MongoDB collection from query.".to_string()); + } + return Ok((default_database.to_string(), collection)); + } + + // database.collection (no shell find) + if let Some(dot) = trimmed.find('.') { + let prefix = trimmed[..dot].trim(); + let suffix = trimmed[dot + 1..].trim(); + let collection = suffix + .split(|c: char| c.is_whitespace() || c == '(') + .next() + .unwrap_or(suffix) + .trim(); + if prefix == "db" { + if collection.is_empty() { + return Err("Could not determine MongoDB collection from query.".to_string()); + } + return Ok((default_database.to_string(), collection.to_string())); + } + if !prefix.is_empty() && !collection.is_empty() { + return Ok((prefix.to_string(), collection.to_string())); + } + } + + // Bare collection name + let collection = trimmed + .split(|c: char| c.is_whitespace() || c == '(') + .next() + .unwrap_or(trimmed) + .trim(); + if collection.is_empty() { + return Err("Could not determine MongoDB collection from query.".to_string()); + } + Ok((default_database.to_string(), collection.to_string())) +} + +fn normalize_to_document(raw: &str) -> Result { + // Try parsing as BSON + if let Ok(doc) = serde_json::from_str::(raw) { + return Ok(doc); + } + // Return empty filter — matches all documents + Ok(doc! {}) +} + +fn bson_to_display_string(value: &mongodb::bson::Bson) -> Option { + use mongodb::bson::Bson; + match value { + Bson::String(s) => Some(s.clone()), + Bson::Int32(n) => Some(n.to_string()), + Bson::Int64(n) => Some(n.to_string()), + Bson::Double(f) => Some(f.to_string()), + Bson::Boolean(b) => Some(b.to_string()), + Bson::ObjectId(oid) => Some(oid.to_hex()), + Bson::DateTime(dt) => Some( + chrono::DateTime::from_timestamp_millis(dt.timestamp_millis()) + .map(|d| d.format("%Y-%m-%d %H:%M:%S").to_string()) + .unwrap_or_else(|| dt.to_string()), + ), + Bson::Null => None, + Bson::Array(arr) => Some(format!("[{} items]", arr.len())), + Bson::Document(_) => Some("{...}".to_string()), + Bson::Binary(_) => Some("".to_string()), + Bson::RegularExpression(re) => Some(format!("/{}/", re.pattern)), + Bson::Decimal128(d) => Some(d.to_string()), + _ => Some(format!("{:?}", value)), + } +} + +fn infer_bson_type(value: &mongodb::bson::Bson) -> &'static str { + use mongodb::bson::Bson; + match value { + Bson::String(_) => "string", + Bson::Int32(_) | Bson::Int64(_) => "integer", + Bson::Double(_) => "double", + Bson::Boolean(_) => "boolean", + Bson::ObjectId(_) => "objectId", + Bson::DateTime(_) => "date", + Bson::Array(_) => "array", + Bson::Document(_) => "object", + Bson::Null => "null", + Bson::Binary(_) => "binary", + Bson::RegularExpression(_) => "regex", + Bson::Decimal128(_) => "decimal", + _ => "unknown", + } +} + +/// Execute a MongoDB find/aggregate query and return tabular results. +#[tauri::command] +pub async fn mongo_run_query( + app: AppHandle, + state: State<'_, AppState>, + input: QueryRequest, +) -> Result { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let client = get_or_create_mongo_client(&app, &state, &connection_id).await?; + + let started_at = Instant::now(); + let raw = input.sql.trim().to_string(); + if raw.is_empty() { + return Err("Enter a MongoDB query or filter.".to_string()); + } + + // Determine the collection name — try extracting from shell syntax + let (collection_name, filter) = if let Some(dot_find) = raw.find(".find(") { + let collection = extract_shell_collection_name(raw[..dot_find].trim()); + let filter = parse_mongo_filter(&raw)?; + (collection, filter) + } else if raw.starts_with('{') { + // No collection specified — use the stored database's first collection as context, + // or just return all databases if none specified. + let db = client.default_database().ok_or_else(|| { + "No default database. Specify a database or use db.collection.find({...}) syntax." + .to_string() + })?; + let names = db + .list_collection_names() + .await + .map_err(|e| format!("Failed to list collections: {}", e))?; + if names.is_empty() { + return Err("No collections found in the current database.".to_string()); + } + let filter = parse_mongo_filter(&raw)?; + (names[0].clone(), filter) + } else { + return Err( + "Use MongoDB shell syntax: db.collection.find({...}) or provide a JSON filter." + .to_string(), + ); + }; + + // Try to determine the database name from the collection + let db = client.default_database().ok_or_else(|| { + "No default database available. Try switching to a database first.".to_string() + })?; + + let collection = db.collection::(&collection_name); + let max_rows = input.max_rows.unwrap_or(MAX_QUERY_ROWS) as i64; + + let mut cursor = collection + .find(filter) + .limit(max_rows) + .await + .map_err(|e| format!("MongoDB query failed: {}", e))?; + + let mut rows: Vec>> = Vec::new(); + let mut columns: Vec = Vec::new(); + let mut columns_set = HashSet::new(); + let mut total = 0usize; + + while let Some(result) = cursor.next().await { + let doc = result.map_err(|e| format!("MongoDB cursor error: {}", e))?; + if total == 0 { + columns = doc.keys().cloned().collect(); + for col in &columns { + columns_set.insert(col.clone()); + } + } else { + // Discover any new fields in subsequent documents + for key in doc.keys() { + if !columns_set.contains(key) { + columns.push(key.clone()); + columns_set.insert(key.clone()); + } + } + } + let mut row = BTreeMap::new(); + for key in &columns { + let value = doc.get(key).and_then(bson_to_display_string); + row.insert(key.clone(), value); + } + rows.push(row); + total += 1; + if total >= max_rows as usize { + break; + } + } + + Ok(QueryResult { + columns, + rows, + row_count: total, + execution_ms: started_at.elapsed().as_millis(), + truncated: false, + command_tag: None, + }) +} + +/// List all collections in the active database. +#[tauri::command] +pub async fn mongo_get_collections( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, connection_id.clone()).await?; + let client = get_or_create_mongo_client(&app, &state, &connection_id).await?; + + let db = client.default_database().ok_or_else(|| { + "No default database available.".to_string() + })?; + + let db_name = db.name().to_string(); + let names = db + .list_collection_names() + .await + .map_err(|e| format!("Failed to list collections: {}", e))?; + + Ok(names + .into_iter() + .map(|name| TableInfo { + schema: db_name.clone(), + name: name.clone(), + preview_query: format!("db.{}.find({{}}).limit(100)", name), + }) + .collect()) +} + +/// Infer the schema of a MongoDB collection by sampling documents. +#[tauri::command] +pub async fn mongo_get_schema( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, + database: String, + collection: String, +) -> Result, String> { + let (connection_id, _engine) = + resolve_connection_engine(&app, &state, connection_id.clone()).await?; + let client = get_or_create_mongo_client(&app, &state, &connection_id).await?; + + let db = client.database(&database); + let coll = db.collection::(&collection); + + let mut cursor = coll + .find(doc! {}) + .limit(200) + .await + .map_err(|e| format!("MongoDB query failed: {}", e))?; + + // Track field → BSON type. First-seen type wins. + let mut type_map: HashMap = HashMap::new(); + // Track field → nullable. If any document lacks the field, it's nullable. + let mut doc_count = 0usize; + let mut field_doc_counts: HashMap = HashMap::new(); + + while let Some(result) = cursor.next().await { + let doc = result.map_err(|e| format!("MongoDB cursor error: {}", e))?; + doc_count += 1; + + for (key, value) in doc.iter() { + *field_doc_counts.entry(key.clone()).or_default() += 1; + type_map + .entry(key.clone()) + .or_insert_with(|| infer_bson_type(value).to_string()); + } + } + + if doc_count == 0 { + return Ok(Vec::new()); + } + + Ok(type_map + .into_iter() + .map(|(name, data_type)| { + let field_count = field_doc_counts.get(&name).copied().unwrap_or(0); + let is_nullable = field_count < doc_count; + ColumnInfo { + table_schema: database.clone(), + table_name: collection.clone(), + column_name: name, + data_type, + is_nullable, + } + }) + .collect()) +} + +// ── MongoDB Export ────────────────────────────────────────────── + +pub async fn mongo_export_csv( + app: &AppHandle, + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + output_path: &str, +) -> Result<(), String> { + let client = crate::db::get_or_create_mongo_client(app, state, connection_id).await?; + let db = client.database(database); + let coll = db.collection::(collection); + + let mut cursor = coll.find(doc! {}).limit(5000).await + .map_err(|e| format!("MongoDB export failed: {}", e))?; + + let mut columns: Vec = Vec::new(); + let mut rows: Vec>> = Vec::new(); + let mut cols_set = HashSet::new(); + + while let Some(result) = cursor.next().await { + let doc = result.map_err(|e| format!("MongoDB cursor error: {}", e))?; + if columns.is_empty() { + columns = doc.keys().cloned().collect(); + for c in &columns { cols_set.insert(c.clone()); } + } else { + for k in doc.keys() { + if !cols_set.contains(k) { + columns.push(k.clone()); + cols_set.insert(k.clone()); + } + } + } + let mut row = BTreeMap::new(); + for key in &columns { + row.insert(key.clone(), doc.get(key).and_then(bson_to_display_string)); + } + rows.push(row); + } + + let mut wtr = csv::Writer::from_path(output_path).map_err(|e| e.to_string())?; + wtr.write_record(&columns).map_err(|e| e.to_string())?; + for row in &rows { + let record: Vec = columns.iter().map(|c| row.get(c).unwrap_or(&None).clone().unwrap_or_default()).collect(); + wtr.write_record(&record).map_err(|e| e.to_string())?; + } + wtr.flush().map_err(|e| e.to_string())?; + Ok(()) +} + +pub async fn mongo_export_json( + app: &AppHandle, + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + output_path: &str, +) -> Result<(), String> { + let client = crate::db::get_or_create_mongo_client(app, state, connection_id).await?; + let db = client.database(database); + let coll = db.collection::(collection); + + let mut cursor = coll.find(doc! {}).limit(5000).await + .map_err(|e| format!("MongoDB export failed: {}", e))?; + + let mut docs: Vec = Vec::new(); + while let Some(result) = cursor.next().await { + let doc = result.map_err(|e| format!("MongoDB cursor error: {}", e))?; + let json_str = serde_json::to_string(&mongodb::bson::to_document(&doc).unwrap_or_default()) + .map_err(|e| e.to_string())?; + docs.push(serde_json::from_str(&json_str).unwrap_or(serde_json::Value::Null)); + } + let content = serde_json::to_string_pretty(&docs).map_err(|e| e.to_string())?; + std::fs::write(output_path, content).map_err(|e| e.to_string())?; + Ok(()) +} + +/// List indexes on a MongoDB collection. +pub async fn mongo_get_table_indexes( + app: &AppHandle, + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, +) -> Result { + use crate::models::IndexInfo; + let client = crate::db::get_or_create_mongo_client(app, state, connection_id).await?; + let db = client.database(database); + let coll = db.collection::(collection); + + let mut cursor = coll.list_indexes().await + .map_err(|e| format!("MongoDB index listing failed: {}", e))?; + + let mut indexes = Vec::new(); + while let Some(result) = cursor.next().await { + let idx = result.map_err(|e| format!("MongoDB index cursor error: {}", e))?; + let name = idx.options.as_ref() + .and_then(|o| o.name.clone()) + .unwrap_or_else(|| { + idx.keys.iter() + .map(|(k, _)| k.clone()) + .collect::>() + .join("_") + }); + let keys_doc = idx.keys.clone(); + let keys: Vec = keys_doc.iter() + .map(|(k, v)| { + let dir = match v.as_i32() { Some(1) => "asc", Some(-1) => "desc", _ => "?" }; + format!("{}({})", k, dir) + }) + .collect(); + indexes.push(IndexInfo { + index_schema: database.to_string(), + index_name: name.clone(), + table_schema: database.to_string(), + table_name: collection.to_string(), + is_unique: false, + is_primary: name == "_id_", + is_valid: true, + is_partial: false, + definition: format!("index {} ({})", name, keys.join(", ")), + index_bytes: 0, + idx_scan: 0, + idx_tup_read: 0, + idx_tup_fetch: 0, + }); + } + + let truncated = indexes.len() > 500; + Ok(TableIndexesResult { + indexes: if truncated { indexes.into_iter().take(500).collect() } else { indexes }, + truncated, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use mongodb::bson::{oid::ObjectId, Bson, DateTime as BsonDateTime}; + + #[test] + fn parse_empty_filter() { + let result = parse_mongo_filter("").unwrap(); + assert!(result.is_empty()); + } + + #[test] + fn parse_whitespace_filter() { + let result = parse_mongo_filter(" ").unwrap(); + assert!(result.is_empty()); + } + + #[test] + fn parse_valid_json() { + let result = parse_mongo_filter(r#"{"status": "active"}"#).unwrap(); + assert_eq!(result.get_str("status").unwrap(), "active"); + } + + #[test] + fn parse_json_with_number() { + let result = parse_mongo_filter(r#"{"age": 25}"#).unwrap(); + assert_eq!(result.get_i32("age").unwrap(), 25); + } + + #[test] + fn parse_json_with_boolean() { + let result = parse_mongo_filter(r#"{"verified": true}"#).unwrap(); + assert!(result.get_bool("verified").unwrap()); + } + + #[test] + fn parse_json_with_null() { + let result = parse_mongo_filter(r#"{"deletedAt": null}"#).unwrap(); + assert!(result.get("deletedAt").unwrap().as_null().is_some()); + } + + #[test] + fn parse_json_with_operator() { + let result = parse_mongo_filter(r#"{"age": {"$gt": 18}}"#).unwrap(); + let age = result.get_document("age").unwrap(); + assert_eq!(age.get_i32("$gt").unwrap(), 18); + } + + #[test] + fn extract_shell_collection_name_from_db_prefix() { + assert_eq!(extract_shell_collection_name("db.users"), "users"); + } + + #[test] + fn extract_shell_collection_name_without_db_prefix() { + assert_eq!(extract_shell_collection_name("users"), "users"); + } + + #[test] + fn extract_shell_collection_name_from_nested_path() { + assert_eq!( + extract_shell_collection_name("db.getSiblingDB('other').orders"), + "orders" + ); + } + + #[test] + fn resolve_mongo_export_target_from_shell_syntax() { + let (db, coll) = + resolve_mongo_export_target(r#"db.users.find({"status": "active"})"#, "mydb").unwrap(); + assert_eq!(db, "mydb"); + assert_eq!(coll, "users"); + } + + #[test] + fn resolve_mongo_export_target_from_database_collection() { + let (db, coll) = resolve_mongo_export_target("analytics.orders", "mydb").unwrap(); + assert_eq!(db, "analytics"); + assert_eq!(coll, "orders"); + } + + #[test] + fn resolve_mongo_export_target_from_bare_collection() { + let (db, coll) = resolve_mongo_export_target("users", "mydb").unwrap(); + assert_eq!(db, "mydb"); + assert_eq!(coll, "users"); + } + + #[test] + fn parse_shell_syntax_find() { + let result = parse_mongo_filter(r#"db.users.find({"status": "active"})"#).unwrap(); + assert_eq!(result.get_str("status").unwrap(), "active"); + } + + #[test] + fn parse_shell_syntax_with_limit() { + let result = parse_mongo_filter(r#"db.users.find({"role": "admin"}).limit(10)"#).unwrap(); + assert_eq!(result.get_str("role").unwrap(), "admin"); + } + + #[test] + fn parse_relaxed_json_single_quotes() { + let result = parse_mongo_filter("{'name': 'John'}").unwrap(); + assert_eq!(result.get_str("name").unwrap(), "John"); + } + + #[test] + fn parse_invalid_json_returns_empty() { + let result = parse_mongo_filter("not json at all").unwrap(); + assert!(result.is_empty()); + } + + #[test] + fn extract_simple_braced_json() { + let result = extract_braced_json(r#"{"key": "value"}"#); + assert_eq!(result.unwrap(), r#"{"key": "value"}"#); + } + + #[test] + fn extract_nested_braced_json() { + let result = extract_braced_json(r#"{"outer": {"inner": true}}"#); + assert!(result.unwrap().contains("inner")); + } + + #[test] + fn extract_braced_json_not_starting_with_brace_returns_none() { + let result = extract_braced_json("not a json"); + assert!(result.is_none()); + } + + #[test] + fn bson_string_display() { + let value = Bson::String("hello".to_string()); + assert_eq!(bson_to_display_string(&value).unwrap(), "hello"); + } + + #[test] + fn bson_int32_display() { + let value = Bson::Int32(42); + assert_eq!(bson_to_display_string(&value).unwrap(), "42"); + } + + #[test] + fn bson_int64_display() { + let value = Bson::Int64(9007199254740991); + assert_eq!(bson_to_display_string(&value).unwrap(), "9007199254740991"); + } + + #[test] + #[allow(clippy::approx_constant)] + fn bson_double_display() { + let value = Bson::Double(3.14); + assert_eq!(bson_to_display_string(&value).unwrap(), "3.14"); + } + + #[test] + fn bson_bool_display() { + assert_eq!(bson_to_display_string(&Bson::Boolean(true)).unwrap(), "true"); + assert_eq!(bson_to_display_string(&Bson::Boolean(false)).unwrap(), "false"); + } + + #[test] + fn bson_objectid_display() { + let oid = ObjectId::new(); + let value = Bson::ObjectId(oid); + let displayed = bson_to_display_string(&value).unwrap(); + assert_eq!(displayed.len(), 24); + assert!(displayed.chars().all(|c| c.is_ascii_hexdigit())); + } + + #[test] + fn bson_null_display() { + assert!(bson_to_display_string(&Bson::Null).is_none()); + } + + #[test] + fn bson_array_display() { + let value = Bson::Array(vec![Bson::Int32(1), Bson::Int32(2), Bson::Int32(3)]); + let result = bson_to_display_string(&value).unwrap(); + assert!(result.contains("3 items")); + } + + #[test] + fn bson_document_display() { + let value = Bson::Document(doc! { "key": "val" }); + assert_eq!(bson_to_display_string(&value).unwrap(), "{...}"); + } + + #[test] + fn bson_binary_display() { + use mongodb::bson::Binary; + let value = Bson::Binary(Binary { + subtype: mongodb::bson::spec::BinarySubtype::Generic, + bytes: vec![0, 1, 2], + }); + assert_eq!(bson_to_display_string(&value).unwrap(), ""); + } + + #[test] + fn bson_datetime_display() { + use chrono::TimeZone; + let dt = chrono::Utc.with_ymd_and_hms(2024, 3, 15, 10, 30, 45).unwrap(); + let bson_dt = BsonDateTime::from_millis(dt.timestamp_millis()); + let value = Bson::DateTime(bson_dt); + let result = bson_to_display_string(&value).unwrap(); + assert!(result.contains("2024-03-15")); + } + + #[test] + fn infer_bson_type_string() { + assert_eq!(infer_bson_type(&Bson::String("hello".to_string())), "string"); + } + + #[test] + fn infer_bson_type_integer() { + assert_eq!(infer_bson_type(&Bson::Int32(1)), "integer"); + assert_eq!(infer_bson_type(&Bson::Int64(1)), "integer"); + } + + #[test] + fn infer_bson_type_double() { + assert_eq!(infer_bson_type(&Bson::Double(1.0)), "double"); + } + + #[test] + fn infer_bson_type_boolean() { + assert_eq!(infer_bson_type(&Bson::Boolean(true)), "boolean"); + } + + #[test] + fn infer_bson_type_objectid() { + assert_eq!(infer_bson_type(&Bson::ObjectId(ObjectId::new())), "objectId"); + } + + #[test] + fn infer_bson_type_date() { + let dt = BsonDateTime::from_millis(0); + assert_eq!(infer_bson_type(&Bson::DateTime(dt)), "date"); + } + + #[test] + fn infer_bson_type_array() { + assert_eq!(infer_bson_type(&Bson::Array(vec![])), "array"); + } + + #[test] + fn infer_bson_type_document() { + assert_eq!(infer_bson_type(&Bson::Document(doc! {})), "object"); + } + + #[test] + fn infer_bson_type_null() { + assert_eq!(infer_bson_type(&Bson::Null), "null"); + } + + #[test] + fn infer_bson_type_regex() { + use mongodb::bson::Regex; + assert_eq!(infer_bson_type(&Bson::RegularExpression(Regex { + pattern: ".*".to_string(), + options: "".to_string(), + })), "regex"); + } + + // ── Cross-engine QueryResult compatibility ────────────────── + + #[test] + fn query_result_structure_is_consistent_across_engines() { + use crate::models::QueryResult; + use std::collections::BTreeMap; + + // Simulate a SQL result (Postgres/MySQL/SQLite format) + let sql_result = QueryResult { + columns: vec!["id".to_string(), "name".to_string()], + rows: vec![ + BTreeMap::from([ + ("id".to_string(), Some("1".to_string())), + ("name".to_string(), Some("Alice".to_string())), + ]), + ], + row_count: 1, + execution_ms: 15, + truncated: false, + command_tag: None, + }; + + // Simulate a MongoDB result (should match exactly the same shape) + let mongo_result = QueryResult { + columns: vec!["_id".to_string(), "name".to_string()], + rows: vec![ + BTreeMap::from([ + ("_id".to_string(), Some("507f1f77bcf86cd799439011".to_string())), + ("name".to_string(), Some("Alice".to_string())), + ]), + ], + row_count: 1, + execution_ms: 18, + truncated: false, + command_tag: None, + }; + + // Both have same shape (structurally identical) + assert_eq!(sql_result.columns.len(), 2); + assert_eq!(mongo_result.columns.len(), 2); + assert_eq!(sql_result.rows.len(), 1); + assert_eq!(mongo_result.rows.len(), 1); + assert!(!sql_result.truncated); + assert!(!mongo_result.truncated); + + // Both are serializable (Tauri IPC sends JSON) + let sql_json = serde_json::to_string(&sql_result).unwrap(); + let mongo_json = serde_json::to_string(&mongo_result).unwrap(); + assert!(sql_json.contains("\"columns\"")); + assert!(mongo_json.contains("\"columns\"")); + assert!(sql_json.contains("\"rows\"")); + assert!(mongo_json.contains("\"rows\"")); + } + + #[test] + fn empty_mongo_result_matches_sql_format() { + use crate::models::QueryResult; + + let empty = QueryResult { + columns: vec![], + rows: vec![], + row_count: 0, + execution_ms: 0, + truncated: false, + command_tag: None, + }; + + // An empty query result should be structurally identical + // regardless of engine — this is what the frontend expects. + assert_eq!(empty.columns.len(), 0); + assert_eq!(empty.rows.len(), 0); + assert_eq!(empty.row_count, 0); + assert!(!empty.truncated); + } + + #[test] + fn mongo_null_values_match_sql_null() { + use std::collections::BTreeMap; + + // SQL NULL and MongoDB null/absent field both map to Option::None + // Both engines produce the same format for the frontend grid. + + let sql_row = BTreeMap::from([ + ("name".to_string(), Some("Alice".to_string())), + ("email".to_string(), None), // SQL NULL + ]); + + let mongo_row = BTreeMap::from([ + ("name".to_string(), Some("Bob".to_string())), + ("email".to_string(), None), // MongoDB field missing or null + ]); + + // Both use Option for values — the frontend treats None as "NULL" + assert_eq!(sql_row.get("email").unwrap(), &None); + assert_eq!(mongo_row.get("email").unwrap(), &None); + assert_eq!(sql_row.get("name").unwrap(), &Some("Alice".to_string())); + assert_eq!(mongo_row.get("name").unwrap(), &Some("Bob".to_string())); + } + + #[test] + fn mongo_dynamic_columns_produces_tabular_rows() { + use std::collections::BTreeMap; + + // When MongoDB documents have different fields, we union the keys + // and fill missing columns with None (same as SQL NULL). + + let columns = vec!["name".to_string(), "extra".to_string()]; + let rows = vec![ + BTreeMap::from([ + ("name".to_string(), Some("doc1".to_string())), + ("extra".to_string(), Some("value1".to_string())), + ]), + BTreeMap::from([ + ("name".to_string(), Some("doc2".to_string())), + ("extra".to_string(), None), // this doc lacks the field + ]), + ]; + + // Every row should have an entry for every column (even if None) + for row in &rows { + for col in &columns { + assert!(row.contains_key(col), "row missing column: {}", col); + } + } + } +} diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs new file mode 100644 index 0000000..2ed85ed --- /dev/null +++ b/src-tauri/src/commands/query.rs @@ -0,0 +1,58 @@ +use tauri::{AppHandle, State}; + +use crate::db::{resolve_connection_engine, AppState, MAX_QUERY_ROWS}; +use crate::engines; +use crate::engines::DatabaseEngineOps; +use crate::models::{ColumnInfo, QueryRequest, QueryResult, SchemaRequest, TableInfo}; + +use super::is_read_only_sql; + +#[tauri::command] +pub async fn run_query( + app: AppHandle, + state: State<'_, AppState>, + input: QueryRequest, +) -> Result { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id).await?; + let sql = input.sql.trim().to_string(); + + if sql.is_empty() { + return Err("Enter a SQL statement before running the query.".to_string()); + } + + if !input.allow_write.unwrap_or(false) && !is_read_only_sql(&sql) { + return Err( + "This statement modifies data or schema. Confirm execution in the editor, \ + or use the model/DDL workflow for schema changes.".to_string(), + ); + } + + let max_query_rows = input.max_rows.unwrap_or(MAX_QUERY_ROWS); + + let ops = engines::get_engine(engine); + ops.run_query(&app, &state, &connection_id, &sql, max_query_rows).await.map_err(String::from) +} + +#[tauri::command] +pub async fn get_tables( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; + + let ops = engines::get_engine(engine); + ops.get_tables(&app, &state, &connection_id).await.map_err(String::from) +} + +#[tauri::command] +pub async fn get_schema( + app: AppHandle, + state: State<'_, AppState>, + input: SchemaRequest, +) -> Result, String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + + let ops = engines::get_engine(engine); + ops.get_schema(&app, &state, &connection_id, &input.table_schema, &input.table_name).await.map_err(String::from) +} diff --git a/src-tauri/src/commands/redis.rs b/src-tauri/src/commands/redis.rs new file mode 100644 index 0000000..5bd03af --- /dev/null +++ b/src-tauri/src/commands/redis.rs @@ -0,0 +1,132 @@ +use std::collections::BTreeMap; +use std::time::Instant; +use tauri::{AppHandle, State}; +use crate::db::{get_or_create_redis_client, resolve_connection_engine, AppState}; +use crate::models::{QueryRequest, QueryResult, TableInfo, ColumnInfo}; + +#[tauri::command] +pub async fn redis_run_query( + app: AppHandle, + state: State<'_, AppState>, + input: QueryRequest, +) -> Result { + let (connection_id, _) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let mut client = get_or_create_redis_client(&app, &state, &connection_id).await?; + let raw = input.sql.trim().to_string(); + if raw.is_empty() { return Err("Enter a Redis command.".to_string()); } + let started_at = Instant::now(); + let parts: Vec<&str> = raw.split_whitespace().collect(); + let cmd = parts.first().map(|s| s.to_uppercase()).unwrap_or_default(); + + match cmd.as_str() { + "PING" => { + let result: String = redis::cmd("PING").query_async(&mut client).await.map_err(|e| e.to_string())?; + Ok(QueryResult { columns: vec!["result".to_string()], rows: vec![BTreeMap::from([("result".to_string(), Some(result))])], row_count: 1, execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: None }) + } + "GET" => { + let key = parts.get(1).unwrap_or(&""); + let result: Option = redis::cmd("GET").arg(key).query_async(&mut client).await.map_err(|e| e.to_string())?; + Ok(QueryResult { columns: vec!["value".to_string()], rows: vec![BTreeMap::from([("value".to_string(), result)])], row_count: 1, execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: None }) + } + "SET" => { + let key = parts.get(1).unwrap_or(&""); + let value = parts.get(2).unwrap_or(&""); + redis::cmd("SET").arg(key).arg(value).query_async::<_, ()>(&mut client).await.map_err(|e| e.to_string())?; + Ok(QueryResult { columns: vec![], rows: vec![], row_count: 0, execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: Some(0) }) + } + "KEYS" | "SCAN" => { + let pattern = parts.get(1).unwrap_or(&"*"); + let keys: Vec = redis::cmd("KEYS").arg(pattern).query_async(&mut client).await.map_err(|e| e.to_string())?; + let rows: Vec>> = keys.into_iter().map(|k| BTreeMap::from([("key".to_string(), Some(k))])).collect(); + Ok(QueryResult { columns: vec!["key".to_string()], rows: rows.clone(), row_count: rows.len(), execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: None }) + } + "HGETALL" => { + let key = parts.get(1).unwrap_or(&""); + let hash: Vec = redis::cmd("HGETALL").arg(key).query_async(&mut client).await.map_err(|e| e.to_string())?; + let mut row = BTreeMap::new(); + let mut cols = Vec::new(); + for chunk in hash.chunks(2) { + if chunk.len() == 2 { + cols.push(chunk[0].clone()); + row.insert(chunk[0].clone(), Some(chunk[1].clone())); + } + } + Ok(QueryResult { columns: cols, rows: vec![row], row_count: 1, execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: None }) + } + "LRANGE" => { + let key = parts.get(1).unwrap_or(&""); + let start: isize = parts.get(2).and_then(|s| s.parse().ok()).unwrap_or(0); + let stop: isize = parts.get(3).and_then(|s| s.parse().ok()).unwrap_or(-1); + let list: Vec = redis::cmd("LRANGE").arg(key).arg(start).arg(stop).query_async(&mut client).await.map_err(|e| e.to_string())?; + let rows: Vec>> = list.into_iter().enumerate().map(|(i, v)| { + let mut map = BTreeMap::new(); + map.insert("index".to_string(), Some(i.to_string())); + map.insert("value".to_string(), Some(v)); + map + }).collect(); + Ok(QueryResult { columns: vec!["index".to_string(), "value".to_string()], rows: rows.clone(), row_count: rows.len(), execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: None }) + } + "SMEMBERS" => { + let key = parts.get(1).unwrap_or(&""); + let members: Vec = redis::cmd("SMEMBERS").arg(key).query_async(&mut client).await.map_err(|e| e.to_string())?; + let rows: Vec>> = members.into_iter().map(|v| BTreeMap::from([("member".to_string(), Some(v))])).collect(); + Ok(QueryResult { columns: vec!["member".to_string()], rows: rows.clone(), row_count: rows.len(), execution_ms: started_at.elapsed().as_millis(), truncated: false, command_tag: None }) + } + _ => Err(format!("Unsupported Redis command: {}. Try GET, SET, KEYS, HGETALL, LRANGE, SMEMBERS, PING.", cmd)) + } +} + +#[tauri::command] +pub async fn redis_get_keys( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, _) = resolve_connection_engine(&app, &state, connection_id).await?; + let mut client = get_or_create_redis_client(&app, &state, &connection_id).await?; + let keys: Vec = redis::cmd("KEYS").arg("*").query_async(&mut client).await.map_err(|e| e.to_string())?; + Ok(keys.into_iter().map(|k| TableInfo { schema: "0".to_string(), name: k.clone(), preview_query: format!("GET {}", k) }).collect()) +} + +/// Infer schema for a Redis key by checking its type and sampling data. +#[tauri::command] +pub async fn redis_get_schema( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, + _table_schema: String, + table_name: String, +) -> Result, String> { + let (connection_id, _) = resolve_connection_engine(&app, &state, connection_id).await?; + let mut client = get_or_create_redis_client(&app, &state, &connection_id).await?; + + let key_type: String = redis::cmd("TYPE").arg(&table_name).query_async(&mut client) + .await.map_err(|e| format!("Redis TYPE failed: {}", e))?; + + match key_type.as_str() { + "string" => Ok(vec![ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.clone(), + column_name: "value".to_string(), data_type: "string".to_string(), is_nullable: true, + }]), + "hash" => { + let fields: Vec = redis::cmd("HKEYS").arg(&table_name).query_async(&mut client) + .await.map_err(|e| format!("Redis HKEYS failed: {}", e))?; + Ok(fields.into_iter().map(|f| ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.clone(), + column_name: f, data_type: "string".to_string(), is_nullable: true, + }).collect()) + } + "list" => Ok(vec![ + ColumnInfo { table_schema: "0".to_string(), table_name: table_name.clone(), column_name: "index".to_string(), data_type: "integer".to_string(), is_nullable: false }, + ColumnInfo { table_schema: "0".to_string(), table_name: table_name.clone(), column_name: "value".to_string(), data_type: "string".to_string(), is_nullable: true }, + ]), + "set" => Ok(vec![ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.clone(), + column_name: "member".to_string(), data_type: "string".to_string(), is_nullable: false, + }]), + _ => Ok(vec![ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.clone(), + column_name: "value".to_string(), data_type: key_type, is_nullable: true, + }]), + } +} diff --git a/src-tauri/src/commands/table_props.rs b/src-tauri/src/commands/table_props.rs new file mode 100644 index 0000000..e743414 --- /dev/null +++ b/src-tauri/src/commands/table_props.rs @@ -0,0 +1,619 @@ +use std::collections::{HashMap, HashSet}; +use tauri::{AppHandle, State}; + +use crate::db::{ + get_or_create_mysql_pool, get_or_create_sqlite_pool, quote_identifier, + require_safe_identifier, resolve_connection_engine, with_pool_client_retry, AppState, +}; +use crate::models::{ + ColumnProperties, DatabaseEngine, ForeignKeyEdge, + IndexInfo, SchemaRequest, TableIndexesResult, TablePropertiesApplyRequest, +}; +use crate::pg_error::map_pg_err; + +use super::{ + MAX_TABLE_INDEX_ROWS, MAX_FOREIGN_KEY_ROWS, + mysql_get_string, mysql_get_optional_string, mysql_get_idx, + sqlite_get_name, sqlite_get_idx, + veloxdb_unique_constraint_name, +}; + +#[tauri::command] +pub async fn get_table_properties( + app: AppHandle, + state: State<'_, AppState>, + input: SchemaRequest, +) -> Result, String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let ctx = input.clone(); + + if engine == DatabaseEngine::Mysql { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + let rows = sqlx::query( + "select c.table_schema, c.table_name, c.column_name, c.data_type, \ + c.is_nullable, c.column_default, c.extra \ + from information_schema.columns c \ + where c.table_schema = ? and c.table_name = ? \ + order by c.ordinal_position", + ).bind(&ctx.table_schema).bind(&ctx.table_name).fetch_all(&pool).await + .map_err(|error| error.to_string())?; + + let pk_rows = sqlx::query( + "select column_name from information_schema.key_column_usage \ + where table_schema = ? and table_name = ? and constraint_name = 'PRIMARY'", + ).bind(&ctx.table_schema).bind(&ctx.table_name).fetch_all(&pool).await + .map_err(|error| error.to_string())?; + let pk_cols: HashSet = pk_rows.into_iter() + .map(|row| mysql_get_string(&row, 0, "column_name", "get_table_properties")) + .collect::, _>>()?; + + let unique_rows = sqlx::query( + "select index_name, column_name, seq_in_index \ + from information_schema.statistics \ + where table_schema = ? and table_name = ? and non_unique = 0 \ + order by index_name, seq_in_index", + ).bind(&ctx.table_schema).bind(&ctx.table_name).fetch_all(&pool).await + .map_err(|error| error.to_string())?; + let mut unique_by_index: HashMap> = HashMap::new(); + for row in unique_rows { + let index_name: String = mysql_get_string(&row, 0, "index_name", "get_table_properties")?; + if index_name == "PRIMARY" { continue; } + let column_name: String = mysql_get_string(&row, 1, "column_name", "get_table_properties")?; + unique_by_index.entry(index_name).or_default().push(column_name); + } + let mut unique_cols: HashSet = HashSet::new(); + let mut composite_unique_cols: HashSet = HashSet::new(); + for cols in unique_by_index.values() { + for col in cols { unique_cols.insert(col.clone()); } + if cols.len() > 1 { for col in cols { composite_unique_cols.insert(col.clone()); } } + } + + let mut properties = Vec::new(); + for row in rows { + let column_name: String = mysql_get_string(&row, 2, "column_name", "get_table_properties")?; + let is_primary_key = pk_cols.contains(&column_name); + let is_unique = is_primary_key || unique_cols.contains(&column_name); + let is_part_of_composite_unique = composite_unique_cols.contains(&column_name); + let extra: String = mysql_get_string(&row, 6, "extra", "get_table_properties")?; + let lower_extra = extra.to_lowercase(); + properties.push(ColumnProperties { + table_schema: mysql_get_string(&row, 0, "table_schema", "get_table_properties")?, + table_name: mysql_get_string(&row, 1, "table_name", "get_table_properties")?, + column_name, + data_type: mysql_get_string(&row, 3, "data_type", "get_table_properties")?, + is_nullable: mysql_get_string(&row, 4, "is_nullable", "get_table_properties")? == "YES", + is_primary_key, is_unique, is_part_of_composite_unique, + column_default: mysql_get_optional_string(&row, 5, "column_default", "get_table_properties")?, + is_identity: lower_extra.contains("auto_increment"), + identity_generation: if lower_extra.contains("auto_increment") { Some("BY DEFAULT".to_string()) } else { None }, + is_generated: if lower_extra.contains("generated") { Some("ALWAYS".to_string()) } else { None }, + }); + } + return Ok(properties); + } + + if engine == DatabaseEngine::Sqlite { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + require_safe_identifier(&ctx.table_name, "table name")?; + let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(&ctx.table_name)); + let rows = sqlx::query(&pragma_sql).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let index_list_sql = format!("PRAGMA index_list(\"{}\");", quote_identifier(&ctx.table_name)); + let index_rows = sqlx::query(&index_list_sql).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let mut unique_cols: HashSet = HashSet::new(); + let mut composite_unique_cols: HashSet = HashSet::new(); + for index in index_rows { + let is_unique = sqlite_get_name::(&index, "unique", "get_table_properties")? == 1; + if !is_unique { continue; } + let origin = sqlite_get_name::(&index, "origin", "get_table_properties")?; + if origin == "pk" { continue; } + let index_name = sqlite_get_name::(&index, "name", "get_table_properties")?; + require_safe_identifier(&index_name, "index name")?; + let info_sql = format!("PRAGMA index_info(\"{}\");", quote_identifier(&index_name)); + let info_rows = sqlx::query(&info_sql).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let mut cols: Vec = Vec::new(); + for info in info_rows { + if let Ok(name) = sqlite_get_name::(&info, "name", "get_table_properties") { cols.push(name); } + } + for col in &cols { unique_cols.insert(col.clone()); } + if cols.len() > 1 { for col in cols { composite_unique_cols.insert(col); } } + } + let mut properties = Vec::new(); + for row in rows { + let column_name: String = sqlite_get_name(&row, "name", "get_table_properties")?; + let is_primary_key = sqlite_get_name::(&row, "pk", "get_table_properties")? == 1; + let is_unique = is_primary_key || unique_cols.contains(&column_name); + let is_part_of_composite_unique = composite_unique_cols.contains(&column_name); + properties.push(ColumnProperties { + table_schema: "main".to_string(), + table_name: ctx.table_name.clone(), + column_name, + data_type: sqlite_get_name(&row, "type", "get_table_properties")?, + is_nullable: sqlite_get_name::(&row, "notnull", "get_table_properties")? == 0, + is_primary_key, is_unique, is_part_of_composite_unique, + column_default: sqlite_get_name::>(&row, "dflt_value", "get_table_properties")?, + is_identity: false, + identity_generation: None, + is_generated: None, + }); + } + return Ok(properties); + } + + if engine == DatabaseEngine::Duckdb { + return crate::commands::duckdb::duckdb_get_table_properties(app, state, ctx).await; + } + + if engine == DatabaseEngine::Mongo || engine == DatabaseEngine::Redis { + return Err("Table properties are not supported for this engine type.".to_string()); + } + + with_pool_client_retry(&app, &state, &connection_id, ctx, |client, input| async move { + let columns = client.query( + "select c.table_schema, c.table_name, c.column_name, c.data_type, \ + c.is_nullable, c.column_default, c.is_identity, c.identity_generation, c.is_generated \ + from information_schema.columns c \ + where c.table_schema = $1 and c.table_name = $2 \ + order by c.ordinal_position", + &[&input.table_schema, &input.table_name], + ).await.map_err(|error| map_pg_err(error, None))?; + + let primary_keys = client.query( + "select kcu.column_name \ + from information_schema.table_constraints tc \ + join information_schema.key_column_usage kcu \ + on tc.constraint_name = kcu.constraint_name \ + and tc.table_schema = kcu.table_schema \ + where tc.table_schema = $1 and tc.table_name = $2 \ + and tc.constraint_type = 'PRIMARY KEY' \ + order by kcu.ordinal_position", + &[&input.table_schema, &input.table_name], + ).await.map_err(|error| map_pg_err(error, None))?; + + let primary_key_columns: HashSet = primary_keys.into_iter() + .map(|row| row.get::<_, String>(0)).collect(); + + let unique_constraints = client.query( + "select tc.constraint_name, kcu.column_name, kcu.ordinal_position \ + from information_schema.table_constraints tc \ + join information_schema.key_column_usage kcu \ + on tc.constraint_name = kcu.constraint_name \ + and tc.table_schema = kcu.table_schema \ + where tc.table_schema = $1 and tc.table_name = $2 \ + and tc.constraint_type = 'UNIQUE' \ + order by tc.constraint_name, kcu.ordinal_position", + &[&input.table_schema, &input.table_name], + ).await.map_err(|error| map_pg_err(error, None))?; + + let mut unique_by_name: HashMap> = HashMap::new(); + for row in unique_constraints { + let constraint_name: String = row.get(0); + let column_name: String = row.get(1); + unique_by_name.entry(constraint_name).or_default().push(column_name); + } + + let mut unique_columns: HashSet = HashSet::new(); + let mut composite_unique_columns: HashSet = HashSet::new(); + for (_constraint_name, cols) in unique_by_name { + for c in &cols { unique_columns.insert(c.clone()); } + if cols.len() > 1 { for c in &cols { composite_unique_columns.insert(c.clone()); } } + } + + Ok(columns.into_iter().map(|row| { + let column_name: String = row.get(2); + let is_primary_key = primary_key_columns.contains(&column_name); + let is_unique = is_primary_key || unique_columns.contains(&column_name); + let is_part_of_composite_unique = composite_unique_columns.contains(&column_name); + ColumnProperties { + table_schema: row.get(0), + table_name: row.get(1), + column_name, + data_type: row.get(3), + is_nullable: row.get::<_, String>(4) == "YES", + is_primary_key, is_unique, is_part_of_composite_unique, + column_default: row.get(5), + is_identity: row.get::<_, Option>(6).as_deref() == Some("YES"), + identity_generation: row.get(7), + is_generated: row.get(8), + } + }).collect()) + }).await.map_err(String::from) +} + +#[tauri::command] +pub async fn apply_table_properties( + app: AppHandle, + state: State<'_, AppState>, + input: TablePropertiesApplyRequest, +) -> Result<(), String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + if engine != DatabaseEngine::Postgres && engine != DatabaseEngine::Duckdb { + return Err(format!( + "Table property editing is not supported for {} connections yet.", + match engine { + DatabaseEngine::Postgres => "PostgreSQL", + DatabaseEngine::Mysql => "MySQL", + DatabaseEngine::Sqlite => "SQLite", + DatabaseEngine::Mongo => "MongoDB", + DatabaseEngine::Duckdb => "DuckDB", + DatabaseEngine::Redis => "Redis", + } + )); + } + + if engine == DatabaseEngine::Duckdb { + return crate::commands::duckdb::duckdb_apply_table_properties(app, state, input).await; + } + + with_pool_client_retry(&app, &state, &connection_id, input, |mut client, input| async move { + let table_schema = input.table_schema; + let table_name = input.table_name; + let columns = input.columns; + + require_safe_identifier(&table_schema, "schema name")?; + require_safe_identifier(&table_name, "table name")?; + + let current_columns = client.query( + "select column_name, is_nullable \ + from information_schema.columns \ + where table_schema = $1 and table_name = $2", + &[&table_schema, &table_name], + ).await.map_err(|error| map_pg_err(error, None))?; + + let mut current_nullable: HashMap = HashMap::new(); + for row in current_columns { + let column_name: String = row.get(0); + let is_nullable = row.get::<_, String>(1) == "YES"; + current_nullable.insert(column_name, is_nullable); + } + + let primary_keys = client.query( + "select kcu.column_name \ + from information_schema.table_constraints tc \ + join information_schema.key_column_usage kcu \ + on tc.constraint_name = kcu.constraint_name \ + and tc.table_schema = kcu.table_schema \ + where tc.table_schema = $1 and tc.table_name = $2 \ + and tc.constraint_type = 'PRIMARY KEY'", + &[&table_schema, &table_name], + ).await.map_err(|error| map_pg_err(error, None))?; + + let primary_key_columns: HashSet = primary_keys.into_iter() + .map(|row| row.get::<_, String>(0)).collect(); + + let unique_constraints = client.query( + "select tc.constraint_name, kcu.column_name, kcu.ordinal_position \ + from information_schema.table_constraints tc \ + join information_schema.key_column_usage kcu \ + on tc.constraint_name = kcu.constraint_name \ + and tc.table_schema = kcu.table_schema \ + where tc.table_schema = $1 and tc.table_name = $2 \ + and tc.constraint_type = 'UNIQUE' \ + order by tc.constraint_name, kcu.ordinal_position", + &[&table_schema, &table_name], + ).await.map_err(|error| map_pg_err(error, None))?; + + let mut unique_by_name: HashMap> = HashMap::new(); + for row in unique_constraints { + let constraint_name: String = row.get(0); + let column_name: String = row.get(1); + unique_by_name.entry(constraint_name).or_default().push(column_name); + } + + let mut composite_unique_columns: HashSet = HashSet::new(); + let mut single_unique_constraint_names_by_column: HashMap> = HashMap::new(); + for (constraint_name, cols) in &unique_by_name { + if cols.len() > 1 { + for c in cols { composite_unique_columns.insert(c.clone()); } + } else if cols.len() == 1 { + let c = &cols[0]; + single_unique_constraint_names_by_column.entry(c.clone()).or_default().push(constraint_name.clone()); + } + } + + let mut desired_by_column: HashMap = HashMap::new(); + for update in columns { + desired_by_column.insert(update.column_name, (update.is_nullable, update.is_unique)); + } + + let txn = client.transaction().await.map_err(|error| map_pg_err(error, None))?; + + for (column_name, (desired_is_nullable, _desired_is_unique)) in &desired_by_column { + let current_is_nullable = current_nullable.get(column_name) + .ok_or_else(|| format!("Unknown column: {}", column_name))?; + if *current_is_nullable == *desired_is_nullable { continue; } + + let qualified_table = format!("\"{}\".\"{}\"", quote_identifier(&table_schema), quote_identifier(&table_name)); + require_safe_identifier(column_name, "column name")?; + let qualified_column = format!("\"{}\"", quote_identifier(column_name)); + + if *desired_is_nullable { + let sql = format!("ALTER TABLE {} ALTER COLUMN {} DROP NOT NULL", qualified_table, qualified_column); + txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + } else { + let sql = format!("ALTER TABLE {} ALTER COLUMN {} SET NOT NULL", qualified_table, qualified_column); + txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + } + } + + for (column_name, (_desired_is_nullable, desired_is_unique)) in &desired_by_column { + let is_primary_key = primary_key_columns.contains(column_name); + let is_part_of_composite_unique = composite_unique_columns.contains(column_name); + + if !*desired_is_unique { + if is_primary_key { + return Err(format!("Cannot disable UNIQUE for primary key column: {}", column_name)); + } + if is_part_of_composite_unique { + return Err(format!("Cannot disable UNIQUE for column in a composite UNIQUE constraint: {}", column_name)); + } + } + + let has_single_unique = single_unique_constraint_names_by_column.get(column_name) + .map(|names| !names.is_empty()).unwrap_or(false); + let current_is_unique = is_primary_key || has_single_unique || is_part_of_composite_unique; + + if *desired_is_unique == current_is_unique { continue; } + + let qualified_table = format!("\"{}\".\"{}\"", quote_identifier(&table_schema), quote_identifier(&table_name)); + require_safe_identifier(column_name, "column name")?; + let qualified_column = format!("\"{}\"", quote_identifier(column_name)); + + if *desired_is_unique { + if current_is_unique { continue; } + let generated_name = veloxdb_unique_constraint_name(&table_name, column_name); + if let Some(existing_cols) = unique_by_name.get(&generated_name) { + if existing_cols.len() != 1 || existing_cols[0] != *column_name { + return Err(format!( + "Cannot create UNIQUE constraint due to name collision ({}). Rename the existing constraint.", + generated_name + )); + } + } + let sql = format!("ALTER TABLE {} ADD CONSTRAINT \"{}\" UNIQUE ({})", + qualified_table, quote_identifier(&generated_name), qualified_column); + txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + } else { + let constraint_names = single_unique_constraint_names_by_column.get(column_name) + .cloned().unwrap_or_default(); + for constraint_name in constraint_names { + require_safe_identifier(&constraint_name, "constraint name")?; + let sql = format!("ALTER TABLE {} DROP CONSTRAINT \"{}\"", + qualified_table, quote_identifier(&constraint_name)); + txn.execute(sql.as_str(), &[]).await.map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + } + } + } + + txn.commit().await.map_err(|error| map_pg_err(error, None))?; + Ok(()) + }).await.map_err(String::from) +} + +#[tauri::command] +pub async fn get_foreign_keys( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result, String> { + let (connection_id, engine) = resolve_connection_engine(&app, &state, connection_id).await?; + + if engine == DatabaseEngine::Mysql { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + let rows = sqlx::query( + "select kcu.table_schema as from_schema, kcu.table_name as from_table, \ + kcu.column_name as from_column, kcu.referenced_table_schema as to_schema, \ + kcu.referenced_table_name as to_table, kcu.referenced_column_name as to_column \ + from information_schema.key_column_usage kcu \ + where kcu.referenced_table_name is not null \ + order by kcu.table_schema, kcu.table_name, kcu.ordinal_position \ + limit ?", + ).bind(MAX_FOREIGN_KEY_ROWS).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let mut edges = Vec::new(); + for row in rows { + edges.push(ForeignKeyEdge { + from_schema: mysql_get_string(&row, 0, "from_schema", "get_foreign_keys")?, + from_table: mysql_get_string(&row, 1, "from_table", "get_foreign_keys")?, + from_column: mysql_get_string(&row, 2, "from_column", "get_foreign_keys")?, + to_schema: mysql_get_string(&row, 3, "to_schema", "get_foreign_keys")?, + to_table: mysql_get_string(&row, 4, "to_table", "get_foreign_keys")?, + to_column: mysql_get_string(&row, 5, "to_column", "get_foreign_keys")?, + }); + } + return Ok(edges); + } + + if engine == DatabaseEngine::Sqlite { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + let tables = sqlx::query( + "select name from sqlite_master \ + where type = 'table' and name not like 'sqlite_%'", + ).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let mut edges = Vec::new(); + for table in tables { + let table_name: String = sqlite_get_idx(&table, 0, "name", "get_foreign_keys")?; + require_safe_identifier(&table_name, "table name")?; + let fk_sql = format!("PRAGMA foreign_key_list(\"{}\");", quote_identifier(&table_name)); + let fk_rows = sqlx::query(&fk_sql).fetch_all(&pool).await.map_err(|error| error.to_string())?; + for row in fk_rows { + edges.push(ForeignKeyEdge { + from_schema: "main".to_string(), from_table: table_name.clone(), + from_column: sqlite_get_name(&row, "from", "get_foreign_keys")?, + to_schema: "main".to_string(), + to_table: sqlite_get_name(&row, "table", "get_foreign_keys")?, + to_column: sqlite_get_name(&row, "to", "get_foreign_keys")?, + }); + if edges.len() >= MAX_FOREIGN_KEY_ROWS as usize { return Ok(edges); } + } + } + return Ok(edges); + } + + if engine == DatabaseEngine::Duckdb { + return crate::commands::duckdb::duckdb_get_foreign_keys(app, state, Some(connection_id)).await; + } + + if engine == DatabaseEngine::Mongo || engine == DatabaseEngine::Redis { + return Ok(Vec::new()); + } + + with_pool_client_retry(&app, &state, &connection_id, (), |client, ()| async move { + let rows = client.query( + "select src_ns.nspname::text as from_schema, src_cls.relname::text as from_table, \ + src_att.attname::text as from_column, tgt_ns.nspname::text as to_schema, \ + tgt_cls.relname::text as to_table, tgt_att.attname::text as to_column \ + from pg_constraint c \ + join pg_class src_cls on src_cls.oid = c.conrelid \ + join pg_namespace src_ns on src_ns.oid = src_cls.relnamespace \ + join pg_class tgt_cls on tgt_cls.oid = c.confrelid \ + join pg_namespace tgt_ns on tgt_ns.oid = tgt_cls.relnamespace \ + cross join lateral unnest(c.conkey, c.confkey) as u(attnum, confattnum) \ + join pg_attribute src_att on src_att.attrelid = c.conrelid \ + and src_att.attnum = u.attnum and not src_att.attisdropped \ + join pg_attribute tgt_att on tgt_att.attrelid = c.confrelid \ + and tgt_att.attnum = u.confattnum and not tgt_att.attisdropped \ + where c.contype = 'f' \ + and src_ns.nspname not in ('pg_catalog', 'information_schema') \ + order by src_ns.nspname, src_cls.relname, c.conname, u.attnum \ + limit $1", + &[&MAX_FOREIGN_KEY_ROWS], + ).await.map_err(|error| error.to_string())?; + + Ok(rows.into_iter().map(|row| ForeignKeyEdge { + from_schema: row.get(0), from_table: row.get(1), from_column: row.get(2), + to_schema: row.get(3), to_table: row.get(4), to_column: row.get(5), + }).collect()) + }).await.map_err(String::from) +} + +#[tauri::command] +pub async fn get_table_indexes( + app: AppHandle, + state: State<'_, AppState>, + input: SchemaRequest, +) -> Result { + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let ctx = input.clone(); + + if engine == DatabaseEngine::Mysql { + let pool = get_or_create_mysql_pool(&app, &state, &connection_id).await?; + let fetch_limit = MAX_TABLE_INDEX_ROWS + 1; + let rows = sqlx::query( + "select table_schema as index_schema, index_name, table_schema, table_name, \ + non_unique = 0 as is_unique, index_name = 'PRIMARY' as is_primary, \ + true as is_valid, false as is_partial, \ + concat(index_name, ' (', group_concat(column_name order by seq_in_index separator ', '), ')') as definition, \ + 0 as index_bytes, 0 as idx_scan, 0 as idx_tup_read, 0 as idx_tup_fetch \ + from information_schema.statistics \ + where table_schema = ? and table_name = ? \ + group by table_schema, table_name, index_name, non_unique \ + order by index_name limit ?", + ).bind(&ctx.table_schema).bind(&ctx.table_name).bind(fetch_limit) + .fetch_all(&pool).await.map_err(|error| error.to_string())?; + let truncated = rows.len() as i64 > MAX_TABLE_INDEX_ROWS; + let mut indexes = Vec::new(); + for row in rows.into_iter().take(MAX_TABLE_INDEX_ROWS as usize) { + indexes.push(IndexInfo { + index_schema: mysql_get_string(&row, 0, "index_schema", "get_table_indexes")?, + index_name: mysql_get_string(&row, 1, "index_name", "get_table_indexes")?, + table_schema: mysql_get_string(&row, 2, "table_schema", "get_table_indexes")?, + table_name: mysql_get_string(&row, 3, "table_name", "get_table_indexes")?, + is_unique: mysql_get_idx(&row, 4, "is_unique", "get_table_indexes")?, + is_primary: mysql_get_idx(&row, 5, "is_primary", "get_table_indexes")?, + is_valid: true, is_partial: false, + definition: mysql_get_string(&row, 8, "definition", "get_table_indexes")?, + index_bytes: 0, idx_scan: 0, idx_tup_read: 0, idx_tup_fetch: 0, + }); + } + return Ok(TableIndexesResult { indexes, truncated }); + } + + if engine == DatabaseEngine::Sqlite { + let pool = get_or_create_sqlite_pool(&app, &state, &connection_id).await?; + require_safe_identifier(&ctx.table_name, "table name")?; + let pragma_sql = format!("PRAGMA index_list(\"{}\");", quote_identifier(&ctx.table_name)); + let rows = sqlx::query(&pragma_sql).fetch_all(&pool).await.map_err(|error| error.to_string())?; + let truncated = rows.len() as i64 > MAX_TABLE_INDEX_ROWS; + let mut indexes = Vec::new(); + for row in rows.into_iter().take(MAX_TABLE_INDEX_ROWS as usize) { + let index_name: String = sqlite_get_name(&row, "name", "get_table_indexes")?; + require_safe_identifier(&index_name, "index name")?; + let index_info_sql = format!("PRAGMA index_info(\"{}\");", quote_identifier(&index_name)); + let index_info_rows = sqlx::query(&index_info_sql).fetch_all(&pool).await + .map_err(|error| error.to_string())?; + let index_columns = index_info_rows.into_iter() + .filter_map(|idx| sqlite_get_name::(&idx, "name", "get_table_indexes").ok()) + .collect::>(); + indexes.push(IndexInfo { + index_schema: "main".to_string(), index_name: index_name.clone(), + table_schema: "main".to_string(), table_name: ctx.table_name.clone(), + is_unique: sqlite_get_name::(&row, "unique", "get_table_indexes")? == 1, + is_primary: sqlite_get_name::(&row, "origin", "get_table_indexes")? == "pk", + is_valid: true, + is_partial: sqlite_get_name::(&row, "partial", "get_table_indexes")? == 1, + definition: if index_columns.is_empty() { + format!("index {}", index_name) + } else { + format!("index {} ({})", index_name, index_columns.join(", ")) + }, + index_bytes: 0, idx_scan: 0, idx_tup_read: 0, idx_tup_fetch: 0, + }); + } + return Ok(TableIndexesResult { indexes, truncated }); + } + + if engine == DatabaseEngine::Duckdb { + return crate::commands::duckdb::duckdb_get_table_indexes(app, state, ctx).await; + } + + if engine == DatabaseEngine::Mongo { + return crate::commands::mongo::mongo_get_table_indexes(&app, &state, &connection_id, &ctx.table_schema, &ctx.table_name).await; + } + + if engine == DatabaseEngine::Redis { + return Ok(TableIndexesResult { indexes: Vec::new(), truncated: false }); + } + + with_pool_client_retry(&app, &state, &connection_id, ctx, |client, input| async move { + let table_schema = input.table_schema; + let table_name = input.table_name; + let fetch_limit = MAX_TABLE_INDEX_ROWS + 1; + + let rows = client.query( + "select ins.nspname::text as index_schema, ic.relname::text as index_name, \ + tn.nspname::text as table_schema, tc.relname::text as table_name, \ + i.indisunique as is_unique, i.indisprimary as is_primary, \ + i.indisvalid as is_valid, (i.indpred is not null) as is_partial, \ + pg_get_indexdef(i.indexrelid) as definition, \ + coalesce(pg_relation_size(i.indexrelid::regclass), 0)::bigint as index_bytes, \ + coalesce(s.idx_scan, 0)::bigint as idx_scan, \ + coalesce(s.idx_tup_read, 0)::bigint as idx_tup_read, \ + coalesce(s.idx_tup_fetch, 0)::bigint as idx_tup_fetch \ + from pg_index i \ + join pg_class ic on ic.oid = i.indexrelid \ + join pg_namespace ins on ins.oid = ic.relnamespace \ + join pg_class tc on tc.oid = i.indrelid \ + join pg_namespace tn on tn.oid = tc.relnamespace \ + left join pg_stat_user_indexes s on s.indexrelid = i.indexrelid \ + where tn.nspname = $1 and tc.relname = $2 \ + and ins.nspname not in ('pg_catalog', 'information_schema') \ + order by ic.relname limit $3", + &[&table_schema, &table_name, &fetch_limit], + ).await.map_err(|error| error.to_string())?; + + let truncated = rows.len() as i64 > MAX_TABLE_INDEX_ROWS; + let take = if truncated { MAX_TABLE_INDEX_ROWS as usize } else { rows.len() }; + let mut indexes = Vec::with_capacity(take); + for row in rows.into_iter().take(take) { + indexes.push(IndexInfo { + index_schema: row.get(0), index_name: row.get(1), + table_schema: row.get(2), table_name: row.get(3), + is_unique: row.get(4), is_primary: row.get(5), + is_valid: row.get(6), is_partial: row.get(7), + definition: row.get(8), + index_bytes: row.get(9), idx_scan: row.get(10), + idx_tup_read: row.get(11), idx_tup_fetch: row.get(12), + }); + } + Ok(TableIndexesResult { indexes, truncated }) + }).await.map_err(String::from) +} diff --git a/src-tauri/src/commands/updater.rs b/src-tauri/src/commands/updater.rs new file mode 100644 index 0000000..d39b57c --- /dev/null +++ b/src-tauri/src/commands/updater.rs @@ -0,0 +1,116 @@ +use serde::Deserialize; +use tauri::{AppHandle, State}; + +use crate::db::AppState; + +const GITHUB_REPO: &str = "abeni16/veloxdb"; + +#[derive(Debug, Clone, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateCheckResult { + pub current_version: String, + pub latest_version: String, + pub has_update: bool, + pub download_url: Option, + pub release_notes: Option, +} + +#[derive(Debug, Deserialize)] +struct GitHubRelease { + tag_name: String, + body: Option, + assets: Vec, +} + +#[derive(Debug, Deserialize)] +struct GitHubAsset { + name: String, + browser_download_url: String, +} + +#[tauri::command] +pub async fn check_for_updates( + app: AppHandle, + state: State<'_, AppState>, +) -> Result { + let current_version = app.package_info().version.to_string(); + + let client = state + .openrouter_client + .get() + .cloned() + .unwrap_or_else(|| { + reqwest::Client::builder() + .user_agent(format!("VeloxDB-Updater/{}", current_version)) + .build() + .expect("failed to build reqwest client") + }); + + let url = format!( + "https://api.github.com/repos/{}/releases/latest", + GITHUB_REPO + ); + + let response = client + .get(&url) + .header("Accept", "application/vnd.github.v3+json") + .send() + .await + .map_err(|e| format!("Failed to check for updates: {}", e))?; + + if !response.status().is_success() { + return Err(format!( + "GitHub API returned {}", + response.status().as_u16() + )); + } + + let release: GitHubRelease = response + .json() + .await + .map_err(|e| format!("Failed to parse release info: {}", e))?; + + let latest_tag = release.tag_name.trim_start_matches('v'); + + let current_ver = semver::Version::parse(¤t_version) + .map_err(|e| format!("Invalid current version: {}", e))?; + let latest_ver = semver::Version::parse(latest_tag) + .map_err(|e| format!("Invalid latest version '{}': {}", latest_tag, e))?; + + let has_update = latest_ver > current_ver; + + let download_url = find_platform_asset(&release.assets); + + Ok(UpdateCheckResult { + current_version, + latest_version: latest_tag.to_string(), + has_update, + download_url, + release_notes: release.body, + }) +} + +fn find_platform_asset(assets: &[GitHubAsset]) -> Option { + for asset in assets { + let name = &asset.name; + #[cfg(target_os = "macos")] + { + if name.ends_with(".dmg") { + return Some(asset.browser_download_url.clone()); + } + } + #[cfg(target_os = "windows")] + { + if name.ends_with(".msi") || name.ends_with(".exe") { + return Some(asset.browser_download_url.clone()); + } + } + #[cfg(target_os = "linux")] + { + if name.ends_with(".AppImage") || name.ends_with(".deb") { + return Some(asset.browser_download_url.clone()); + } + } + } + None +} diff --git a/src-tauri/src/commands/veloxy.rs b/src-tauri/src/commands/veloxy.rs new file mode 100644 index 0000000..63ba22d --- /dev/null +++ b/src-tauri/src/commands/veloxy.rs @@ -0,0 +1,440 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use serde_json::Value; +use tauri::{AppHandle, State}; + +use crate::db::{ + load_connection, resolve_connection_engine, AppState, MAX_QUERY_ROWS, +}; +use crate::models::{ + AskVeloxyChatRequest, AskVeloxyChatResponse, AskVeloxyConversationMessage, + AskVeloxyConversationResponse, AskVeloxyRequest, AskVeloxyResponse, + AskVeloxyTokenStats, VeloxyStreamChunk, +}; + +use super::{ + ask_veloxy_context_cache_key, ask_veloxy_conversation_key, + build_schema_context, classify_sql_intent, emit_veloxy_stream_chunk, + estimate_tokens, extract_openrouter_message_content, now_epoch_seconds, + normalize_openrouter_base, parse_ask_veloxy_chat_content, + parse_ask_veloxy_json, parse_ask_veloxy_suggestions, + stream_openrouter_chat_completion, truncate_on_char_boundary, validate_generated_sql, + ASK_VELOXY_MAX_HISTORY_MESSAGES, ASK_VELOXY_PROMPT_CHAR_BUDGET, +}; +use crate::models::DatabaseEngine; +use super::editor_meta::{ + fetch_foreign_keys_for_connection, fetch_query_editor_metadata_for_connection, +}; +use crate::models::AskVeloxyDbContextCache; + +async fn get_or_build_ask_veloxy_db_context( + app: &AppHandle, + state: &AppState, + connection_id: &str, + engine: DatabaseEngine, +) -> Result { + let stored_connection = load_connection(app, connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let cache_key = ask_veloxy_context_cache_key(connection_id, &stored_connection.database); + if let Some(cached) = state.ask_veloxy_db_context_cache.read().await.get(&cache_key).cloned() { + return Ok(cached); + } + + let metadata = fetch_query_editor_metadata_for_connection(app, state, connection_id, engine).await?; + let foreign_keys = fetch_foreign_keys_for_connection(app, state, connection_id, engine).await?; + let cache = AskVeloxyDbContextCache { + database_name: stored_connection.database, + engine, metadata, foreign_keys, + }; + state.ask_veloxy_db_context_cache.write().await.insert(cache_key, cache.clone()); + Ok(cache) +} + +#[tauri::command] +pub async fn cancel_veloxy_request(state: State<'_, AppState>) -> Result<(), String> { + if let Some(cancel) = state.veloxy_cancel.read().await.as_ref() { + cancel.store(true, Ordering::Relaxed); + } + Ok(()) +} + +#[tauri::command] +pub async fn chat_with_db( + app: AppHandle, + state: State<'_, AppState>, + input: AskVeloxyChatRequest, +) -> Result { + let natural_prompt = input.natural_prompt.trim(); + if natural_prompt.is_empty() { + return Err("Ask Veloxy prompt cannot be empty.".to_string()); + } + if input.provider_config.api_key.trim().is_empty() { + return Err("OpenRouter API key is required.".to_string()); + } + if input.provider_config.model.trim().is_empty() { + return Err("OpenRouter model is required.".to_string()); + } + + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let stored_connection = load_connection(&app, &connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let db_context = get_or_build_ask_veloxy_db_context(&app, &state, &connection_id, engine).await?; + let schema_context = build_schema_context(&db_context, natural_prompt, input.target_table.as_ref()); + let conversation_key = ask_veloxy_conversation_key(&connection_id, &stored_connection.database); + let history = state.ask_veloxy_conversations.read().await.get(&conversation_key).cloned().unwrap_or_default(); + + let history_block = history.iter().rev().take(8).rev() + .map(|message| format!("{}: {}", message.role, message.text)) + .collect::>().join("\n"); + + let mut user_prompt = format!( + "Engine: {:?}\nDatabase: {}\nTask: {}\nMaxRows: {}\nRecentConversation:\n{}\nSchemaContext:\n{}\n", + db_context.engine, db_context.database_name, natural_prompt, + input.max_rows.unwrap_or(MAX_QUERY_ROWS), history_block, schema_context + ); + truncate_on_char_boundary(&mut user_prompt, ASK_VELOXY_PROMPT_CHAR_BUDGET); + + let system_prompt = "You are Ask Veloxy chat mode. Return JSON when possible with keys: message (string), suggestions (array of strings), sqlDraft (string optional), needsSqlGeneration (boolean), needsClarification (boolean), warnings (array of strings). If JSON is not possible, return helpful plain text."; + let base_url = normalize_openrouter_base(input.provider_config.base_url.as_deref()); + let endpoint = format!("{}/chat/completions", base_url); + let client = state.openrouter_client.get_or_init(reqwest::Client::new); + let request_id = input.request_id.clone() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| format!("req-{}", uuid::Uuid::new_v4())); + + let cancel = Arc::new(AtomicBool::new(false)); + { + let mut guard = state.veloxy_cancel.write().await; + *guard = Some(cancel.clone()); + } + + let (message_content, hit_token_limit) = stream_openrouter_chat_completion( + &app, client, &endpoint, + input.provider_config.api_key.trim(), + input.provider_config.model.trim(), + system_prompt, &user_prompt, &request_id, cancel.clone(), + ).await?; + + { + let mut guard = state.veloxy_cancel.write().await; + *guard = None; + } + + let (message, suggestions, mut warnings, sql_draft, needs_sql_generation, needs_clarification) = + parse_ask_veloxy_chat_content(&message_content); + + if cancel.load(Ordering::Relaxed) { warnings.push("Stopped early.".to_string()); } + if hit_token_limit { + warnings.push(format!( + "Response may be truncated (model output limit of {} tokens).", + super::ASK_VELOXY_MAX_CHAT_TOKENS + )); + } + + emit_veloxy_stream_chunk(&app, VeloxyStreamChunk { + request_id: request_id.clone(), + delta: String::new(), + done: true, + message: Some(message.clone()), + suggestions: suggestions.clone(), + warnings: warnings.clone(), + sql_draft: sql_draft.clone(), + needs_sql_generation, + needs_clarification, + }); + + { + let mut conversations = state.ask_veloxy_conversations.write().await; + let bucket = conversations.entry(conversation_key).or_default(); + bucket.push(AskVeloxyConversationMessage { + id: format!("msg-{}", uuid::Uuid::new_v4()), + role: "user".to_string(), mode: "chat".to_string(), + text: natural_prompt.to_string(), created_at: now_epoch_seconds(), + sql_draft: None, + }); + bucket.push(AskVeloxyConversationMessage { + id: format!("msg-{}", uuid::Uuid::new_v4()), + role: "assistant".to_string(), mode: "chat".to_string(), + text: message.clone(), created_at: now_epoch_seconds(), + sql_draft: sql_draft.clone(), + }); + if bucket.len() > ASK_VELOXY_MAX_HISTORY_MESSAGES { + let remove_count = bucket.len() - ASK_VELOXY_MAX_HISTORY_MESSAGES; + bucket.drain(0..remove_count); + } + } + + Ok(AskVeloxyChatResponse { + message, suggestions, warnings, sql_draft, needs_sql_generation, needs_clarification, + }) +} + +#[tauri::command] +pub async fn load_veloxy_conversation( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result { + let (resolved_connection_id, _) = resolve_connection_engine(&app, &state, connection_id).await?; + let stored_connection = load_connection(&app, &resolved_connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let key = ask_veloxy_conversation_key(&resolved_connection_id, &stored_connection.database); + let messages = state.ask_veloxy_conversations.read().await.get(&key).cloned().unwrap_or_default(); + Ok(AskVeloxyConversationResponse { messages }) +} + +#[tauri::command] +pub async fn clear_veloxy_conversation( + app: AppHandle, + state: State<'_, AppState>, + connection_id: Option, +) -> Result<(), String> { + let (resolved_connection_id, _) = resolve_connection_engine(&app, &state, connection_id).await?; + let stored_connection = load_connection(&app, &resolved_connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let key = ask_veloxy_conversation_key(&resolved_connection_id, &stored_connection.database); + state.ask_veloxy_conversations.write().await.remove(&key); + Ok(()) +} + +#[tauri::command] +pub async fn generate_sql_from_nl( + app: AppHandle, + state: State<'_, AppState>, + input: AskVeloxyRequest, +) -> Result { + let natural_prompt = input.natural_prompt.trim(); + if natural_prompt.is_empty() { + return Err("Ask Veloxy prompt cannot be empty.".to_string()); + } + if input.provider_config.api_key.trim().is_empty() { + return Err("OpenRouter API key is required.".to_string()); + } + if input.provider_config.model.trim().is_empty() { + return Err("OpenRouter model is required.".to_string()); + } + + let (connection_id, engine) = resolve_connection_engine(&app, &state, input.connection_id.clone()).await?; + let db_context = get_or_build_ask_veloxy_db_context(&app, &state, &connection_id, engine).await?; + let schema_context = build_schema_context(&db_context, natural_prompt, input.target_table.as_ref()); + + let mut user_prompt = format!( + "Engine: {:?}\nDatabase: {}\nTask: {}\nMaxRows: {}\nSchemaContext:\n{}\n", + db_context.engine, db_context.database_name, natural_prompt, + input.max_rows.unwrap_or(MAX_QUERY_ROWS), schema_context + ); + truncate_on_char_boundary(&mut user_prompt, ASK_VELOXY_PROMPT_CHAR_BUDGET); + + let system_prompt = "You are Ask Veloxy. Return JSON only with keys: sql (string), intent (string), confidence (number 0..1), explanation (string), suggestions (array of short strings), warnings (array of strings). Generate exactly one SQL statement, keep explanation concise, and never include markdown."; + let base_url = normalize_openrouter_base(input.provider_config.base_url.as_deref()); + let endpoint = format!("{}/chat/completions", base_url); + + let client = state.openrouter_client.get_or_init(reqwest::Client::new); + let response = client.post(&endpoint) + .header("Authorization", format!("Bearer {}", input.provider_config.api_key.trim())) + .header("Content-Type", "application/json") + .json(&serde_json::json!({ + "model": input.provider_config.model.trim(), + "temperature": 0.1, + "max_tokens": 500, + "messages": [ + { "role": "system", "content": system_prompt }, + { "role": "user", "content": user_prompt } + ] + })) + .send().await.map_err(|error| format!("OpenRouter request failed: {}", error))?; + + let status = response.status(); + let payload = response.json::().await + .map_err(|error| format!("Invalid OpenRouter JSON response: {}", error))?; + if !status.is_success() { + let message = payload.get("error").and_then(|error| error.get("message")) + .and_then(Value::as_str).unwrap_or("Unknown OpenRouter error"); + return Err(format!("OpenRouter error ({}): {}", status.as_u16(), message)); + } + + let message_content = extract_openrouter_message_content(&payload)?; + let generated = parse_ask_veloxy_json(&message_content)?; + let sql = generated.get("sql").and_then(Value::as_str).unwrap_or_default().trim().to_string(); + validate_generated_sql(&sql)?; + + let mut warnings = generated.get("warnings").and_then(Value::as_array) + .map(|items| items.iter().filter_map(Value::as_str).map(str::to_string).collect::>()) + .unwrap_or_default(); + + let intent = generated.get("intent").and_then(Value::as_str).map(str::to_string) + .unwrap_or_else(|| classify_sql_intent(&sql)); + let confidence = generated.get("confidence").and_then(Value::as_f64).unwrap_or(0.6).clamp(0.0, 1.0); + let explanation = generated.get("explanation").and_then(Value::as_str).map(str::trim) + .filter(|value| !value.is_empty()).map(|value| { + let mut truncated = value.to_string(); + truncate_on_char_boundary(&mut truncated, 350); + truncated + }); + let suggestions = parse_ask_veloxy_suggestions(&generated); + + if intent != "select" { + warnings.push("Generated SQL is not read-only. Review before execution.".to_string()); + } + if confidence < 0.5 { + warnings.push("Low confidence result. Review SQL carefully.".to_string()); + } + + let token_stats = AskVeloxyTokenStats { + schema_chars: schema_context.len(), + schema_tokens_estimate: estimate_tokens(schema_context.len()), + prompt_chars: user_prompt.len() + system_prompt.len(), + prompt_tokens_estimate: estimate_tokens(user_prompt.len() + system_prompt.len()), + }; + + Ok(AskVeloxyResponse { + sql, intent, confidence, explanation, suggestions, warnings, token_stats, + }) +} + +#[cfg(test)] +mod tests { + use super::super::{ + build_schema_context, classify_sql_intent, database_name_from_mysql_value, + decode_mysql_bytes_as_string, extract_openrouter_stream_delta, mysql_decode_error, + parse_ask_veloxy_json, sqlite_decode_error, streaming_display_text, + validate_generated_sql, is_read_only_sql, + }; + use crate::models::{ + AskVeloxyDbContextCache, DatabaseEngine, QueryEditorColumn, QueryEditorMetadata, + QueryEditorTable, + }; + + #[test] + fn streaming_display_text_extracts_partial_json_message() { + let partial = r#"{ "message": "The messages table has relationships with:\n- delivery_reports"#; + let display = streaming_display_text(partial); + assert!(display.contains("messages table")); + assert!(display.contains("delivery_reports")); + } + + #[test] + fn streaming_display_text_returns_plain_text_directly() { + assert_eq!(streaming_display_text("Hello from Veloxy"), "Hello from Veloxy"); + } + + #[test] + fn extract_openrouter_stream_delta_reads_content() { + let data = r#"{"choices":[{"delta":{"content":"Hello"}}]}"#; + assert_eq!(extract_openrouter_stream_delta(data).as_deref(), Some("Hello")); + } + + #[test] + fn database_name_from_mysql_value_rejects_empty() { + assert!(database_name_from_mysql_value(None, "list_databases").is_err()); + assert!(database_name_from_mysql_value(Some(String::new()), "list_databases").is_err()); + } + + #[test] + fn database_name_from_mysql_value_accepts_non_empty() { + let name = database_name_from_mysql_value(Some("my_app".to_string()), "list_databases").expect("name"); + assert_eq!(name, "my_app"); + } + + #[test] + fn decode_mysql_bytes_as_string_uses_utf8_text() { + assert_eq!(decode_mysql_bytes_as_string(b"my_schema"), "my_schema"); + } + + #[test] + fn mysql_decode_error_is_explicit() { + let message = mysql_decode_error("get_tables", "table_schema", Some(0), "mismatched types"); + assert!(message.contains("MySQL decode error")); + assert!(message.contains("get_tables")); + assert!(message.contains("table_schema")); + } + + #[test] + fn sqlite_decode_error_is_explicit() { + let message = sqlite_decode_error("get_schema", "name", Some(0), "unsupported value type"); + assert!(message.contains("SQLite decode error")); + assert!(message.contains("get_schema")); + assert!(message.contains("name")); + } + + #[test] + fn schema_context_is_bounded() { + let columns = (0..40).map(|idx| QueryEditorColumn { + name: format!("column_{}", idx), data_type: "text".to_string(), + }).collect::>(); + let tables = (0..20).map(|idx| QueryEditorTable { + schema: "public".to_string(), name: format!("events_{}", idx), columns: columns.clone(), + }).collect::>(); + let metadata = QueryEditorMetadata { + tables, functions: Vec::new(), + truncated_tables: false, truncated_columns: false, truncated_functions: false, + }; + let db_context = AskVeloxyDbContextCache { + database_name: "test".to_string(), engine: DatabaseEngine::Postgres, + metadata, foreign_keys: Vec::new(), + }; + let context = build_schema_context(&db_context, "show events", None); + assert!(!context.is_empty()); + assert!(context.len() <= super::super::ASK_VELOXY_SCHEMA_CHAR_BUDGET); + } + + #[test] + fn ask_veloxy_json_parser_handles_embedded_block() { + let content = "Here is the output {\"sql\":\"select 1\",\"intent\":\"select\",\"confidence\":0.9,\"warnings\":[]}"; + let parsed = parse_ask_veloxy_json(content).expect("json should parse"); + assert_eq!(parsed.get("sql").and_then(|v| v.as_str()), Some("select 1")); + } + + #[test] + fn sql_validation_rejects_multi_statement() { + assert!(validate_generated_sql("select 1; select 2;").is_err()); + } + + #[test] + fn sql_intent_classifier_recognizes_update() { + assert_eq!(classify_sql_intent("UPDATE foo SET bar = 1"), "update"); + } + + #[test] + fn read_only_check_allows_selects_and_explain() { + assert!(is_read_only_sql("SELECT 1")); + assert!(is_read_only_sql("EXPLAIN ANALYZE SELECT * FROM t")); + assert!(is_read_only_sql("WITH x AS (SELECT 1) SELECT * FROM x")); + assert!(is_read_only_sql("BEGIN; SELECT 1; COMMIT;")); + } + + #[test] + fn read_only_check_blocks_writes() { + assert!(!is_read_only_sql("DELETE FROM t")); + assert!(!is_read_only_sql("DROP TABLE t")); + assert!(!is_read_only_sql("BEGIN; UPDATE t SET a = 1; COMMIT;")); + assert!(!is_read_only_sql("SELECT 1; DELETE FROM t")); + assert!(!is_read_only_sql("")); + } + + #[test] + fn mysql_timestamp_formats_as_datetime_string() { + let dt = chrono::DateTime::parse_from_rfc3339("2024-03-15T10:30:45Z").unwrap() + .with_timezone(&chrono::Utc); + assert_eq!(dt.format("%Y-%m-%d %H:%M:%S").to_string(), "2024-03-15 10:30:45"); + } + + #[test] + fn mysql_datetime_formats_as_naive_datetime_string() { + let dt = chrono::NaiveDateTime::parse_from_str("2024-03-15 10:30:45", "%Y-%m-%d %H:%M:%S").unwrap(); + assert_eq!(dt.format("%Y-%m-%d %H:%M:%S").to_string(), "2024-03-15 10:30:45"); + } + + #[test] + fn mysql_date_formats_as_iso_date() { + let d = chrono::NaiveDate::from_ymd_opt(2024, 3, 15).unwrap(); + assert_eq!(d.to_string(), "2024-03-15"); + } + + #[test] + fn mysql_time_formats_as_iso_time() { + let t = chrono::NaiveTime::from_hms_opt(10, 30, 45).unwrap(); + assert_eq!(t.to_string(), "10:30:45"); + } +} diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 1b618ab..b700321 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -1,790 +1,7 @@ -use std::collections::HashMap; -use std::sync::atomic::AtomicBool; -use std::sync::{Arc, OnceLock}; -use std::time::Duration; +mod connection_pool; +mod identifier_guard; +mod retry; -use deadpool_postgres::{ - Client, Config as PostgresConfig, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime, - SslMode as DeadpoolSslMode, Timeouts, -}; -use sqlx::{MySqlPool, SqlitePool}; -use tauri::AppHandle; -use tauri_plugin_store::StoreExt; -use tokio::sync::RwLock; -use tokio_postgres::NoTls; -use rustls::pki_types::{CertificateDer, PrivateKeyDer}; - -use crate::models::{ - AskVeloxyConversationMessage, AskVeloxyDbContextCache, ConnectionInput, ConnectionSslMode, - ConnectionSummary, DatabaseEngine, StoredConnection, -}; -use crate::ssh_tunnel::SshTunnel; -use crate::credentials; - -pub const CONNECTION_STORE_PATH: &str = "connections.json"; -pub const MAX_QUERY_ROWS: usize = 1000; - -const APP_NAME: &str = "VeloxDB"; -const CONNECT_TIMEOUT_SECS: u64 = 12; -const KEEPALIVES_IDLE_SECS: u64 = 60; -const POOL_MAX_SIZE: usize = 6; -const POOL_WAIT_SECS: u64 = 30; -const POOL_CREATE_SECS: u64 = 15; -const POOL_RECYCLE_SECS: u64 = 15; -pub const DEFAULT_MYSQL_PORT: u16 = 3306; - -fn deadpool_ssl_mode(mode: ConnectionSslMode) -> DeadpoolSslMode { - match mode { - ConnectionSslMode::Disable => DeadpoolSslMode::Disable, - ConnectionSslMode::Prefer => DeadpoolSslMode::Prefer, - ConnectionSslMode::Require => DeadpoolSslMode::Require, - } -} - -#[derive(Default)] -pub struct AppState { - pub pools: RwLock>, - pub mysql_pools: RwLock>, - pub sqlite_pools: RwLock>, - pub active_connection_id: RwLock>, - pub ssh_tunnels: RwLock>, - pub ask_veloxy_db_context_cache: RwLock>, - pub ask_veloxy_conversations: RwLock>>, - pub veloxy_cancel: RwLock>>, - pub openrouter_client: OnceLock, -} - -fn load_pem_certs(path: &str) -> Result>, String> { - let data = std::fs::read(path).map_err(|e| format!("Failed to read cert file {}: {}", path, e))?; - let mut certs = Vec::new(); - for result in rustls_pemfile::certs(&mut &data[..]) { - let cert = result.map_err(|e| format!("Failed to parse cert from {}: {}", path, e))?; - certs.push(cert); - } - if certs.is_empty() { - return Err(format!("No certificates found in {}", path)); - } - Ok(certs) -} - -fn load_pem_key(path: &str) -> Result, String> { - let data = std::fs::read(path).map_err(|e| format!("Failed to read key file {}: {}", path, e))?; - for result in rustls_pemfile::read_all(&mut &data[..]) { - match result { - Ok(rustls_pemfile::Item::Pkcs1Key(key)) => return Ok(key.into()), - Ok(rustls_pemfile::Item::Pkcs8Key(key)) => return Ok(key.into()), - Ok(rustls_pemfile::Item::Sec1Key(key)) => return Ok(key.into()), - Err(e) => return Err(format!("Failed to parse key from {}: {}", path, e)), - _ => continue, - } - } - Err(format!("No private key found in {}", path)) -} - -fn tls_connector_with_params(extra_params: &HashMap) -> Result { - let mut root_store = rustls::RootCertStore::empty(); - - let has_custom_root = extra_params.contains_key("sslrootcert"); - if has_custom_root { - let path = extra_params.get("sslrootcert").unwrap(); - let certs = load_pem_certs(path)?; - for cert in certs { - root_store.add(cert).map_err(|e| format!("Failed to add root cert: {}", e))?; - } - } else { - root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); - } - - let config_builder = rustls::ClientConfig::builder() - .with_root_certificates(root_store); - - let has_client_cert = extra_params.contains_key("sslcert") && extra_params.contains_key("sslkey"); - let config = if has_client_cert { - let cert_path = extra_params.get("sslcert").unwrap(); - let key_path = extra_params.get("sslkey").unwrap(); - let certs = load_pem_certs(cert_path)?; - let key = load_pem_key(key_path)?; - config_builder - .with_client_auth_cert(certs, key) - .map_err(|e| format!("Failed to configure client auth: {}", e))? - } else { - config_builder - .with_no_client_auth() - }; - - Ok(tokio_postgres_rustls::MakeRustlsConnect::new(config)) -} - -fn tls_connector() -> Result { - static CACHE: OnceLock> = OnceLock::new(); - CACHE - .get_or_init(|| { - let mut root_store = rustls::RootCertStore::empty(); - root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); - let config = rustls::ClientConfig::builder() - .with_root_certificates(root_store) - .with_no_client_auth(); - Ok(tokio_postgres_rustls::MakeRustlsConnect::new(config)) - }) - .clone() -} - -fn apply_extra_params( - config: &mut PostgresConfig, - extra_params: &HashMap, -) { - let mut remaining_opts: Vec = Vec::new(); - - for (key, value) in extra_params { - match key.as_str() { - "connect_timeout" => { - if let Ok(secs) = value.parse::() { - config.connect_timeout = Some(Duration::from_secs(secs)); - } - } - "application_name" => { - config.application_name = Some(value.clone()); - } - "keepalives_idle" => { - if let Ok(secs) = value.parse::() { - config.keepalives_idle = Some(Duration::from_secs(secs)); - } - } - "options" => { - if !value.is_empty() { - remaining_opts.push(value.clone()); - } - } - "keepalives" => { - if let Ok(v) = value.parse::() { - config.keepalives = Some(v != 0); - } - } - // TLS params handled by tls_connector_with_params - "sslrootcert" | "sslcert" | "sslkey" => {} - _ => { - let sanitized_value = value.replace('\\', "\\\\").replace('\'', "\\'"); - remaining_opts.push(format!("-c {}={}", key, sanitized_value)); - } - } - } - - if !remaining_opts.is_empty() { - let existing = config.options.clone().unwrap_or_default(); - let merged = if existing.is_empty() { - remaining_opts.join(" ") - } else { - format!("{} {}", existing, remaining_opts.join(" ")) - }; - config.options = Some(merged); - } -} - -pub fn build_pool_custom(host: &str, port: u16, input: &ConnectionInput) -> Result { - let mut config = PostgresConfig::new(); - config.host = Some(host.to_string()); - config.port = Some(port); - config.dbname = Some(input.database.clone()); - config.user = Some(input.user.clone()); - config.password = Some(input.password.clone()); - config.application_name = Some(APP_NAME.to_string()); - config.connect_timeout = Some(Duration::from_secs(CONNECT_TIMEOUT_SECS)); - config.keepalives = Some(true); - config.keepalives_idle = Some(Duration::from_secs(KEEPALIVES_IDLE_SECS)); - config.ssl_mode = Some(deadpool_ssl_mode(input.ssl_mode)); - config.manager = Some(ManagerConfig { - recycling_method: RecyclingMethod::Verified, - }); - - let mut pool_config = PoolConfig::new(POOL_MAX_SIZE); - pool_config.timeouts = Timeouts { - wait: Some(Duration::from_secs(POOL_WAIT_SECS)), - create: Some(Duration::from_secs(POOL_CREATE_SECS)), - recycle: Some(Duration::from_secs(POOL_RECYCLE_SECS)), - }; - config.pool = Some(pool_config); - - if let Some(ref extra) = input.extra_params { - apply_extra_params(&mut config, extra); - } - - let has_tls_params = input.extra_params.as_ref().map_or(false, |e| { - e.contains_key("sslrootcert") || e.contains_key("sslcert") || e.contains_key("sslkey") - }); - - match input.ssl_mode { - ConnectionSslMode::Disable => config - .create_pool(Some(Runtime::Tokio1), NoTls) - .map_err(|error| error.to_string()), - ConnectionSslMode::Prefer | ConnectionSslMode::Require => { - let tls = if has_tls_params { - tls_connector_with_params(input.extra_params.as_ref().unwrap())? - } else { - tls_connector()? - }; - config - .create_pool(Some(Runtime::Tokio1), tls) - .map_err(|error| error.to_string()) - } - } -} - -pub fn build_pool(input: &ConnectionInput) -> Result { - build_pool_custom(&input.host, input.port, input) -} - -fn mysql_url(host: &str, port: u16, input: &ConnectionInput) -> String { - let password = urlencoding::encode(&input.password); - let user = urlencoding::encode(&input.user); - let database = urlencoding::encode(&input.database); - let ssl_param = match input.ssl_mode { - ConnectionSslMode::Disable => "ssl-mode=DISABLED", - ConnectionSslMode::Prefer => "ssl-mode=PREFERRED", - ConnectionSslMode::Require => "ssl-mode=REQUIRED", - }; - format!( - "mysql://{}:{}@{}:{}/{}?{}", - user, password, host, port, database, ssl_param - ) -} - -fn sqlite_url(input: &ConnectionInput) -> Result { - let path = input - .file_path - .clone() - .unwrap_or_else(|| input.database.clone()); - if path.trim().is_empty() { - return Err("SQLite file path is required.".to_string()); - } - if path == ":memory:" { - return Ok("sqlite::memory:".to_string()); - } - if path.starts_with('/') { - Ok(format!("sqlite://{}", path)) - } else { - Ok(format!("sqlite://{}", path)) - } -} - -pub async fn build_mysql_pool_custom(host: &str, port: u16, input: &ConnectionInput) -> Result { - let url = mysql_url(host, port, input); - sqlx::mysql::MySqlPoolOptions::new() - .max_connections(POOL_MAX_SIZE as u32) - .acquire_timeout(Duration::from_secs(POOL_WAIT_SECS)) - .connect(&url) - .await - .map_err(|error| error.to_string()) -} - -pub async fn build_mysql_pool(input: &ConnectionInput) -> Result { - let port = if input.port == 0 { DEFAULT_MYSQL_PORT } else { input.port }; - build_mysql_pool_custom(&input.host, port, input).await -} - -pub async fn build_sqlite_pool(input: &ConnectionInput) -> Result { - use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqliteSynchronous}; - - let url = sqlite_url(input)?; - let options: SqliteConnectOptions = url - .parse() - .map_err(|error: sqlx::Error| error.to_string())?; - let options = options - .journal_mode(SqliteJournalMode::Wal) - .synchronous(SqliteSynchronous::Normal); - - sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(POOL_MAX_SIZE as u32) - .acquire_timeout(Duration::from_secs(POOL_WAIT_SECS)) - .connect_with(options) - .await - .map_err(|error| error.to_string()) -} - -/// Heuristic for transport-level failures where discarding the pool and opening -/// a new TCP session may succeed (sleep/VPN blips, idle disconnects). -fn is_retryable_connection_error(message: &str) -> bool { - let m = message.to_lowercase(); - m.contains("broken pipe") - || m.contains("connection reset") - || m.contains("connection refused") - || m.contains("unexpected eof") - || m.contains("unexpected end of file") - || m.contains("error communicating with the server") - || m.contains("connection closed") - || m.contains("closed the connection") - || m.contains("server closed the connection") - || m.contains("eof has been reached") - || m.contains("could not receive data from server") - || m.contains("could not send data to server") - || m.contains("timeout occurred while waiting") - || m.contains("timeout occurred while creating") - || m.contains("timeout occurred while recycling") -} - -pub async fn drop_pool(state: &AppState, connection_id: &str) { - state.pools.write().await.remove(connection_id); - state.mysql_pools.write().await.remove(connection_id); - state.sqlite_pools.write().await.remove(connection_id); - state - .ask_veloxy_db_context_cache - .write() - .await - .retain(|cache_key, _| !cache_key.starts_with(&format!("{}::", connection_id))); - state - .ask_veloxy_conversations - .write() - .await - .retain(|conversation_key, _| !conversation_key.starts_with(&format!("{}::", connection_id))); - if let Some(mut tunnel) = state.ssh_tunnels.write().await.remove(connection_id) { - tunnel.close().await; - } -} - -pub async fn disconnect_connection(state: &AppState, connection_id: &str) { - drop_pool(state, connection_id).await; - let mut active = state.active_connection_id.write().await; - if active.as_deref() == Some(connection_id) { - *active = None; - } -} - -/// Runs `operation` with a pooled client. On a retryable connection error, drops -/// the cached pool once and retries the whole operation (at most one extra attempt). -pub async fn with_pool_client_retry( - app: &AppHandle, - state: &AppState, - connection_id: &str, - ctx: C, - mut operation: F, -) -> Result -where - C: Clone + Send, - F: FnMut(Client, C) -> Fut, - Fut: std::future::Future> + Send, - T: Send, -{ - let mut dropped_pool = false; - - loop { - let pool = get_or_create_pool(app, state, connection_id).await?; - let client = match pool.get().await { - Ok(client) => client, - Err(error) => { - let message = error.to_string(); - if !dropped_pool && is_retryable_connection_error(&message) { - drop_pool(state, connection_id).await; - dropped_pool = true; - continue; - } - return Err(message); - } - }; - - match operation(client, ctx.clone()).await { - Ok(value) => return Ok(value), - Err(message) => { - if !dropped_pool && is_retryable_connection_error(&message) { - drop_pool(state, connection_id).await; - dropped_pool = true; - continue; - } - return Err(message); - } - } - } -} - -pub async fn resolve_connection_id( - state: &AppState, - requested_connection_id: Option, -) -> Result { - if let Some(connection_id) = requested_connection_id { - return Ok(connection_id); - } - - state - .active_connection_id - .read() - .await - .clone() - .ok_or_else(|| "Connect to a database before running this action.".to_string()) -} - -pub async fn get_or_create_pool( - app: &AppHandle, - state: &AppState, - connection_id: &str, -) -> Result { - if let Some(pool) = state.pools.read().await.get(connection_id).cloned() { - return Ok(pool); - } - - let stored_connection = load_connection(app, connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - - let input = stored_connection.to_input(); - - let resolved_port = if input.port == 0 { DEFAULT_MYSQL_PORT } else { input.port }; - let (host, port) = if let Some(ref ssh_config) = input.ssh_config { - if ssh_config.is_active() { - let tunnel = SshTunnel::connect(ssh_config, &input.host, resolved_port).await?; - let local_port = tunnel.local_port; - state - .ssh_tunnels - .write() - .await - .insert(connection_id.to_string(), tunnel); - ("127.0.0.1".to_string(), local_port) - } else { - (input.host.clone(), resolved_port) - } - } else { - (input.host.clone(), resolved_port) - }; - - let pool = build_pool_custom(&host, port, &input)?; - - state - .pools - .write() - .await - .insert(connection_id.to_string(), pool.clone()); - - Ok(pool) -} - -pub async fn get_or_create_mysql_pool( - app: &AppHandle, - state: &AppState, - connection_id: &str, -) -> Result { - if let Some(pool) = state.mysql_pools.read().await.get(connection_id).cloned() { - return Ok(pool); - } - - let stored_connection = load_connection(app, connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - - let input = stored_connection.to_input(); - let resolved_port = if input.port == 0 { - DEFAULT_MYSQL_PORT - } else { - input.port - }; - - let (host, port) = if let Some(ref ssh_config) = input.ssh_config { - if ssh_config.is_active() { - let tunnel = SshTunnel::connect(ssh_config, &input.host, resolved_port).await?; - let local_port = tunnel.local_port; - state - .ssh_tunnels - .write() - .await - .insert(connection_id.to_string(), tunnel); - ("127.0.0.1".to_string(), local_port) - } else { - (input.host.clone(), resolved_port) - } - } else { - (input.host.clone(), resolved_port) - }; - - let pool = build_mysql_pool_custom(&host, port, &input).await?; - state - .mysql_pools - .write() - .await - .insert(connection_id.to_string(), pool.clone()); - Ok(pool) -} - -pub async fn get_or_create_sqlite_pool( - app: &AppHandle, - state: &AppState, - connection_id: &str, -) -> Result { - if let Some(pool) = state.sqlite_pools.read().await.get(connection_id).cloned() { - return Ok(pool); - } - - let stored_connection = load_connection(app, connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - - let input = stored_connection.to_input(); - let pool = build_sqlite_pool(&input).await?; - - state - .sqlite_pools - .write() - .await - .insert(connection_id.to_string(), pool.clone()); - - Ok(pool) -} - -/// Drops cached pools and re-validates the connection (used by sidebar refresh). -pub async fn refresh_connection_pools( - app: &AppHandle, - state: &AppState, - connection_id: &str, -) -> Result<(), String> { - let stored = load_connection(app, connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - let engine = stored.engine; - - drop_pool(state, connection_id).await; - - match engine { - DatabaseEngine::Postgres => { - with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { - client - .simple_query("select 1") - .await - .map_err(|error| error.to_string())?; - Ok(()) - }) - .await?; - } - DatabaseEngine::Mysql => { - let pool = get_or_create_mysql_pool(app, state, connection_id).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - } - DatabaseEngine::Sqlite => { - let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; - sqlx::query("select 1") - .execute(&pool) - .await - .map_err(|error| error.to_string())?; - } - } - - Ok(()) -} - -pub async fn resolve_connection_engine( - app: &AppHandle, - state: &AppState, - requested_connection_id: Option, -) -> Result<(String, DatabaseEngine), String> { - let connection_id = resolve_connection_id(state, requested_connection_id).await?; - let stored = load_connection(app, &connection_id)? - .ok_or_else(|| "Stored connection details were not found.".to_string())?; - Ok((connection_id, stored.engine)) -} - -pub fn list_connections(app: &AppHandle) -> Result, String> { - let store = app - .store(CONNECTION_STORE_PATH) - .map_err(|error| error.to_string())?; - - let mut connections = store - .entries() - .into_iter() - .map(|(_, value)| { - serde_json::from_value::(value).map_err(|error| error.to_string()) - }) - .collect::, _>>()?; - - connections.sort_by(|left, right| { - right - .connected_at - .cmp(&left.connected_at) - .then_with(|| left.name.cmp(&right.name)) - }); - - Ok(connections - .into_iter() - .map(|connection| connection.summary()) - .collect()) -} - -pub fn persist_connection(app: &AppHandle, connection: &StoredConnection) -> Result<(), String> { - let store = app - .store(CONNECTION_STORE_PATH) - .map_err(|error| error.to_string())?; - - store.set( - connection.id.clone(), - serde_json::to_value(connection).map_err(|error| error.to_string())?, - ); - - store.save().map_err(|error| error.to_string())?; - - Ok(()) -} - -pub fn persist_connection_with_password( - app: &AppHandle, - connection: &StoredConnection, - password: &str, -) -> Result<(), String> { - let mut stored = connection.clone(); - stored.password = Some(password.to_string()); - persist_connection(app, &stored)?; - if connection.engine != DatabaseEngine::Sqlite { - credentials::store_password(&connection.id, password)?; - } - Ok(()) -} - -/// Renames a persisted connection and returns the updated summary. -pub fn rename_connection_in_store( - app: &AppHandle, - connection_id: &str, - new_name: &str, -) -> Result { - let mut stored = load_connection(app, connection_id)? - .ok_or_else(|| format!("Connection {} not found.", connection_id))?; - stored.name = new_name.to_string(); - persist_connection(app, &stored)?; - Ok(stored.summary()) -} - -pub fn delete_connection_from_store(app: &AppHandle, connection_id: &str) -> Result<(), String> { - let store = app - .store(CONNECTION_STORE_PATH) - .map_err(|error| error.to_string())?; - store.delete(connection_id); - store.save().map_err(|error| error.to_string())?; - Ok(()) -} - -pub fn load_connection( - app: &AppHandle, - connection_id: &str, -) -> Result, String> { - let store = app - .store(CONNECTION_STORE_PATH) - .map_err(|error| error.to_string())?; - - let mut connection: StoredConnection = match store - .get(connection_id) - .map(|value| serde_json::from_value::(value).map_err(|error| error.to_string())) - .transpose()? - { - Some(conn) => conn, - None => return Ok(None), - }; - - let json_password = connection.password.clone(); - - if connection.engine == DatabaseEngine::Sqlite { - connection.password = Some(String::new()); - return Ok(Some(connection)); - } - - match credentials::get_password(connection_id) { - Ok(Some(password)) => { - if json_password.is_some() { - let clean = StoredConnection { - password: None, - ..connection.clone() - }; - let _ = persist_connection(app, &clean); - } - connection.password = Some(password); - } - Ok(None) => { - if let Some(ref pwd) = json_password { - if let Err(e) = credentials::store_password(connection_id, pwd) { - log::warn!("Failed to migrate password to keychain for {}: {}", connection_id, e); - } - connection.password = Some(pwd.clone()); - } else { - return Err(format!( - "No saved password found for connection {}. Re-enter your password to reconnect.", - connection.name - )); - } - } - Err(e) => { - log::warn!("Failed to read password from keychain for {}: {}", connection_id, e); - if let Some(ref pwd) = json_password { - connection.password = Some(pwd.clone()); - } else { - return Err(format!( - "Failed to read saved password for connection {} (keychain error). Re-enter your password to reconnect.", - connection.name - )); - } - } - } - - Ok(Some(connection)) -} - -pub fn quote_identifier(value: &str) -> String { - value.replace('"', "\"\"") -} - -/// Whether `name` is safe to interpolate into dynamic SQL (e.g. SQLite `PRAGMA` -/// statements, where bind parameters are not allowed). Restricted to ASCII -/// alphanumerics and underscores so it cannot terminate or escape a statement. -pub fn is_safe_identifier(name: &str) -> bool { - !name.is_empty() - && name.len() <= 128 - && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') -} - -/// Validates an identifier before it is interpolated into dynamic SQL. Returns -/// the identifier unchanged when safe, or a descriptive error otherwise. -pub fn require_safe_identifier<'a>(name: &'a str, context: &str) -> Result<&'a str, String> { - if is_safe_identifier(name) { - Ok(name) - } else { - Err(format!("Invalid identifier for {}: {:?}", context, name)) - } -} - -#[cfg(test)] -mod tests { - use super::{is_safe_identifier, mysql_url, require_safe_identifier}; - use crate::models::{ConnectionInput, ConnectionSslMode, DatabaseEngine}; - - fn mysql_input(ssl_mode: ConnectionSslMode) -> ConnectionInput { - ConnectionInput { - id: None, - name: "test".to_string(), - engine: DatabaseEngine::Mysql, - host: "localhost".to_string(), - port: 3306, - database: "app".to_string(), - file_path: None, - user: "root".to_string(), - password: "pw".to_string(), - ssl_mode, - ssh_config: None, - extra_params: None, - } - } - - #[test] - fn mysql_url_maps_ssl_mode() { - assert!(mysql_url("localhost", 3306, &mysql_input(ConnectionSslMode::Disable)) - .ends_with("?ssl-mode=DISABLED")); - assert!(mysql_url("localhost", 3306, &mysql_input(ConnectionSslMode::Prefer)) - .ends_with("?ssl-mode=PREFERRED")); - assert!(mysql_url("localhost", 3306, &mysql_input(ConnectionSslMode::Require)) - .ends_with("?ssl-mode=REQUIRED")); - } - - #[test] - fn rejects_sql_injection_in_identifiers() { - assert!(!is_safe_identifier("\"; DROP TABLE users; --")); - assert!(!is_safe_identifier("foo; DELETE FROM bar")); - assert!(!is_safe_identifier("a\0b")); - assert!(!is_safe_identifier("")); - assert!(require_safe_identifier("foo);", "table name").is_err()); - } - - #[test] - fn accepts_plain_identifiers() { - assert!(is_safe_identifier("users")); - assert!(is_safe_identifier("_internal")); - assert!(is_safe_identifier("Table123")); - assert_eq!(require_safe_identifier("users", "table name"), Ok("users")); - } -} +pub use connection_pool::*; +pub use identifier_guard::*; +pub use retry::*; diff --git a/src-tauri/src/db/connection_pool.rs b/src-tauri/src/db/connection_pool.rs new file mode 100644 index 0000000..8e3684a --- /dev/null +++ b/src-tauri/src/db/connection_pool.rs @@ -0,0 +1,1008 @@ +use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use deadpool_postgres::{ + Config as PostgresConfig, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime, + SslMode as DeadpoolSslMode, Timeouts, +}; +use sqlx::{MySqlPool, SqlitePool}; +use tauri::AppHandle; +use tauri_plugin_store::StoreExt; +use tokio::sync::RwLock; +use tokio_postgres::NoTls; +use rustls::pki_types::{CertificateDer, PrivateKeyDer}; + +use mongodb::{Client as MongoClient, bson::doc}; +use duckdb::Connection as DuckConnection; +use crate::models::{ + AskVeloxyConversationMessage, AskVeloxyDbContextCache, ConnectionInput, ConnectionSslMode, + ConnectionSummary, DatabaseEngine, StoredConnection, +}; +use crate::ssh_tunnel::SshTunnel; +use crate::credentials; + +use super::retry; + +use crate::error::VeloxError; + +pub const CONNECTION_STORE_PATH: &str = "connections.json"; +pub const MAX_QUERY_ROWS: usize = 1000; + +const APP_NAME: &str = "VeloxDB"; +const CONNECT_TIMEOUT_SECS: u64 = 12; +const KEEPALIVES_IDLE_SECS: u64 = 60; +const POOL_MAX_SIZE: usize = 6; +const POOL_WAIT_SECS: u64 = 30; +const POOL_CREATE_SECS: u64 = 15; +const POOL_RECYCLE_SECS: u64 = 15; +pub const DEFAULT_MYSQL_PORT: u16 = 3306; +pub const DEFAULT_MONGO_PORT: u16 = 27017; +pub const DEFAULT_REDIS_PORT: u16 = 6379; + +fn deadpool_ssl_mode(mode: ConnectionSslMode) -> DeadpoolSslMode { + match mode { + ConnectionSslMode::Disable => DeadpoolSslMode::Disable, + ConnectionSslMode::Prefer => DeadpoolSslMode::Prefer, + ConnectionSslMode::Require => DeadpoolSslMode::Require, + } +} + +#[derive(Default)] +pub struct AppState { + pub pools: RwLock>, + pub mysql_pools: RwLock>, + pub sqlite_pools: RwLock>, + pub mongo_clients: RwLock>, + pub duckdb_connections: RwLock>>, + pub redis_clients: RwLock>, + pub active_connection_id: RwLock>, + pub ssh_tunnels: RwLock>, + pub ask_veloxy_db_context_cache: RwLock>, + pub ask_veloxy_conversations: RwLock>>, + pub veloxy_cancel: RwLock>>, + pub openrouter_client: OnceLock, +} + +fn load_pem_certs(path: &str) -> Result>, VeloxError> { + let data = std::fs::read(path).map_err(|e| format!("Failed to read cert file {}: {}", path, e))?; + let mut certs = Vec::new(); + for result in rustls_pemfile::certs(&mut &data[..]) { + let cert = result.map_err(|e| format!("Failed to parse cert from {}: {}", path, e))?; + certs.push(cert); + } + if certs.is_empty() { + return Err(VeloxError::Configuration(format!("No certificates found in {}", path))); + } + Ok(certs) +} + +fn load_pem_key(path: &str) -> Result, VeloxError> { + let data = std::fs::read(path).map_err(|e| format!("Failed to read key file {}: {}", path, e))?; + for result in rustls_pemfile::read_all(&mut &data[..]) { + match result { + Ok(rustls_pemfile::Item::Pkcs1Key(key)) => return Ok(key.into()), + Ok(rustls_pemfile::Item::Pkcs8Key(key)) => return Ok(key.into()), + Ok(rustls_pemfile::Item::Sec1Key(key)) => return Ok(key.into()), + Err(e) => return Err(VeloxError::Configuration(format!("Failed to parse key from {}: {}", path, e))), + _ => continue, + } + } + Err(VeloxError::Configuration(format!("No private key found in {}", path))) +} + +fn tls_connector_with_params(extra_params: &HashMap) -> Result { + let mut root_store = rustls::RootCertStore::empty(); + + let has_custom_root = extra_params.contains_key("sslrootcert"); + if has_custom_root { + let path = extra_params.get("sslrootcert").unwrap(); + let certs = load_pem_certs(path)?; + for cert in certs { + root_store.add(cert).map_err(|e| format!("Failed to add root cert: {}", e))?; + } + } else { + root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + } + + let config_builder = rustls::ClientConfig::builder() + .with_root_certificates(root_store); + + let has_client_cert = extra_params.contains_key("sslcert") && extra_params.contains_key("sslkey"); + let config = if has_client_cert { + let cert_path = extra_params.get("sslcert").unwrap(); + let key_path = extra_params.get("sslkey").unwrap(); + let certs = load_pem_certs(cert_path)?; + let key = load_pem_key(key_path)?; + config_builder + .with_client_auth_cert(certs, key) + .map_err(|e| format!("Failed to configure client auth: {}", e))? + } else { + config_builder + .with_no_client_auth() + }; + + Ok(tokio_postgres_rustls::MakeRustlsConnect::new(config)) +} + +fn tls_connector() -> Result { + static CACHE: OnceLock> = OnceLock::new(); + CACHE + .get_or_init(|| { + let mut root_store = rustls::RootCertStore::empty(); + root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + let config = rustls::ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + Ok(tokio_postgres_rustls::MakeRustlsConnect::new(config)) + }) + .clone() +} + +fn apply_extra_params( + config: &mut PostgresConfig, + extra_params: &HashMap, +) { + let mut remaining_opts: Vec = Vec::new(); + + for (key, value) in extra_params { + match key.as_str() { + "connect_timeout" => { + if let Ok(secs) = value.parse::() { + config.connect_timeout = Some(Duration::from_secs(secs)); + } + } + "application_name" => { + config.application_name = Some(value.clone()); + } + "keepalives_idle" => { + if let Ok(secs) = value.parse::() { + config.keepalives_idle = Some(Duration::from_secs(secs)); + } + } + "options" => { + if !value.is_empty() { + remaining_opts.push(value.clone()); + } + } + "keepalives" => { + if let Ok(v) = value.parse::() { + config.keepalives = Some(v != 0); + } + } + // TLS params handled by tls_connector_with_params + "sslrootcert" | "sslcert" | "sslkey" => {} + _ => { + let sanitized_value = value.replace('\\', "\\\\").replace('\'', "\\'"); + remaining_opts.push(format!("-c {}={}", key, sanitized_value)); + } + } + } + + if !remaining_opts.is_empty() { + let existing = config.options.clone().unwrap_or_default(); + let merged = if existing.is_empty() { + remaining_opts.join(" ") + } else { + format!("{} {}", existing, remaining_opts.join(" ")) + }; + config.options = Some(merged); + } +} + +pub fn build_pool_custom(host: &str, port: u16, input: &ConnectionInput) -> Result { + let mut config = PostgresConfig::new(); + config.host = Some(host.to_string()); + config.port = Some(port); + config.dbname = Some(input.database.clone()); + config.user = Some(input.user.clone()); + config.password = Some(input.password.clone()); + config.application_name = Some(APP_NAME.to_string()); + config.connect_timeout = Some(Duration::from_secs(CONNECT_TIMEOUT_SECS)); + config.keepalives = Some(true); + config.keepalives_idle = Some(Duration::from_secs(KEEPALIVES_IDLE_SECS)); + config.ssl_mode = Some(deadpool_ssl_mode(input.ssl_mode)); + config.manager = Some(ManagerConfig { + recycling_method: RecyclingMethod::Verified, + }); + + let mut pool_config = PoolConfig::new(POOL_MAX_SIZE); + pool_config.timeouts = Timeouts { + wait: Some(Duration::from_secs(POOL_WAIT_SECS)), + create: Some(Duration::from_secs(POOL_CREATE_SECS)), + recycle: Some(Duration::from_secs(POOL_RECYCLE_SECS)), + }; + config.pool = Some(pool_config); + + if let Some(ref extra) = input.extra_params { + apply_extra_params(&mut config, extra); + } + + let has_tls_params = input.extra_params.as_ref().is_some_and(|e| { + e.contains_key("sslrootcert") || e.contains_key("sslcert") || e.contains_key("sslkey") + }); + + match input.ssl_mode { + ConnectionSslMode::Disable => config + .create_pool(Some(Runtime::Tokio1), NoTls) + .map_err(|error| VeloxError::Connection(error.to_string())), + ConnectionSslMode::Prefer | ConnectionSslMode::Require => { + let tls = if has_tls_params { + tls_connector_with_params(input.extra_params.as_ref().unwrap())? + } else { + tls_connector()? + }; + config + .create_pool(Some(Runtime::Tokio1), tls) + .map_err(|error| VeloxError::Connection(error.to_string())) + } + } +} + +pub fn build_pool(input: &ConnectionInput) -> Result { + build_pool_custom(&input.host, input.port, input) +} + +pub fn mysql_url(host: &str, port: u16, input: &ConnectionInput) -> String { + let password = urlencoding::encode(&input.password); + let user = urlencoding::encode(&input.user); + let database = urlencoding::encode(&input.database); + let ssl_param = match input.ssl_mode { + ConnectionSslMode::Disable => "ssl-mode=DISABLED", + ConnectionSslMode::Prefer => "ssl-mode=PREFERRED", + ConnectionSslMode::Require => "ssl-mode=REQUIRED", + }; + format!( + "mysql://{}:{}@{}:{}/{}?{}", + user, password, host, port, database, ssl_param + ) +} + +fn sqlite_url(input: &ConnectionInput) -> Result { + let path = input + .file_path + .clone() + .unwrap_or_else(|| input.database.clone()); + if path.trim().is_empty() { + return Err(VeloxError::Validation("SQLite file path is required.".to_string())); + } + if path == ":memory:" { + return Ok("sqlite::memory:".to_string()); + } + Ok(format!("sqlite://{}", path)) +} + +pub async fn build_mysql_pool_custom(host: &str, port: u16, input: &ConnectionInput) -> Result { + let url = mysql_url(host, port, input); + sqlx::mysql::MySqlPoolOptions::new() + .max_connections(POOL_MAX_SIZE as u32) + .acquire_timeout(Duration::from_secs(POOL_WAIT_SECS)) + .connect(&url) + .await + .map_err(|error| VeloxError::Connection(error.to_string())) +} + +pub async fn build_mysql_pool(input: &ConnectionInput) -> Result { + let port = if input.port == 0 { DEFAULT_MYSQL_PORT } else { input.port }; + build_mysql_pool_custom(&input.host, port, input).await +} + +pub async fn build_sqlite_pool(input: &ConnectionInput) -> Result { + use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqliteSynchronous}; + + let url = sqlite_url(input)?; + let options: SqliteConnectOptions = url + .parse() + .map_err(|error: sqlx::Error| VeloxError::Connection(error.to_string()))?; + let options = options + .journal_mode(SqliteJournalMode::Wal) + .synchronous(SqliteSynchronous::Normal); + + sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(POOL_MAX_SIZE as u32) + .acquire_timeout(Duration::from_secs(POOL_WAIT_SECS)) + .connect_with(options) + .await + .map_err(|error| VeloxError::Connection(error.to_string())) +} + +// ── MongoDB ───────────────────────────────────────────────────── + +pub fn build_mongo_connection_string(input: &ConnectionInput) -> String { + let scheme = if input.srv_enabled { "mongodb+srv" } else { "mongodb" }; + let host = if input.host.is_empty() { "localhost" } else { &input.host }; + let port = if input.port == 0 { DEFAULT_MONGO_PORT } else { input.port }; + let mut uri = if input.user.is_empty() { + if input.srv_enabled { + format!("{}://{}/", scheme, host) + } else { + format!("{}://{}:{}/", scheme, host, port) + } + } else { + if input.srv_enabled { + format!( + "{}://{}:{}@{}/", + scheme, + urlencoding::encode(&input.user), + urlencoding::encode(&input.password), + host, + ) + } else { + format!( + "{}://{}:{}@{}:{}/", + scheme, + urlencoding::encode(&input.user), + urlencoding::encode(&input.password), + host, + port, + ) + } + }; + let database = if input.database.is_empty() { "admin" } else { &input.database }; + uri.push_str(database); + if let Some(ref params) = input.extra_params { + let qs: Vec = params.iter().map(|(k, v)| format!("{}={}", k, v)).collect(); + if !qs.is_empty() { + uri.push('?'); + uri.push_str(&qs.join("&")); + } + } + uri +} + +pub async fn get_or_create_mongo_client( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result { + { + let clients = state.mongo_clients.read().await; + if let Some(client) = clients.get(connection_id) { + return Ok(client.clone()); + } + } + let stored = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + let uri = build_mongo_connection_string(&stored.to_input()); + let client = MongoClient::with_uri_str(&uri) + .await + .map_err(|e| format!("MongoDB connection failed: {}", e))?; + client + .database("admin") + .run_command(doc! { "ping": 1 }) + .await + .map_err(|e| format!("MongoDB ping failed: {}", e))?; + state.mongo_clients.write().await.insert(connection_id.to_string(), client.clone()); + Ok(client) +} + +// ── DuckDB ────────────────────────────────────────────────────── + +pub fn build_duckdb_connection(input: &ConnectionInput) -> Result { + let path = input.file_path.as_deref().unwrap_or(":memory:"); + let conn = if path == ":memory:" || path.is_empty() { + duckdb::Connection::open_in_memory() + .map_err(|e| format!("DuckDB in-memory connection failed: {}", e))? + } else { + duckdb::Connection::open(path) + .map_err(|e| format!("DuckDB connection to '{}' failed: {}", path, e))? + }; + // Verify the connection works + conn.execute_batch("SELECT 1") + .map_err(|e| format!("DuckDB verification query failed: {}", e))?; + Ok(conn) +} + +pub async fn get_or_create_duckdb_connection( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result<(), VeloxError> { + { + let conns = state.duckdb_connections.read().await; + if conns.contains_key(connection_id) { + return Ok(()); + } + } + let stored = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + let conn = build_duckdb_connection(&stored.to_input())?; + state.duckdb_connections.write().await.insert(connection_id.to_string(), tokio::sync::Mutex::new(conn)); + Ok(()) +} + +// ── Redis ─────────────────────────────────────────────────────── + +pub fn build_redis_url(input: &ConnectionInput) -> String { + let host = if input.host.is_empty() { "127.0.0.1" } else { &input.host }; + let port = if input.port == 0 { DEFAULT_REDIS_PORT } else { input.port }; + if input.user.is_empty() { + format!("redis://{}:{}", host, port) + } else { + format!( + "redis://{}:{}@{}:{}", + urlencoding::encode(&input.user), + urlencoding::encode(&input.password), + host, + port, + ) + } +} + +pub async fn get_or_create_redis_client( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result { + { + let clients = state.redis_clients.read().await; + if let Some(client) = clients.get(connection_id) { + return Ok(client.clone()); + } + } + let stored = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + let url = build_redis_url(&stored.to_input()); + let client = redis::Client::open(url) + .map_err(|e| format!("Redis connection failed: {}", e))?; + let conn = redis::aio::ConnectionManager::new(client) + .await + .map_err(|e| format!("Redis connection failed: {}", e))?; + // Verify with PING + let mut verify = conn.clone(); + redis::cmd("PING") + .query_async::<_, String>(&mut verify) + .await + .map_err(|e| format!("Redis ping failed: {}", e))?; + state.redis_clients.write().await.insert(connection_id.to_string(), conn.clone()); + Ok(conn) +} + +pub async fn drop_pool(state: &AppState, connection_id: &str) { + state.pools.write().await.remove(connection_id); + state.mysql_pools.write().await.remove(connection_id); + state.sqlite_pools.write().await.remove(connection_id); + state.mongo_clients.write().await.remove(connection_id); + state.duckdb_connections.write().await.remove(connection_id); + state.redis_clients.write().await.remove(connection_id); + state + .ask_veloxy_db_context_cache + .write() + .await + .retain(|cache_key, _| !cache_key.starts_with(&format!("{}::", connection_id))); + state + .ask_veloxy_conversations + .write() + .await + .retain(|conversation_key, _| !conversation_key.starts_with(&format!("{}::", connection_id))); + if let Some(mut tunnel) = state.ssh_tunnels.write().await.remove(connection_id) { + tunnel.close().await; + } +} + +pub async fn disconnect_connection(state: &AppState, connection_id: &str) { + drop_pool(state, connection_id).await; + let mut active = state.active_connection_id.write().await; + if active.as_deref() == Some(connection_id) { + *active = None; + } +} + +pub async fn resolve_connection_id( + state: &AppState, + requested_connection_id: Option, +) -> Result { + if let Some(connection_id) = requested_connection_id { + return Ok(connection_id); + } + + state + .active_connection_id + .read() + .await + .clone() + .ok_or_else(|| VeloxError::Validation("Connect to a database before running this action.".to_string())) +} + +pub async fn get_or_create_pool( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result { + if let Some(pool) = state.pools.read().await.get(connection_id).cloned() { + return Ok(pool); + } + + let stored_connection = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + + let input = stored_connection.to_input(); + + let resolved_port = if input.port == 0 { DEFAULT_MYSQL_PORT } else { input.port }; + let (host, port) = if let Some(ref ssh_config) = input.ssh_config { + if ssh_config.is_active() { + let tunnel = SshTunnel::connect(ssh_config, &input.host, resolved_port).await?; + let local_port = tunnel.local_port; + state + .ssh_tunnels + .write() + .await + .insert(connection_id.to_string(), tunnel); + ("127.0.0.1".to_string(), local_port) + } else { + (input.host.clone(), resolved_port) + } + } else { + (input.host.clone(), resolved_port) + }; + + let pool = build_pool_custom(&host, port, &input)?; + + state + .pools + .write() + .await + .insert(connection_id.to_string(), pool.clone()); + + Ok(pool) +} + +pub async fn get_or_create_mysql_pool( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result { + if let Some(pool) = state.mysql_pools.read().await.get(connection_id).cloned() { + return Ok(pool); + } + + let stored_connection = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + + let input = stored_connection.to_input(); + let resolved_port = if input.port == 0 { + DEFAULT_MYSQL_PORT + } else { + input.port + }; + + let (host, port) = if let Some(ref ssh_config) = input.ssh_config { + if ssh_config.is_active() { + let tunnel = SshTunnel::connect(ssh_config, &input.host, resolved_port).await?; + let local_port = tunnel.local_port; + state + .ssh_tunnels + .write() + .await + .insert(connection_id.to_string(), tunnel); + ("127.0.0.1".to_string(), local_port) + } else { + (input.host.clone(), resolved_port) + } + } else { + (input.host.clone(), resolved_port) + }; + + let pool = build_mysql_pool_custom(&host, port, &input).await?; + state + .mysql_pools + .write() + .await + .insert(connection_id.to_string(), pool.clone()); + Ok(pool) +} + +pub async fn get_or_create_sqlite_pool( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result { + if let Some(pool) = state.sqlite_pools.read().await.get(connection_id).cloned() { + return Ok(pool); + } + + let stored_connection = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + + let input = stored_connection.to_input(); + let pool = build_sqlite_pool(&input).await?; + + state + .sqlite_pools + .write() + .await + .insert(connection_id.to_string(), pool.clone()); + + Ok(pool) +} + +/// Drops cached pools and re-validates the connection (used by sidebar refresh). +pub async fn refresh_connection_pools( + app: &AppHandle, + state: &AppState, + connection_id: &str, +) -> Result<(), VeloxError> { + let stored = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + let engine = stored.engine; + + drop_pool(state, connection_id).await; + + match engine { + DatabaseEngine::Postgres => { + retry::with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { + client + .simple_query("select 1") + .await + .map_err(|error| error.to_string())?; + Ok(()) + }) + .await?; + } + DatabaseEngine::Mysql => { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + sqlx::query("select 1") + .execute(&pool) + .await + .map_err(|error| error.to_string())?; + } + DatabaseEngine::Sqlite => { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + sqlx::query("select 1") + .execute(&pool) + .await + .map_err(|error| error.to_string())?; + } + DatabaseEngine::Mongo => { + let client = get_or_create_mongo_client(app, state, connection_id).await?; + client.database("admin").run_command(doc! { "ping": 1 }).await + .map_err(|e| format!("MongoDB ping failed: {}", e))?; + } + DatabaseEngine::Duckdb => { + get_or_create_duckdb_connection(app, state, connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn = conns.get(connection_id).ok_or(VeloxError::Connection("DuckDB connection not found".to_string()))?; + let conn = conn.lock().await; + conn.execute_batch("SELECT 1").map_err(|e| format!("DuckDB ping failed: {}", e))?; + } + DatabaseEngine::Redis => { + let mut client = get_or_create_redis_client(app, state, connection_id).await?; + redis::cmd("PING").query_async::<_, String>(&mut client).await + .map_err(|e| format!("Redis ping failed: {}", e))?; + } + } + + Ok(()) +} + +pub async fn resolve_connection_engine( + app: &AppHandle, + state: &AppState, + requested_connection_id: Option, +) -> Result<(String, DatabaseEngine), String> { + let connection_id = resolve_connection_id(state, requested_connection_id).await?; + let stored = load_connection(app, &connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + Ok((connection_id, stored.engine)) +} + +pub fn list_connections(app: &AppHandle) -> Result, VeloxError> { + let store = app + .store(CONNECTION_STORE_PATH) + .map_err(|error| error.to_string())?; + + let mut connections = store + .entries() + .into_iter() + .map(|(_, value)| { + serde_json::from_value::(value).map_err(|error| error.to_string()) + }) + .collect::, _>>()?; + + connections.sort_by(|left, right| { + right + .connected_at + .cmp(&left.connected_at) + .then_with(|| left.name.cmp(&right.name)) + }); + + Ok(connections + .into_iter() + .map(|connection| connection.summary()) + .collect()) +} + +pub fn persist_connection(app: &AppHandle, connection: &StoredConnection) -> Result<(), VeloxError> { + let store = app + .store(CONNECTION_STORE_PATH) + .map_err(|error| error.to_string())?; + + store.set( + connection.id.clone(), + serde_json::to_value(connection).map_err(|error| error.to_string())?, + ); + + store.save().map_err(|error| error.to_string())?; + + Ok(()) +} + +pub fn persist_connection_with_password( + app: &AppHandle, + connection: &StoredConnection, + password: &str, +) -> Result<(), VeloxError> { + let mut stored = connection.clone(); + stored.password = Some(password.to_string()); + persist_connection(app, &stored)?; + if connection.engine != DatabaseEngine::Sqlite { + credentials::store_password(&connection.id, password)?; + } + Ok(()) +} + +/// Renames a persisted connection and returns the updated summary. +pub fn rename_connection_in_store( + app: &AppHandle, + connection_id: &str, + new_name: &str, +) -> Result { + let mut stored = load_connection(app, connection_id)? + .ok_or_else(|| format!("Connection {} not found.", connection_id))?; + stored.name = new_name.to_string(); + persist_connection(app, &stored)?; + Ok(stored.summary()) +} + +pub fn delete_connection_from_store(app: &AppHandle, connection_id: &str) -> Result<(), VeloxError> { + let store = app + .store(CONNECTION_STORE_PATH) + .map_err(|error| error.to_string())?; + store.delete(connection_id); + store.save().map_err(|error| error.to_string())?; + Ok(()) +} + +pub fn load_connection( + app: &AppHandle, + connection_id: &str, +) -> Result, VeloxError> { + let store = app + .store(CONNECTION_STORE_PATH) + .map_err(|error| error.to_string())?; + + let mut connection: StoredConnection = match store + .get(connection_id) + .map(|value| serde_json::from_value::(value).map_err(|error| error.to_string())) + .transpose()? + { + Some(conn) => conn, + None => return Ok(None), + }; + + let json_password = connection.password.clone(); + + if connection.engine == DatabaseEngine::Sqlite { + connection.password = Some(String::new()); + return Ok(Some(connection)); + } + + match credentials::get_password(connection_id) { + Ok(Some(password)) => { + if json_password.is_some() { + let clean = StoredConnection { + password: None, + ..connection.clone() + }; + let _ = persist_connection(app, &clean); + } + connection.password = Some(password); + } + Ok(None) => { + if let Some(ref pwd) = json_password { + if let Err(e) = credentials::store_password(connection_id, pwd) { + log::warn!("Failed to migrate password to keychain for {}: {}", connection_id, e); + } + connection.password = Some(pwd.clone()); + } else { + return Err(VeloxError::Credential(format!( + "No saved password found for connection {}. Re-enter your password to reconnect.", + connection.name + ))); + } + } + Err(e) => { + log::warn!("Failed to read password from keychain for {}: {}", connection_id, e); + if let Some(ref pwd) = json_password { + connection.password = Some(pwd.clone()); + } else { + return Err(VeloxError::Credential(format!( + "Failed to read saved password for connection {} (keychain error). Re-enter your password to reconnect.", + connection.name + ))); + } + } + } + + Ok(Some(connection)) +} + +#[cfg(test)] +mod tests { + use super::{build_mongo_connection_string, mysql_url}; + use crate::models::{ConnectionInput, ConnectionSslMode, DatabaseEngine}; + use std::collections::HashMap; + + fn mysql_input(ssl_mode: ConnectionSslMode) -> ConnectionInput { + ConnectionInput { + id: None, + name: "test".to_string(), + engine: DatabaseEngine::Mysql, + host: "localhost".to_string(), + port: 3306, + database: "app".to_string(), + file_path: None, + user: "root".to_string(), + password: "pw".to_string(), + srv_enabled: false, + ssl_mode, + ssh_config: None, + extra_params: None, + } + } + + #[allow(dead_code)] + fn postgres_input() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test".to_string(), + engine: DatabaseEngine::Postgres, + host: "localhost".to_string(), + port: 5432, + database: "app".to_string(), + file_path: None, + user: "postgres".to_string(), + password: "secret".to_string(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Prefer, + ssh_config: None, + extra_params: None, + } + } + + #[allow(dead_code)] + fn sqlite_input() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test".to_string(), + engine: DatabaseEngine::Sqlite, + host: String::new(), + port: 0, + database: String::new(), + file_path: Some("/tmp/test.db".to_string()), + user: String::new(), + password: String::new(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ssh_config: None, + extra_params: None, + } + } + + fn mongo_input() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test".to_string(), + engine: DatabaseEngine::Mongo, + host: "localhost".to_string(), + port: 27017, + database: "admin".to_string(), + file_path: None, + user: String::new(), + password: String::new(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ssh_config: None, + extra_params: None, + } + } + + fn mariadb_input() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test".to_string(), + engine: DatabaseEngine::Mysql, + host: "localhost".to_string(), + port: 3306, + database: "app".to_string(), + file_path: None, + user: "root".to_string(), + password: "pw".to_string(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Prefer, + ssh_config: None, + extra_params: None, + } + } + + #[test] + fn mysql_url_maps_ssl_mode() { + assert!(mysql_url("localhost", 3306, &mysql_input(ConnectionSslMode::Disable)) + .ends_with("?ssl-mode=DISABLED")); + assert!(mysql_url("localhost", 3306, &mysql_input(ConnectionSslMode::Prefer)) + .ends_with("?ssl-mode=PREFERRED")); + assert!(mysql_url("localhost", 3306, &mysql_input(ConnectionSslMode::Require)) + .ends_with("?ssl-mode=REQUIRED")); + } + + #[test] + fn mariadb_uses_mysql_url_format() { + // MariaDB is wire-compatible with MySQL — uses the same URL format + let url = mysql_url("localhost", 3306, &mariadb_input()); + assert!(url.starts_with("mysql://")); + assert!(url.contains("root:pw@localhost:3306/app")); + } + + #[test] + fn mongo_connection_string_no_auth() { + let uri = build_mongo_connection_string(&mongo_input()); + assert_eq!(uri, "mongodb://localhost:27017/admin"); + } + + #[test] + fn mongo_connection_string_with_auth() { + let input = ConnectionInput { + user: "admin".to_string(), + password: "secret".to_string(), + ..mongo_input() + }; + let uri = build_mongo_connection_string(&input); + assert!(uri.contains("admin:secret@")); + assert!(uri.contains("localhost:27017/")); + } + + #[test] + fn mongo_connection_string_with_custom_port() { + let input = ConnectionInput { + port: 27018, + ..mongo_input() + }; + let uri = build_mongo_connection_string(&input); + assert_eq!(uri, "mongodb://localhost:27018/admin"); + } + + #[test] + fn mongo_connection_string_with_extra_params() { + let mut params = HashMap::new(); + params.insert("replicaSet".to_string(), "rs0".to_string()); + params.insert("authSource".to_string(), "admin".to_string()); + let input = ConnectionInput { + extra_params: Some(params), + ..mongo_input() + }; + let uri = build_mongo_connection_string(&input); + assert!(uri.contains("replicaSet=rs0")); + assert!(uri.contains("authSource=admin")); + } + + #[test] + fn mongo_connection_defaults_to_admin_database() { + let input = ConnectionInput { + database: String::new(), + ..mongo_input() + }; + let uri = build_mongo_connection_string(&input); + assert!(uri.ends_with("/admin")); + } + + #[test] + fn mongo_connection_defaults_to_localhost() { + let input = ConnectionInput { + host: String::new(), + ..mongo_input() + }; + let uri = build_mongo_connection_string(&input); + assert!(uri.contains("localhost")); + } +} diff --git a/src-tauri/src/db/identifier_guard.rs b/src-tauri/src/db/identifier_guard.rs new file mode 100644 index 0000000..fcf718e --- /dev/null +++ b/src-tauri/src/db/identifier_guard.rs @@ -0,0 +1,49 @@ +use crate::error::VeloxError; + +pub fn quote_identifier(value: &str) -> String { + value.replace('"', "\"\"") +} + +/// Whether `name` is safe to interpolate into dynamic SQL (e.g. SQLite `PRAGMA` +/// statements, where bind parameters are not allowed). Restricted to ASCII +/// alphanumerics and underscores so it cannot terminate or escape a statement. +pub fn is_safe_identifier(name: &str) -> bool { + !name.is_empty() + && name.len() <= 128 + && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') +} + +/// Validates an identifier before it is interpolated into dynamic SQL. Returns +/// the identifier unchanged when safe, or a descriptive error otherwise. +pub fn require_safe_identifier<'a>(name: &'a str, context: &str) -> Result<&'a str, VeloxError> { + if is_safe_identifier(name) { + Ok(name) + } else { + Err(VeloxError::Validation(format!( + "Invalid identifier for {}: {:?}", + context, name + ))) + } +} + +#[cfg(test)] +mod tests { + use super::{is_safe_identifier, require_safe_identifier}; + + #[test] + fn rejects_sql_injection_in_identifiers() { + assert!(!is_safe_identifier("\"; DROP TABLE users; --")); + assert!(!is_safe_identifier("foo; DELETE FROM bar")); + assert!(!is_safe_identifier("a\0b")); + assert!(!is_safe_identifier("")); + assert!(require_safe_identifier("foo);", "table name").is_err()); + } + + #[test] + fn accepts_plain_identifiers() { + assert!(is_safe_identifier("users")); + assert!(is_safe_identifier("_internal")); + assert!(is_safe_identifier("Table123")); + assert_eq!(require_safe_identifier("users", "table name").unwrap(), "users"); + } +} diff --git a/src-tauri/src/db/retry.rs b/src-tauri/src/db/retry.rs new file mode 100644 index 0000000..924bf5d --- /dev/null +++ b/src-tauri/src/db/retry.rs @@ -0,0 +1,72 @@ +use deadpool_postgres::Client; +use tauri::AppHandle; + +use super::connection_pool::{self, AppState}; +use crate::error::VeloxError; + +/// Heuristic for transport-level failures where discarding the pool and opening +/// a new TCP session may succeed (sleep/VPN blips, idle disconnects). +pub fn is_retryable_connection_error(message: &str) -> bool { + let m = message.to_lowercase(); + m.contains("broken pipe") + || m.contains("connection reset") + || m.contains("connection refused") + || m.contains("unexpected eof") + || m.contains("unexpected end of file") + || m.contains("error communicating with the server") + || m.contains("connection closed") + || m.contains("closed the connection") + || m.contains("server closed the connection") + || m.contains("eof has been reached") + || m.contains("could not receive data from server") + || m.contains("could not send data to server") + || m.contains("timeout occurred while waiting") + || m.contains("timeout occurred while creating") + || m.contains("timeout occurred while recycling") +} + +/// Runs `operation` with a pooled client. On a retryable connection error, drops +/// the cached pool once and retries the whole operation (at most one extra attempt). +pub async fn with_pool_client_retry( + app: &AppHandle, + state: &AppState, + connection_id: &str, + ctx: C, + mut operation: F, +) -> Result +where + C: Clone + Send, + F: FnMut(Client, C) -> Fut, + Fut: std::future::Future> + Send, + T: Send, +{ + let mut dropped_pool = false; + + loop { + let pool = connection_pool::get_or_create_pool(app, state, connection_id).await?; + let client = match pool.get().await { + Ok(client) => client, + Err(error) => { + let message = error.to_string(); + if !dropped_pool && is_retryable_connection_error(&message) { + connection_pool::drop_pool(state, connection_id).await; + dropped_pool = true; + continue; + } + return Err(VeloxError::Connection(message)); + } + }; + + match operation(client, ctx.clone()).await { + Ok(value) => return Ok(value), + Err(message) => { + if !dropped_pool && is_retryable_connection_error(&message) { + connection_pool::drop_pool(state, connection_id).await; + dropped_pool = true; + continue; + } + return Err(VeloxError::Query(message)); + } + } + } +} diff --git a/src-tauri/src/engines/duckdb.rs b/src-tauri/src/engines/duckdb.rs new file mode 100644 index 0000000..3db3813 --- /dev/null +++ b/src-tauri/src/engines/duckdb.rs @@ -0,0 +1,190 @@ +use std::collections::BTreeMap; +use std::time::Instant; + +use tauri::AppHandle; + +use crate::db::{ + build_duckdb_connection, get_or_create_duckdb_connection, AppState, +}; +use crate::error::VeloxError; +use crate::models::{ColumnInfo, ConnectionInput, DatabaseInfo, QueryResult, TableInfo}; + +use super::DatabaseEngineOps; + +pub struct DuckdbEngine; + +impl DatabaseEngineOps for DuckdbEngine { + async fn connect( + &self, + _app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError> { + let conn = build_duckdb_connection(input)?; + state.duckdb_connections.write().await.insert( + connection_id.to_string(), + tokio::sync::Mutex::new(conn), + ); + Ok(()) + } + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError> { + get_or_create_duckdb_connection(app, state, connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn = conns.get(connection_id) + .ok_or_else(|| VeloxError::Connection("DuckDB connection not found".to_string()))?; + let conn = conn.lock().await; + conn.execute_batch("SELECT 1") + .map_err(|e| VeloxError::Connection(format!("DuckDB ping failed: {}", e)))?; + Ok(()) + } + + async fn run_query( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_rows: usize, + ) -> Result { + get_or_create_duckdb_connection(app, state, connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(connection_id) + .ok_or_else(|| VeloxError::Connection("DuckDB connection not found".to_string()))?; + let conn = conn_mutex.lock().await; + + let is_select = sql.trim().to_uppercase().starts_with("SELECT") + || sql.trim().to_uppercase().starts_with("WITH"); + + let started_at = Instant::now(); + + if is_select { + let stmt = conn.prepare(sql) + .map_err(|e| VeloxError::Query(e.to_string()))?; + let col_count = stmt.column_count(); + let columns: Vec = (0..col_count) + .map(|i| stmt.column_name(i).map_or("?".to_string(), |v| v.to_string())) + .collect(); + + let mut rows = Vec::new(); + let mut stmt = conn.prepare(sql) + .map_err(|e| VeloxError::Query(e.to_string()))?; + let results = stmt.query_map([], |row| { + let vals: Vec> = (0..col_count).map(|i| { + row.get::<_, Option>(i).ok().flatten() + }).collect(); + Ok(vals) + }).map_err(|e| VeloxError::Query(e.to_string()))?; + + let mut total = 0usize; + for result in results { + let vals = result.map_err(|e| VeloxError::Query(e.to_string()))?; + if rows.len() < max_rows { + let mut map = BTreeMap::new(); + for (i, col) in columns.iter().enumerate() { + map.insert(col.clone(), vals[i].clone()); + } + rows.push(map); + } + total += 1; + } + Ok(QueryResult { + columns, + row_count: rows.len(), + rows, + execution_ms: started_at.elapsed().as_millis(), + truncated: total > max_rows, + command_tag: None, + }) + } else { + let affected = conn.execute(sql, []) + .map_err(|e| VeloxError::Query(e.to_string()))?; + Ok(QueryResult { + columns: vec![], + row_count: affected, + rows: vec![], + execution_ms: started_at.elapsed().as_millis(), + truncated: false, + command_tag: Some(affected as u64), + }) + } + } + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + get_or_create_duckdb_connection(app, state, connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(connection_id) + .ok_or_else(|| VeloxError::Connection("DuckDB connection not found".to_string()))?; + let conn = conn_mutex.lock().await; + + let mut stmt = conn.prepare( + "SELECT table_name FROM information_schema.tables WHERE table_schema='main' ORDER BY table_name" + ).map_err(|e| VeloxError::Query(e.to_string()))?; + + let names: Vec = stmt.query_map([], |row| row.get(0)) + .map_err(|e| VeloxError::Query(e.to_string()))? + .filter_map(|r| r.ok()) + .collect(); + + Ok(names.into_iter().map(|name| TableInfo { + schema: "main".to_string(), + name: name.clone(), + preview_query: format!("SELECT * FROM \"{}\" LIMIT 100;", name), + }).collect()) + } + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + _table_schema: &str, + table_name: &str, + ) -> Result, VeloxError> { + get_or_create_duckdb_connection(app, state, connection_id).await?; + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(connection_id) + .ok_or_else(|| VeloxError::Connection("DuckDB connection not found".to_string()))?; + let conn = conn_mutex.lock().await; + + let mut stmt = conn.prepare( + "SELECT column_name, data_type FROM information_schema.columns \ + WHERE table_name=? ORDER BY ordinal_position" + ).map_err(|e| VeloxError::Query(e.to_string()))?; + + let cols: Vec<(String, String)> = stmt.query_map( + duckdb::params![table_name], + |row| Ok((row.get(0)?, row.get(1)?)), + ).map_err(|e| VeloxError::Query(e.to_string()))? + .filter_map(|r| r.ok()) + .collect(); + + Ok(cols.into_iter().map(|(name, dtype)| ColumnInfo { + table_schema: "main".to_string(), + table_name: table_name.to_string(), + column_name: name, + data_type: dtype, + is_nullable: true, + }).collect()) + } + + async fn list_databases( + &self, + _app: &AppHandle, + _state: &AppState, + _connection_id: &str, + ) -> Result, VeloxError> { + Ok(vec![DatabaseInfo { name: "main".to_string() }]) + } +} diff --git a/src-tauri/src/engines/mod.rs b/src-tauri/src/engines/mod.rs new file mode 100644 index 0000000..b5d778d --- /dev/null +++ b/src-tauri/src/engines/mod.rs @@ -0,0 +1,130 @@ +mod duckdb; +mod mongo; +mod mysql; +mod postgres; +mod redis; +mod sqlite; + +use tauri::AppHandle; + +use crate::db::AppState; +use crate::error::VeloxError; +use crate::models::{ + ColumnInfo, ConnectionInput, DatabaseEngine, DatabaseInfo, QueryResult, TableInfo, +}; + +use duckdb::DuckdbEngine; +use mongo::MongoEngine; +use mysql::MySqlEngine; +use postgres::PostgresEngine; +use redis::RedisEngine; +use sqlite::SqliteEngine; + +/// Per-engine operations implemented by each engine unit struct. +/// Dispatch goes through [`EngineDispatcher`] (static enum, no boxing). +#[allow(async_fn_in_trait)] +pub trait DatabaseEngineOps: Send + Sync { + async fn connect( + &self, + app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError>; + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError>; + + async fn run_query( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_rows: usize, + ) -> Result; + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError>; + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + table_schema: &str, + table_name: &str, + ) -> Result, VeloxError>; + + async fn list_databases( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError>; +} + +/// Static-dispatch enum over all engine implementations. +/// Use [`get_engine`] to construct the right variant. +pub enum EngineDispatcher { + Postgres(PostgresEngine), + Mysql(MySqlEngine), + Sqlite(SqliteEngine), + Mongo(MongoEngine), + Duckdb(DuckdbEngine), + Redis(RedisEngine), +} + +macro_rules! delegate { + ($self:ident, $method:ident, $($arg:expr),* $(,)?) => { + match $self { + EngineDispatcher::Postgres(e) => e.$method($($arg),*).await, + EngineDispatcher::Mysql(e) => e.$method($($arg),*).await, + EngineDispatcher::Sqlite(e) => e.$method($($arg),*).await, + EngineDispatcher::Mongo(e) => e.$method($($arg),*).await, + EngineDispatcher::Duckdb(e) => e.$method($($arg),*).await, + EngineDispatcher::Redis(e) => e.$method($($arg),*).await, + } + }; +} + +impl DatabaseEngineOps for EngineDispatcher { + async fn connect(&self, app: &AppHandle, state: &AppState, input: &ConnectionInput, connection_id: &str) -> Result<(), VeloxError> { + delegate!(self, connect, app, state, input, connection_id) + } + async fn ping(&self, app: &AppHandle, state: &AppState, connection_id: &str) -> Result<(), VeloxError> { + delegate!(self, ping, app, state, connection_id) + } + async fn run_query(&self, app: &AppHandle, state: &AppState, connection_id: &str, sql: &str, max_rows: usize) -> Result { + delegate!(self, run_query, app, state, connection_id, sql, max_rows) + } + async fn get_tables(&self, app: &AppHandle, state: &AppState, connection_id: &str) -> Result, VeloxError> { + delegate!(self, get_tables, app, state, connection_id) + } + async fn get_schema(&self, app: &AppHandle, state: &AppState, connection_id: &str, table_schema: &str, table_name: &str) -> Result, VeloxError> { + delegate!(self, get_schema, app, state, connection_id, table_schema, table_name) + } + async fn list_databases(&self, app: &AppHandle, state: &AppState, connection_id: &str) -> Result, VeloxError> { + delegate!(self, list_databases, app, state, connection_id) + } +} + +/// Return the right engine dispatcher for the given discriminant. +pub fn get_engine(engine: DatabaseEngine) -> EngineDispatcher { + match engine { + DatabaseEngine::Postgres => EngineDispatcher::Postgres(PostgresEngine), + DatabaseEngine::Mysql => EngineDispatcher::Mysql(MySqlEngine), + DatabaseEngine::Sqlite => EngineDispatcher::Sqlite(SqliteEngine), + DatabaseEngine::Mongo => EngineDispatcher::Mongo(MongoEngine), + DatabaseEngine::Duckdb => EngineDispatcher::Duckdb(DuckdbEngine), + DatabaseEngine::Redis => EngineDispatcher::Redis(RedisEngine), + } +} diff --git a/src-tauri/src/engines/mongo.rs b/src-tauri/src/engines/mongo.rs new file mode 100644 index 0000000..3df786b --- /dev/null +++ b/src-tauri/src/engines/mongo.rs @@ -0,0 +1,287 @@ +use std::collections::BTreeMap; +use std::time::Instant; + +use mongodb::bson::{self, doc, Document}; +use tauri::AppHandle; + +use crate::db::{ + build_mongo_connection_string, get_or_create_mongo_client, load_connection, AppState, +}; +use crate::error::VeloxError; +use crate::models::{ColumnInfo, ConnectionInput, DatabaseInfo, QueryResult, TableInfo}; + +use super::DatabaseEngineOps; + +pub struct MongoEngine; + +impl DatabaseEngineOps for MongoEngine { + async fn connect( + &self, + _app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError> { + let uri = build_mongo_connection_string(input); + let client = mongodb::Client::with_uri_str(&uri) + .await + .map_err(|e| VeloxError::Connection(format!("MongoDB connection failed: {}", e)))?; + client + .database("admin") + .run_command(doc! { "ping": 1 }) + .await + .map_err(|e| VeloxError::Connection(format!("MongoDB ping failed: {}", e)))?; + state.mongo_clients.write().await.insert(connection_id.to_string(), client); + Ok(()) + } + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError> { + let client = get_or_create_mongo_client(app, state, connection_id).await?; + client.database("admin").run_command(doc! { "ping": 1 }).await + .map_err(|e| VeloxError::Connection(format!("MongoDB ping failed: {}", e)))?; + Ok(()) + } + + async fn run_query( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_rows: usize, + ) -> Result { + let client = get_or_create_mongo_client(app, state, connection_id).await?; + let stored = load_connection(app, connection_id)? + .ok_or_else(|| VeloxError::Connection("Stored connection details were not found.".to_string()))?; + let db = client.default_database().ok_or_else(|| { + VeloxError::Connection("No default database for MongoDB connection.".to_string()) + })?; + + // Parse: "db.collection.find({...})" or just collection name + let trimmed = sql.trim(); + let (db_name, coll_name, filter) = parse_mongo_query(trimmed, &stored.database); + + let db = if db_name == stored.database { db } else { client.database(&db_name) }; + let collection = db.collection::(&coll_name); + + let started_at = Instant::now(); + let mut cursor = collection + .find(filter) + .limit(max_rows as i64) + .await + .map_err(|e| VeloxError::Query(format!("MongoDB find failed: {}", e)))?; + + let mut rows: Vec>> = Vec::new(); + let mut columns_set = std::collections::HashSet::new(); + let mut columns = Vec::new(); + + while let Ok(true) = cursor.advance().await { + let doc = cursor.deserialize_current() + .map_err(|e| VeloxError::Query(format!("MongoDB cursor error: {}", e)))?; + let mut row = BTreeMap::new(); + for (key, value) in &doc { + if !columns_set.contains(key) { + columns.push(key.clone()); + columns_set.insert(key.clone()); + } + row.insert(key.clone(), bson_to_display(value)); + } + for col in &columns { + row.entry(col.clone()).or_insert(None); + } + rows.push(row); + } + + Ok(QueryResult { + row_count: rows.len(), + rows, + columns, + execution_ms: started_at.elapsed().as_millis(), + truncated: false, + command_tag: None, + }) + } + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + let client = get_or_create_mongo_client(app, state, connection_id).await?; + let db = client.default_database().ok_or_else(|| { + VeloxError::Connection("No default database for MongoDB connection.".to_string()) + })?; + let db_name = db.name().to_string(); + let names = db.list_collection_names().await + .map_err(|e| VeloxError::Query(format!("Failed to list MongoDB collections: {}", e)))?; + Ok(names.into_iter().map(|name| TableInfo { + schema: db_name.clone(), + name: name.clone(), + preview_query: format!("{}.find({{}}).limit(100)", name), + }).collect()) + } + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + ) -> Result, VeloxError> { + let client = get_or_create_mongo_client(app, state, connection_id).await?; + let db = client.database(database); + let coll = db.collection::(collection); + let mut cursor = coll.find(doc! {}).limit(200).await + .map_err(|e| VeloxError::Query(format!("MongoDB cursor error: {}", e)))?; + + let mut field_counts: BTreeMap = BTreeMap::new(); + let mut field_types: BTreeMap = BTreeMap::new(); + while let Ok(true) = cursor.advance().await { + let doc = cursor.deserialize_current() + .map_err(|e| VeloxError::Query(format!("MongoDB cursor error: {}", e)))?; + for (key, value) in &doc { + let entry = field_counts.entry(key.clone()).or_default(); + *entry += 1; + // Track the first non-null BSON type seen for each field + let btype = infer_bson_type(value); + if btype != "null" { + field_types.entry(key.clone()).or_insert(btype); + } + } + } + + let total = field_counts.values().max().copied().unwrap_or(1).max(1); + Ok(field_counts.into_iter().map(|(name, count)| { + let btype = field_types.get(&name).copied().unwrap_or("unknown"); + ColumnInfo { + table_schema: database.to_string(), + table_name: collection.to_string(), + column_name: name, + data_type: format!("{} ({}% present)", btype, count * 100 / total), + is_nullable: count < total, + } + }).collect()) + } + + async fn list_databases( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + let client = get_or_create_mongo_client(app, state, connection_id).await?; + let db_names = client.list_database_names().await + .map_err(|e| VeloxError::Query(format!("Failed to list MongoDB databases: {}", e)))?; + Ok(db_names.into_iter().map(|name| DatabaseInfo { name }).collect()) + } +} + +fn infer_bson_type(value: &mongodb::bson::Bson) -> &'static str { + use mongodb::bson::Bson; + match value { + Bson::String(_) => "string", + Bson::Int32(_) | Bson::Int64(_) => "integer", + Bson::Double(_) => "double", + Bson::Boolean(_) => "boolean", + Bson::ObjectId(_) => "objectId", + Bson::DateTime(_) => "date", + Bson::Array(_) => "array", + Bson::Document(_) => "object", + Bson::Null => "null", + Bson::Binary(_) => "binary", + Bson::RegularExpression(_) => "regex", + Bson::Decimal128(_) => "decimal", + _ => "unknown", + } +} + +fn parse_mongo_query(sql: &str, default_db: &str) -> (String, String, Document) { + let trimmed = sql.trim(); + // "db.collection.find({...})" format + if let Some(rest) = trimmed.strip_prefix("db.") { + if let Some(dot) = rest.find('.') { + let collection = &rest[..dot]; + let after = &rest[dot + 1..]; + if let Some(filter_part) = after.strip_prefix("find(") { + let filter_str = filter_part.trim_end_matches(')').trim(); + let filter = parse_filter(filter_str); + return (default_db.to_string(), collection.to_string(), filter); + } + } + } + // "database.collection" format + if let Some(dot) = trimmed.find('.') { + let db = &trimmed[..dot]; + let coll = &trimmed[dot + 1..]; + return (db.to_string(), coll.to_string(), doc! {}); + } + // Just collection name + (default_db.to_string(), trimmed.to_string(), doc! {}) +} + +fn parse_filter(input: &str) -> Document { + let trimmed = input.trim(); + if trimmed.is_empty() || trimmed == "{}" { + return doc! {}; + } + // Try JSON + if let Ok(doc) = bson::from_slice(trimmed.as_bytes()) { + return doc; + } + // Try as relaxed JSON (single quotes) + let json_like = trimmed.replace('\'', "\""); + if let Ok(doc) = bson::from_slice(json_like.as_bytes()) { + return doc; + } + // Key-value pairs: "key: value, key2: value2" + let mut doc = Document::new(); + for part in trimmed.split(',') { + let part = part.trim(); + if let Some(colon) = part.find(':') { + let key = part[..colon].trim().trim_matches('"').trim_matches('\''); + let val = part[colon + 1..].trim().trim_matches('"').trim_matches('\''); + if let Ok(n) = val.parse::() { + doc.insert(key, n); + } else if val == "true" { + doc.insert(key, true); + } else if val == "false" { + doc.insert(key, false); + } else { + doc.insert(key, val); + } + } + } + doc +} + +fn bson_to_display(value: &bson::Bson) -> Option { + match value { + bson::Bson::Double(v) => Some(v.to_string()), + bson::Bson::String(v) => Some(v.clone()), + bson::Bson::Boolean(v) => Some(v.to_string()), + bson::Bson::Null => None, + bson::Bson::Int32(v) => Some(v.to_string()), + bson::Bson::Int64(v) => Some(v.to_string()), + bson::Bson::ObjectId(v) => Some(v.to_hex()), + bson::Bson::DateTime(v) => Some(v.to_string()), + bson::Bson::Binary(v) => Some(hex::encode(&v.bytes)), + bson::Bson::Array(arr) => { + let items: Vec = arr.iter().filter_map(bson_to_display).collect(); + Some(format!("[{}]", items.join(", "))) + } + bson::Bson::Document(doc) => { + let items: Vec = doc.iter() + .map(|(k, v)| format!("{}: {}", k, bson_to_display(v).unwrap_or_else(|| "null".to_string()))) + .collect(); + Some(format!("{{{}}}", items.join(", "))) + } + _ => Some(format!("{:?}", value)), + } +} diff --git a/src-tauri/src/engines/mysql.rs b/src-tauri/src/engines/mysql.rs new file mode 100644 index 0000000..24071c7 --- /dev/null +++ b/src-tauri/src/engines/mysql.rs @@ -0,0 +1,166 @@ +use tauri::AppHandle; + +use crate::db::{ + build_mysql_pool, build_mysql_pool_custom, get_or_create_mysql_pool, load_connection, + AppState, DEFAULT_MYSQL_PORT, +}; +use crate::error::VeloxError; +use crate::models::{ColumnInfo, ConnectionInput, DatabaseInfo, QueryResult, TableInfo}; +use crate::ssh_tunnel::SshTunnel; + +use super::DatabaseEngineOps; + +use crate::commands::{ + mysql_get_string, mysql_database_name_from_row, run_query_mysql_or_sqlite, +}; + +pub struct MySqlEngine; + +impl DatabaseEngineOps for MySqlEngine { + async fn connect( + &self, + _app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError> { + let pool = if let Some(ref ssh_config) = input.ssh_config { + if ssh_config.is_active() { + let remote_port = if input.port == 0 { DEFAULT_MYSQL_PORT } else { input.port }; + let tunnel = SshTunnel::connect(ssh_config, &input.host, remote_port) + .await + .map_err(|e| VeloxError::Connection(format!("SSH tunnel failed: {}", e)))?; + let local_port = tunnel.local_port; + state.ssh_tunnels.write().await.insert(connection_id.to_string(), tunnel); + build_mysql_pool_custom("127.0.0.1", local_port, input).await? + } else { + build_mysql_pool(input).await? + } + } else { + build_mysql_pool(input).await? + }; + + sqlx::query("select 1").execute(&pool).await + .map_err(|e| VeloxError::Connection(e.to_string()))?; + state.mysql_pools.write().await.insert(connection_id.to_string(), pool); + Ok(()) + } + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError> { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + sqlx::query("select 1").execute(&pool).await + .map_err(|e| VeloxError::Connection(e.to_string()))?; + Ok(()) + } + + async fn run_query( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_rows: usize, + ) -> Result { + run_query_mysql_or_sqlite( + app, + state, + connection_id, + sql, + max_rows, + crate::models::DatabaseEngine::Mysql, + ) + .await + .map_err(VeloxError::from) + } + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + let database = load_connection(app, connection_id)? + .map(|connection| connection.database).unwrap_or_default(); + let rows = sqlx::query( + "select table_schema, table_name \ + from information_schema.tables \ + where table_type = 'BASE TABLE' \ + and table_schema = ? \ + and table_schema not in ('information_schema', 'mysql', 'performance_schema', 'sys') \ + order by table_schema, table_name", + ).bind(&database).fetch_all(&pool).await + .map_err(|e| VeloxError::Query(e.to_string()))?; + let mut tables = Vec::new(); + for row in rows { + let schema: String = mysql_get_string(&row, 0, "table_schema", "get_tables") + .map_err(VeloxError::from)?; + let name: String = mysql_get_string(&row, 1, "table_name", "get_tables") + .map_err(VeloxError::from)?; + tables.push(TableInfo { + preview_query: format!("select * from `{}`.`{}` limit 100;", schema, name), + schema, name, + }); + } + Ok(tables) + } + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + table_schema: &str, + table_name: &str, + ) -> Result, VeloxError> { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + let rows = sqlx::query( + "select table_schema, table_name, column_name, data_type, is_nullable \ + from information_schema.columns \ + where table_schema = ? and table_name = ? \ + order by ordinal_position", + ).bind(table_schema).bind(table_name) + .fetch_all(&pool).await + .map_err(|e| VeloxError::Query(e.to_string()))?; + let mut columns = Vec::new(); + for row in rows { + columns.push(ColumnInfo { + table_schema: mysql_get_string(&row, 0, "table_schema", "get_schema") + .map_err(VeloxError::from)?, + table_name: mysql_get_string(&row, 1, "table_name", "get_schema") + .map_err(VeloxError::from)?, + column_name: mysql_get_string(&row, 2, "column_name", "get_schema") + .map_err(VeloxError::from)?, + data_type: mysql_get_string(&row, 3, "data_type", "get_schema") + .map_err(VeloxError::from)?, + is_nullable: mysql_get_string(&row, 4, "is_nullable", "get_schema") + .map_err(VeloxError::from)? == "YES", + }); + } + Ok(columns) + } + + async fn list_databases( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + let pool = get_or_create_mysql_pool(app, state, connection_id).await?; + let rows = sqlx::query("show databases") + .fetch_all(&pool).await + .map_err(|e| VeloxError::Query(e.to_string()))?; + let mut databases = Vec::with_capacity(rows.len()); + for row in rows { + let name = mysql_database_name_from_row(&row, "list_databases") + .map_err(VeloxError::from)?; + databases.push(DatabaseInfo { name }); + } + Ok(databases) + } +} diff --git a/src-tauri/src/engines/postgres.rs b/src-tauri/src/engines/postgres.rs new file mode 100644 index 0000000..96faa3a --- /dev/null +++ b/src-tauri/src/engines/postgres.rs @@ -0,0 +1,203 @@ +use std::time::Instant; + +use tauri::AppHandle; +use tokio_postgres::SimpleQueryMessage; + +use crate::db::{ + build_pool, build_pool_custom, drop_pool, quote_identifier, + with_pool_client_retry, AppState, +}; +use crate::error::VeloxError; +use crate::models::{ColumnInfo, ConnectionInput, DatabaseInfo, QueryResult, TableInfo}; +use crate::pg_error::map_pg_err; +use crate::ssh_tunnel::SshTunnel; + +use super::DatabaseEngineOps; + +pub struct PostgresEngine; + +impl DatabaseEngineOps for PostgresEngine { + async fn connect( + &self, + _app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError> { + let pool = if let Some(ref ssh_config) = input.ssh_config { + if ssh_config.is_active() { + let tunnel = SshTunnel::connect(ssh_config, &input.host, input.port) + .await + .map_err(|e| VeloxError::Connection(format!("SSH tunnel failed: {}", e)))?; + let local_port = tunnel.local_port; + state.ssh_tunnels.write().await.insert(connection_id.to_string(), tunnel); + build_pool_custom("127.0.0.1", local_port, input)? + } else { + build_pool(input)? + } + } else { + build_pool(input)? + }; + + let client = pool.get().await.map_err(|e| { + VeloxError::Connection(e.to_string()) + })?; + + if let Err(e) = client.simple_query("select 1").await { + drop_pool(state, connection_id).await; + return Err(VeloxError::Postgres(crate::error::PgError::from_error(e, None))); + } + + state.pools.write().await.insert(connection_id.to_string(), pool); + Ok(()) + } + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError> { + with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { + client.simple_query("select 1").await.map_err(|error| error.to_string())?; + Ok(()) + }) + .await + } + + async fn run_query( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_rows: usize, + ) -> Result { + with_pool_client_retry(app, state, connection_id, sql.to_string(), |client, sql| async move { + let started_at = Instant::now(); + let messages = client.simple_query(&sql).await + .map_err(|error| map_pg_err(error, Some(sql.as_str())))?; + + let mut columns = Vec::new(); + let mut rows = Vec::new(); + let mut total_rows = 0usize; + let mut command_tag = None; + + for message in messages { + match message { + SimpleQueryMessage::RowDescription(description) => { + if columns.is_empty() { + columns = description.iter() + .map(|column| column.name().to_string()).collect(); + } + } + SimpleQueryMessage::Row(row) => { + total_rows += 1; + if columns.is_empty() { + columns = row.columns().iter() + .map(|column| column.name().to_string()).collect(); + } + if rows.len() >= max_rows { continue; } + let mut mapped_row = std::collections::BTreeMap::new(); + for (index, column_name) in columns.iter().enumerate() { + mapped_row.insert(column_name.clone(), row.get(index).map(str::to_owned)); + } + rows.push(mapped_row); + } + SimpleQueryMessage::CommandComplete(count) => { + command_tag = Some(count); + } + _ => {} + } + } + + Ok(QueryResult { + columns, + row_count: rows.len(), + rows, + execution_ms: started_at.elapsed().as_millis(), + truncated: total_rows > max_rows, + command_tag, + }) + }).await + } + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { + let rows = client.query( + "select t.table_schema, t.table_name \ + from information_schema.tables t \ + join pg_catalog.pg_class c on c.oid = (quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::regclass \ + where t.table_type = 'BASE TABLE' \ + and t.table_schema not in ('pg_catalog', 'information_schema') \ + and c.relispartition = false \ + order by t.table_schema, t.table_name", + &[], + ).await.map_err(|error| map_pg_err(error, None))?; + + Ok(rows.into_iter().map(|row| { + let schema: String = row.get(0); + let name: String = row.get(1); + let preview_query = format!( + "select * from \"{}\".\"{}\" limit 100;", + quote_identifier(&schema), quote_identifier(&name) + ); + TableInfo { schema, name, preview_query } + }).collect()) + }).await + } + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + table_schema: &str, + table_name: &str, + ) -> Result, VeloxError> { + let ts = table_schema.to_string(); + let tn = table_name.to_string(); + with_pool_client_retry(app, state, connection_id, (ts, tn), |client, (schema, name)| async move { + let rows = client.query( + "select table_schema, table_name, column_name, data_type, is_nullable \ + from information_schema.columns \ + where table_schema = $1 and table_name = $2 \ + order by ordinal_position", + &[&schema, &name], + ).await.map_err(|error| map_pg_err(error, None))?; + + Ok(rows.into_iter().map(|row| ColumnInfo { + table_schema: row.get(0), + table_name: row.get(1), + column_name: row.get(2), + data_type: row.get(3), + is_nullable: row.get::<_, String>(4) == "YES", + }).collect()) + }).await + } + + async fn list_databases( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + with_pool_client_retry(app, state, connection_id, (), |client, ()| async move { + let rows = client.query( + "select datname from pg_database \ + where datistemplate = false and has_database_privilege(datname, 'CONNECT') \ + order by datname", + &[], + ).await.map_err(|error| map_pg_err(error, None))?; + Ok(rows.into_iter().map(|row| { + let name: String = row.get(0); + DatabaseInfo { name } + }).collect()) + }).await + } +} diff --git a/src-tauri/src/engines/redis.rs b/src-tauri/src/engines/redis.rs new file mode 100644 index 0000000..34a450e --- /dev/null +++ b/src-tauri/src/engines/redis.rs @@ -0,0 +1,120 @@ +use tauri::AppHandle; + +use crate::db::{build_redis_url, get_or_create_redis_client, AppState}; +use crate::error::VeloxError; +use crate::models::{ColumnInfo, ConnectionInput, DatabaseInfo, QueryResult, TableInfo}; + +use super::DatabaseEngineOps; + +pub struct RedisEngine; + +impl DatabaseEngineOps for RedisEngine { + async fn connect( + &self, + _app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError> { + let url = build_redis_url(input); + let client = redis::Client::open(url) + .map_err(|e| VeloxError::Connection(format!("Redis connection failed: {}", e)))?; + let mut conn = redis::aio::ConnectionManager::new(client).await + .map_err(|e| VeloxError::Connection(format!("Redis connection failed: {}", e)))?; + redis::cmd("PING").query_async::<_, String>(&mut conn).await + .map_err(|e| VeloxError::Connection(format!("Redis ping failed: {}", e)))?; + state.redis_clients.write().await.insert(connection_id.to_string(), conn); + Ok(()) + } + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError> { + let mut client = get_or_create_redis_client(app, state, connection_id).await?; + redis::cmd("PING").query_async::<_, String>(&mut client).await + .map_err(|e| VeloxError::Connection(format!("Redis ping failed: {}", e)))?; + Ok(()) + } + + async fn run_query( + &self, + _app: &AppHandle, + _state: &AppState, + _connection_id: &str, + _sql: &str, + _max_rows: usize, + ) -> Result { + Err(VeloxError::Validation( + "Redis uses its own command path. Use the Redis-specific UI.".to_string(), + )) + } + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + let mut client = get_or_create_redis_client(app, state, connection_id).await?; + let keys: Vec = redis::cmd("KEYS").arg("*").query_async(&mut client).await + .map_err(|e| VeloxError::Query(format!("Redis KEYS failed: {}", e)))?; + Ok(keys.into_iter().map(|k| TableInfo { + schema: "0".to_string(), + name: k.clone(), + preview_query: format!("GET {}", k), + }).collect()) + } + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + _table_schema: &str, + table_name: &str, + ) -> Result, VeloxError> { + let mut client = get_or_create_redis_client(app, state, connection_id).await?; + + let key_type: String = redis::cmd("TYPE").arg(table_name).query_async(&mut client) + .await.map_err(|e| VeloxError::Query(format!("Redis TYPE failed: {}", e)))?; + + match key_type.as_str() { + "string" => Ok(vec![ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.to_string(), + column_name: "value".to_string(), data_type: "string".to_string(), is_nullable: true, + }]), + "hash" => { + let fields: Vec = redis::cmd("HKEYS").arg(table_name).query_async(&mut client) + .await.map_err(|e| VeloxError::Query(format!("Redis HKEYS failed: {}", e)))?; + Ok(fields.into_iter().map(|f| ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.to_string(), + column_name: f, data_type: "string".to_string(), is_nullable: true, + }).collect()) + } + "list" => Ok(vec![ + ColumnInfo { table_schema: "0".to_string(), table_name: table_name.to_string(), column_name: "index".to_string(), data_type: "integer".to_string(), is_nullable: false }, + ColumnInfo { table_schema: "0".to_string(), table_name: table_name.to_string(), column_name: "value".to_string(), data_type: "string".to_string(), is_nullable: true }, + ]), + "set" => Ok(vec![ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.to_string(), + column_name: "member".to_string(), data_type: "string".to_string(), is_nullable: false, + }]), + _ => Ok(vec![ColumnInfo { + table_schema: "0".to_string(), table_name: table_name.to_string(), + column_name: "value".to_string(), data_type: key_type, is_nullable: true, + }]), + } + } + + async fn list_databases( + &self, + _app: &AppHandle, + _state: &AppState, + _connection_id: &str, + ) -> Result, VeloxError> { + Ok(vec![DatabaseInfo { name: "0".to_string() }]) + } +} diff --git a/src-tauri/src/engines/sqlite.rs b/src-tauri/src/engines/sqlite.rs new file mode 100644 index 0000000..c4440ef --- /dev/null +++ b/src-tauri/src/engines/sqlite.rs @@ -0,0 +1,129 @@ +use tauri::AppHandle; + +use crate::db::{ + build_sqlite_pool, get_or_create_sqlite_pool, quote_identifier, + require_safe_identifier, AppState, +}; +use crate::error::VeloxError; +use crate::models::{ColumnInfo, ConnectionInput, DatabaseInfo, QueryResult, TableInfo}; + +use super::DatabaseEngineOps; + +use crate::commands::{sqlite_get_idx, sqlite_get_name, run_query_mysql_or_sqlite}; + +pub struct SqliteEngine; + +impl DatabaseEngineOps for SqliteEngine { + async fn connect( + &self, + _app: &AppHandle, + state: &AppState, + input: &ConnectionInput, + connection_id: &str, + ) -> Result<(), VeloxError> { + let pool = build_sqlite_pool(input).await?; + sqlx::query("select 1").execute(&pool).await + .map_err(|e| VeloxError::Connection(e.to_string()))?; + state.sqlite_pools.write().await.insert(connection_id.to_string(), pool); + Ok(()) + } + + async fn ping( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result<(), VeloxError> { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + sqlx::query("select 1").execute(&pool).await + .map_err(|e| VeloxError::Connection(e.to_string()))?; + Ok(()) + } + + async fn run_query( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + sql: &str, + max_rows: usize, + ) -> Result { + run_query_mysql_or_sqlite( + app, + state, + connection_id, + sql, + max_rows, + crate::models::DatabaseEngine::Sqlite, + ) + .await + .map_err(VeloxError::from) + } + + async fn get_tables( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + ) -> Result, VeloxError> { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + let rows = sqlx::query( + "select name from sqlite_master \ + where type = 'table' and name not like 'sqlite_%' order by name", + ).fetch_all(&pool).await + .map_err(|e| VeloxError::Query(e.to_string()))?; + let mut tables = Vec::new(); + for row in rows { + let name: String = sqlite_get_idx(&row, 0, "name", "get_tables") + .map_err(VeloxError::from)?; + require_safe_identifier(&name, "table name")?; + tables.push(TableInfo { + schema: "main".to_string(), + preview_query: format!("select * from \"{}\" limit 100;", quote_identifier(&name)), + name, + }); + } + Ok(tables) + } + + async fn get_schema( + &self, + app: &AppHandle, + state: &AppState, + connection_id: &str, + _table_schema: &str, + table_name: &str, + ) -> Result, VeloxError> { + let pool = get_or_create_sqlite_pool(app, state, connection_id).await?; + require_safe_identifier(table_name, "table name")?; + let pragma_sql = format!("PRAGMA table_info(\"{}\");", quote_identifier(table_name)); + let rows = sqlx::query(&pragma_sql).fetch_all(&pool).await + .map_err(|e| VeloxError::Query(e.to_string()))?; + let mut columns = Vec::new(); + for row in rows { + let col_name: String = sqlite_get_name(&row, "name", "get_schema") + .map_err(VeloxError::from)?; + let col_type: String = sqlite_get_name(&row, "type", "get_schema") + .map_err(VeloxError::from)?; + let notnull: i64 = sqlite_get_name(&row, "notnull", "get_schema") + .map_err(VeloxError::from)?; + columns.push(ColumnInfo { + table_schema: "main".to_string(), + table_name: table_name.to_string(), + column_name: col_name, + data_type: col_type, + is_nullable: notnull == 0, + }); + } + Ok(columns) + } + + async fn list_databases( + &self, + _app: &AppHandle, + _state: &AppState, + _connection_id: &str, + ) -> Result, VeloxError> { + Ok(vec![DatabaseInfo { name: "main".to_string() }]) + } +} diff --git a/src-tauri/src/error.rs b/src-tauri/src/error.rs new file mode 100644 index 0000000..95abb7c --- /dev/null +++ b/src-tauri/src/error.rs @@ -0,0 +1,80 @@ +/// Centralised error type for VeloxDB. +/// +/// Replaces ad-hoc `Result` in the connection pool, retry, and +/// identifier-guard modules. Tauri command signatures are unchanged (still +/// `Result`); the `From for String` impl bridges the two +/// worlds so that `?` works in both directions. +#[derive(Debug, Clone, PartialEq, thiserror::Error)] +pub enum VeloxError { + /// Network or transport-level failures (timeout, refused, broken pipe, …). + #[error("{0}")] + Connection(String), + + /// A database returned an error for a query. + #[error("{0}")] + Query(String), + + /// A PostgreSQL error that has been formatted by `pg_error` (severity, + /// SQLSTATE, DETAIL, HINT, caret marker). + #[error("{0}")] + Postgres(PgError), + + /// TLS / SSH / pool configuration is invalid. + #[error("{0}")] + Configuration(String), + + /// User-provided input failed validation (identifier safety, SQL checks, …). + #[error("{0}")] + Validation(String), + + /// OS keychain access failed or a credential was missing. + #[error("{0}")] + Credential(String), + + /// Catch-all for unexpected internal states. + #[error("{0}")] + Internal(String), +} + +/// Formatted PostgreSQL error string (produced by `pg_error::map_pg_err`). +#[derive(Debug, Clone, PartialEq, thiserror::Error)] +#[error("{0}")] +pub struct PgError(pub String); + +impl PgError { + pub fn from_error(error: tokio_postgres::Error, sql: Option<&str>) -> Self { + PgError(crate::pg_error::map_pg_err(error, sql)) + } +} + +// ── Two-way bridge with `String` ──────────────────────────────── +// +// Existing code (especially Tauri commands) returns `Result<_, String>`. +// New internal functions return `Result<_, VeloxError>`. +// These impls let `?` work seamlessly in both directions. + +impl From for String { + fn from(e: VeloxError) -> Self { + e.to_string() + } +} + +impl From for VeloxError { + fn from(s: String) -> Self { + VeloxError::Internal(s) + } +} + +// ── From impls for common library error types ─────────────────── + +impl From for VeloxError { + fn from(e: std::io::Error) -> Self { + VeloxError::Internal(e.to_string()) + } +} + +impl From for VeloxError { + fn from(e: serde_json::Error) -> Self { + VeloxError::Internal(e.to_string()) + } +} diff --git a/src-tauri/src/export.rs b/src-tauri/src/export.rs index f1fa29f..5e4b568 100644 --- a/src-tauri/src/export.rs +++ b/src-tauri/src/export.rs @@ -9,7 +9,7 @@ use sqlx::sqlite::SqliteRow; use tokio_postgres::SimpleQueryMessage; use crate::db::{ - get_or_create_mysql_pool, get_or_create_sqlite_pool, resolve_connection_engine, + get_or_create_mysql_pool, get_or_create_sqlite_pool, load_connection, resolve_connection_engine, with_pool_client_retry, AppState, }; use crate::models::DatabaseEngine; @@ -374,6 +374,43 @@ pub async fn export_results_csv( lines } } + DatabaseEngine::Mongo => { + let stored = load_connection(app, &connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let (database, collection) = + crate::commands::mongo::resolve_mongo_export_target(&sql, &stored.database)?; + return crate::commands::mongo::mongo_export_csv( + app, + state, + &connection_id, + &database, + &collection, + &input.output_path, + ) + .await; + } + DatabaseEngine::Duckdb => { + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(&connection_id).ok_or("DuckDB connection not found")?; + let conn = conn_mutex.lock().await; + let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?; + let cols: Vec = stmt.column_names().iter().map(|c| c.to_string()).collect(); + let mut wtr = csv::Writer::from_path(&input.output_path).map_err(|e| e.to_string())?; + wtr.write_record(&cols).map_err(|e| e.to_string())?; + let rows = stmt.query_map([], |row| { + let vals: Vec = cols.iter().enumerate().map(|(i, _)| { + row.get::<_, Option>(i).ok().flatten().unwrap_or_default() + }).collect(); + Ok(vals) + }).map_err(|e| e.to_string())?; + for row in rows { + let vals = row.map_err(|e| e.to_string())?; + wtr.write_record(&vals).map_err(|e| e.to_string())?; + } + wtr.flush().map_err(|e| e.to_string())?; + return Ok(()); + } + DatabaseEngine::Redis => return Err("Not supported for Redis.".to_string()), }; let content = lines.join("\n") + "\n"; @@ -480,6 +517,44 @@ pub async fn export_results_json( } result } + DatabaseEngine::Mongo => { + let stored = load_connection(app, &connection_id)? + .ok_or_else(|| "Stored connection details were not found.".to_string())?; + let (database, collection) = + crate::commands::mongo::resolve_mongo_export_target(&sql, &stored.database)?; + return crate::commands::mongo::mongo_export_json( + app, + state, + &connection_id, + &database, + &collection, + &input.output_path, + ) + .await; + } + DatabaseEngine::Duckdb => { + let conns = state.duckdb_connections.read().await; + let conn_mutex = conns.get(&connection_id).ok_or("DuckDB connection not found")?; + let conn = conn_mutex.lock().await; + let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?; + let cols: Vec = stmt.column_names().iter().map(|c| c.to_string()).collect(); + let mut rows = Vec::new(); + let results = stmt.query_map([], |row| { + let mut map = serde_json::Map::new(); + for (i, col) in cols.iter().enumerate() { + let val: Option = row.get(i).ok().flatten(); + map.insert(col.clone(), serde_json::Value::String(val.unwrap_or_default())); + } + Ok(serde_json::Value::Object(map)) + }).map_err(|e| e.to_string())?; + for r in results { + rows.push(r.map_err(|e| e.to_string())?); + } + let content = serde_json::to_string_pretty(&rows).map_err(|e| e.to_string())?; + fs::write(&input.output_path, content).map_err(|e| e.to_string())?; + return Ok(()); + } + DatabaseEngine::Redis => return Err("Not supported for Redis.".to_string()), }; let content = format!("[\n{}\n]\n", rows.join(",\n")); diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e5f991e..1b0531e 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1,18 +1,23 @@ mod commands; mod credentials; -mod db; +pub mod db; +pub mod engines; +pub mod error; mod export; -mod models; +pub mod models; mod pg_error; -mod sql_split; +pub mod sql_split; mod ssh_tunnel; use commands::{ - apply_table_properties, cancel_veloxy_request, chat_with_db, clear_veloxy_conversation, connect_db, delete_connection, - delete_openrouter_api_key, disconnect_db, execute_ddl_statement, execute_ddl_transaction, export_diagram_png, + apply_table_properties, cancel_veloxy_request, chat_with_db, check_for_updates, clear_veloxy_conversation, + connect_db, delete_connection, + delete_openrouter_api_key, disconnect_db, duckdb_get_schema, duckdb_get_tables, duckdb_run_query, + execute_ddl_statement, execute_ddl_transaction, export_diagram_png, export_results_csv_command, export_results_json_command, generate_sql_from_nl, get_foreign_keys, get_openrouter_api_key, get_query_editor_metadata, get_schema, get_table_indexes, get_table_properties, get_tables, - lint_sql, list_connections_command, list_databases, load_veloxy_conversation, ping_connection, + lint_sql, list_connections_command, list_databases, load_veloxy_conversation, mongo_get_collections, + mongo_get_schema, mongo_run_query, ping_connection, redis_get_keys, redis_get_schema, redis_run_query, refresh_connection, rename_connection, run_query, save_base64_png, save_text_file, set_active_connection, store_openrouter_api_key, switch_database, }; @@ -48,6 +53,7 @@ pub fn run() { } }) .invoke_handler(tauri::generate_handler![ + check_for_updates, connect_db, disconnect_db, rename_connection, @@ -81,7 +87,16 @@ pub fn run() { clear_veloxy_conversation, store_openrouter_api_key, get_openrouter_api_key, - delete_openrouter_api_key + delete_openrouter_api_key, + mongo_run_query, + mongo_get_collections, + mongo_get_schema, + duckdb_run_query, + duckdb_get_tables, + duckdb_get_schema, + redis_run_query, + redis_get_keys, + redis_get_schema ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/src-tauri/src/models.rs b/src-tauri/src/models.rs index 3f3c876..c373a90 100644 --- a/src-tauri/src/models.rs +++ b/src-tauri/src/models.rs @@ -19,6 +19,8 @@ fn default_connection_ssl_mode() -> ConnectionSslMode { ConnectionSslMode::Prefer } +fn is_false(v: &bool) -> bool { !*v } + #[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum DatabaseEngine { @@ -26,6 +28,9 @@ pub enum DatabaseEngine { Postgres, Mysql, Sqlite, + Mongo, + Duckdb, + Redis, } fn default_database_engine() -> DatabaseEngine { @@ -69,7 +74,7 @@ impl SshConfig { && !self.host.is_empty() && !self.user.is_empty() && (self.auth_method != SshAuthMethod::Password - || self.password.as_ref().map_or(false, |p| !p.is_empty())) + || self.password.as_ref().is_some_and(|p| !p.is_empty())) } } @@ -89,6 +94,8 @@ pub struct ConnectionInput { pub password: String, #[serde(default = "default_connection_ssl_mode")] pub ssl_mode: ConnectionSslMode, + #[serde(default, skip_serializing_if = "is_false")] + pub srv_enabled: bool, #[serde(default)] pub ssh_config: Option, #[serde(default)] @@ -113,6 +120,8 @@ pub struct StoredConnection { pub connected_at: String, #[serde(default = "default_connection_ssl_mode")] pub ssl_mode: ConnectionSslMode, + #[serde(default, skip_serializing_if = "is_false")] + pub srv_enabled: bool, #[serde(default, skip_serializing_if = "Option::is_none")] pub ssh_config: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -491,6 +500,7 @@ impl StoredConnection { password: None, connected_at: timestamp_string(), ssl_mode: input.ssl_mode, + srv_enabled: input.srv_enabled, ssh_config: input.ssh_config, extra_params: input.extra_params, } @@ -526,6 +536,7 @@ impl StoredConnection { user: self.user.clone(), password: self.password.clone().unwrap_or_default(), ssl_mode: self.ssl_mode, + srv_enabled: self.srv_enabled, ssh_config: self.ssh_config.clone(), extra_params: self.extra_params.clone(), } diff --git a/src-tauri/tests/integration_tests.rs b/src-tauri/tests/integration_tests.rs new file mode 100644 index 0000000..80e9ea8 --- /dev/null +++ b/src-tauri/tests/integration_tests.rs @@ -0,0 +1,1006 @@ +//! Baseline regression tests for VeloxDB refactor. +//! +//! These tests exercise the public library API that backs every Tauri command. +//! They guard against accidental behavior changes during the refactor. +//! +//! Engines covered: +//! SQLite – full round-trip (in-memory): connect → DDL → DML → query → schema → disconnect +//! DuckDB – full round-trip (in-memory): connect → DDL → DML → query → schema → disconnect +//! PostgreSQL – connection builder, error formatting, identifier validation +//! MySQL – connection/URL builder, port defaults +//! MongoDB – connection string builder, port defaults +//! Redis – URL builder, port defaults + +use std::collections::HashMap; + +use sqlx::Column; +use sqlx::Row; +use veloxdb_lib::db::{ + build_duckdb_connection, build_mongo_connection_string, build_redis_url, + build_sqlite_pool, is_retryable_connection_error, is_safe_identifier, mysql_url, + quote_identifier, require_safe_identifier, +}; +use veloxdb_lib::models::{ + ConnectionInput, ConnectionSslMode, DatabaseEngine, StoredConnection, +}; +use veloxdb_lib::sql_split::split_sql_statements; + +// ── Helper builders ──────────────────────────────────────────── + +fn sqlite_input_memory() -> ConnectionInput { + ConnectionInput { + id: Some("test-sqlite".into()), + name: "test-sqlite".into(), + engine: DatabaseEngine::Sqlite, + host: String::new(), + port: 0, + database: String::new(), + file_path: Some(":memory:".into()), + user: String::new(), + password: String::new(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ssh_config: None, + extra_params: None, + } +} + +fn duckdb_input_memory() -> ConnectionInput { + ConnectionInput { + id: Some("test-duckdb".into()), + name: "test-duckdb".into(), + engine: DatabaseEngine::Duckdb, + host: String::new(), + port: 0, + database: String::new(), + file_path: Some(":memory:".into()), + user: String::new(), + password: String::new(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ssh_config: None, + extra_params: None, + } +} + +fn mysql_input_default() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test-mysql".into(), + engine: DatabaseEngine::Mysql, + host: "db.example.com".into(), + port: 3306, + database: "app".into(), + file_path: None, + user: "root".into(), + password: "secret".into(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Prefer, + ssh_config: None, + extra_params: None, + } +} + +fn postgres_input_default() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test-pg".into(), + engine: DatabaseEngine::Postgres, + host: "pg.example.com".into(), + port: 5432, + database: "app".into(), + file_path: None, + user: "postgres".into(), + password: "secret".into(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Prefer, + ssh_config: None, + extra_params: None, + } +} + +fn mongo_input_default() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test-mongo".into(), + engine: DatabaseEngine::Mongo, + host: "mongo.example.com".into(), + port: 27017, + database: "admin".into(), + file_path: None, + user: "admin".into(), + password: "secret".into(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ssh_config: None, + extra_params: None, + } +} + +fn redis_input_default() -> ConnectionInput { + ConnectionInput { + id: None, + name: "test-redis".into(), + engine: DatabaseEngine::Redis, + host: "redis.example.com".into(), + port: 6379, + database: "0".into(), + file_path: None, + user: String::new(), + password: String::new(), + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ssh_config: None, + extra_params: None, + } +} + +// ── SQLite full round-trip ───────────────────────────────────── + +#[tokio::test] +async fn sqlite_connect_query_schema_disconnect() { + let input = sqlite_input_memory(); + let pool = build_sqlite_pool(&input) + .await + .expect("SQLite pool should build"); + + // DDL + let mut conn = pool.acquire().await.expect("acquire connection"); + sqlx::query("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL, age INTEGER)") + .execute(&mut *conn) + .await + .expect("create table"); + + // DML (insert) + let result = sqlx::query("INSERT INTO users (name, age) VALUES ('Alice', 30)") + .execute(&mut *conn) + .await + .expect("insert"); + assert_eq!(result.rows_affected(), 1); + + // Insert another + sqlx::query("INSERT INTO users (name, age) VALUES ('Bob', 25)") + .execute(&mut *conn) + .await + .expect("insert bob"); + + // SELECT + let rows = sqlx::query("SELECT id, name, age FROM users ORDER BY id") + .fetch_all(&mut *conn) + .await + .expect("select"); + assert_eq!(rows.len(), 2); + + // Verify column names and values + let col_names: Vec = rows[0] + .columns() + .iter() + .map(|c| c.name().to_string()) + .collect(); + assert_eq!(col_names, vec!["id", "name", "age"]); + + let name_a: String = rows[0].try_get("name").expect("get name"); + let age_a: i64 = rows[0].try_get("age").expect("get age"); + assert_eq!(name_a, "Alice"); + assert_eq!(age_a, 30); + + let name_b: String = rows[1].try_get("name").expect("get name"); + let age_b: i64 = rows[1].try_get("age").expect("get age"); + assert_eq!(name_b, "Bob"); + assert_eq!(age_b, 25); + + // Schema inspection — list tables + let tables = sqlx::query( + "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name", + ) + .fetch_all(&mut *conn) + .await + .expect("list tables"); + let table_names: Vec = tables + .iter() + .map(|r| r.try_get::("name").unwrap()) + .collect(); + assert!(table_names.contains(&"users".to_string())); + + // Schema inspection — PRAGMA table_info + let columns = sqlx::query("PRAGMA table_info('users')") + .fetch_all(&mut *conn) + .await + .expect("pragma"); + assert_eq!(columns.len(), 3); + let col0: String = columns[0].try_get("name").expect("col0 name"); + assert_eq!(col0, "id"); + + drop(conn); +} + +#[tokio::test] +async fn sqlite_empty_result_set() { + let input = sqlite_input_memory(); + let pool = build_sqlite_pool(&input).await.expect("pool"); + let mut conn = pool.acquire().await.expect("acquire"); + + sqlx::query("CREATE TABLE t (x INTEGER)") + .execute(&mut *conn) + .await + .expect("create"); + + let rows = sqlx::query("SELECT * FROM t WHERE x > 999") + .fetch_all(&mut *conn) + .await + .expect("select"); + assert!(rows.is_empty()); + + drop(conn); +} + +#[tokio::test] +async fn sqlite_nullable_columns() { + let input = sqlite_input_memory(); + let pool = build_sqlite_pool(&input).await.expect("pool"); + let mut conn = pool.acquire().await.expect("acquire"); + + sqlx::query("CREATE TABLE t (id INTEGER, maybe TEXT)") + .execute(&mut *conn) + .await + .expect("create"); + + sqlx::query("INSERT INTO t (id) VALUES (1)") + .execute(&mut *conn) + .await + .expect("insert"); + + let rows = sqlx::query("SELECT * FROM t") + .fetch_all(&mut *conn) + .await + .expect("select"); + + let id: i64 = rows[0].try_get("id").expect("id"); + let maybe: Option = rows[0].try_get("maybe").expect("maybe"); + assert_eq!(id, 1); + assert!(maybe.is_none()); + + drop(conn); +} + +#[tokio::test] +async fn sqlite_multiple_statements_within_tx() { + let input = sqlite_input_memory(); + let pool = build_sqlite_pool(&input).await.expect("pool"); + let mut conn = pool.acquire().await.expect("acquire"); + + sqlx::query( + "CREATE TABLE a (x INTEGER); CREATE TABLE b (y TEXT);", + ) + .execute(&mut *conn) + .await + .expect("multi-ddl"); + + let tables = sqlx::query( + "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name", + ) + .fetch_all(&mut *conn) + .await + .expect("list"); + let names: Vec = tables.iter().map(|r| r.try_get("name").unwrap()).collect(); + assert!(names.contains(&"a".to_string())); + assert!(names.contains(&"b".to_string())); + + drop(conn); +} + +// ── DuckDB full round-trip ───────────────────────────────────── + +#[test] +fn duckdb_connect_query_schema() { + let input = duckdb_input_memory(); + let conn = build_duckdb_connection(&input).expect("DuckDB connect"); + + // DDL — DuckDB does not auto-increment INTEGER PRIMARY KEY; use explicit IDs + conn.execute_batch( + "CREATE TABLE products (id INTEGER PRIMARY KEY, name VARCHAR, price DOUBLE)", + ) + .expect("create table"); + + // DML — provide explicit IDs + let inserted = conn + .execute( + "INSERT INTO products (id, name, price) VALUES (1, 'Widget', 9.99), (2, 'Gadget', 19.50)", + [], + ) + .expect("insert"); + assert_eq!(inserted, 2); + + // SELECT + let mut stmt = conn + .prepare("SELECT id, name, price FROM products ORDER BY id") + .expect("prepare"); + let rows: Vec<(i32, String, f64)> = stmt + .query_map([], |row| { + Ok((row.get(0)?, row.get(1)?, row.get(2)?)) + }) + .expect("query") + .filter_map(|r| r.ok()) + .collect(); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].1, "Widget"); + assert_eq!(rows[0].2, 9.99); + assert_eq!(rows[1].1, "Gadget"); + assert_eq!(rows[1].2, 19.50); + + // Schema — list tables + let mut stmt = conn + .prepare("SELECT table_name FROM information_schema.tables WHERE table_schema='main' ORDER BY table_name") + .expect("list tables"); + let table_names: Vec = stmt + .query_map([], |row| row.get(0)) + .expect("query") + .filter_map(|r| r.ok()) + .collect(); + assert!(table_names.contains(&"products".to_string())); + + // Schema — describe columns + let mut stmt = conn + .prepare("SELECT column_name, data_type FROM information_schema.columns WHERE table_name='products' ORDER BY ordinal_position") + .expect("list columns"); + let cols: Vec<(String, String)> = stmt + .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) + .expect("query") + .filter_map(|r| r.ok()) + .collect(); + assert_eq!(cols.len(), 3); + assert_eq!(cols[0].0, "id"); + assert_eq!(cols[1].0, "name"); + assert_eq!(cols[2].0, "price"); +} + +#[test] +fn duckdb_empty_table() { + let input = duckdb_input_memory(); + let conn = build_duckdb_connection(&input).expect("connect"); + + conn.execute_batch("CREATE TABLE e (x INTEGER)") + .expect("create"); + + let mut stmt = conn.prepare("SELECT * FROM e").expect("prepare"); + let rows: Vec = stmt + .query_map([], |row| row.get(0)) + .expect("query") + .filter_map(|r| r.ok()) + .collect(); + assert!(rows.is_empty()); +} + +#[test] +fn duckdb_type_preservation() { + let input = duckdb_input_memory(); + let conn = build_duckdb_connection(&input).expect("connect"); + + conn.execute_batch( + "CREATE TABLE types (i BIGINT, f DOUBLE, s VARCHAR)", + ) + .expect("create"); + + conn.execute( + "INSERT INTO types VALUES (42, 3.14, 'hello')", + [], + ) + .expect("insert"); + + // DuckDB uses its own timestamp type; verify core types only + let mut stmt = conn.prepare("SELECT i, f, s FROM types").expect("prepare"); + let row: (i64, f64, String) = stmt + .query_row([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?))) + .expect("query"); + + assert_eq!(row.0, 42); + #[allow(clippy::approx_constant)] + { + assert!((row.1 - 3.14).abs() < 0.001); + } + assert_eq!(row.2, "hello"); + + // Verify boolean handling works (0/1 integer equivalent) + conn.execute_batch("CREATE TABLE flags (active INTEGER)") + .expect("create flags"); + conn.execute("INSERT INTO flags VALUES (1)", []) + .expect("insert flag"); + + let mut stmt = conn.prepare("SELECT active FROM flags").expect("prepare flags"); + let active: i32 = stmt + .query_row([], |row| row.get(0)) + .expect("query flag"); + assert_eq!(active, 1); +} + +// ── PostgreSQL — connection builder / error formatting ──────── + +/// Pool creation is lazy — the pool object is created without connecting. +/// This test verifies the config is accepted and a pool is returned. +#[test] +fn postgres_build_pool_creates_lazy_pool() { + let input = ConnectionInput { + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ..postgres_input_default() + }; + // Pool creation should succeed (lazy — no connection attempted yet) + let pool = veloxdb_lib::db::build_pool(&input) + .expect("Pool creation with Disable SSL should succeed"); + // Verify we can close it without errors + pool.close(); +} + +#[test] +fn postgres_build_pool_prefer_ssl() { + let input = ConnectionInput { + srv_enabled: false, + ssl_mode: ConnectionSslMode::Prefer, + ..postgres_input_default() + }; + let pool = veloxdb_lib::db::build_pool(&input) + .expect("Pool creation with Prefer SSL should succeed"); + pool.close(); +} + +#[test] +fn postgres_build_pool_require_ssl() { + let input = ConnectionInput { + srv_enabled: false, + ssl_mode: ConnectionSslMode::Require, + ..postgres_input_default() + }; + let pool = veloxdb_lib::db::build_pool(&input) + .expect("Pool creation with Require SSL should succeed"); + pool.close(); +} + +#[test] +fn postgres_tls_connector_is_cached() { + // Verify the OnceLock TLS connector works repeatedly + let input = ConnectionInput { + srv_enabled: false, + ssl_mode: ConnectionSslMode::Prefer, + ..postgres_input_default() + }; + + let pool1 = veloxdb_lib::db::build_pool(&input).expect("first pool"); + let pool2 = veloxdb_lib::db::build_pool(&input).expect("second pool"); + + // Both should produce valid pool objects (TLS connector is cached via OnceLock) + pool1.close(); + pool2.close(); +} + +// ── MySQL — URL builder / port defaults ─────────────────────── + +#[test] +fn mysql_url_includes_ssl_mode() { + let input = mysql_input_default(); + let url = mysql_url(&input.host, input.port, &input); + assert!(url.starts_with("mysql://")); + assert!(url.contains("@db.example.com:3306/app")); + assert!(url.contains("ssl-mode=PREFERRED")); +} + +#[test] +fn mysql_url_disable_ssl() { + let input = ConnectionInput { + srv_enabled: false, + ssl_mode: ConnectionSslMode::Disable, + ..mysql_input_default() + }; + let url = mysql_url(&input.host, input.port, &input); + assert!(url.ends_with("ssl-mode=DISABLED")); +} + +#[test] +fn mysql_url_require_ssl() { + let input = ConnectionInput { + srv_enabled: false, + ssl_mode: ConnectionSslMode::Require, + ..mysql_input_default() + }; + let url = mysql_url(&input.host, input.port, &input); + assert!(url.ends_with("ssl-mode=REQUIRED")); +} + +#[test] +fn mysql_url_encodes_special_chars() { + let input = ConnectionInput { + host: "db.example.com".into(), + port: 3306, + database: "my-app".into(), + user: "user@domain".into(), + password: "p@ss:word".into(), + ..mysql_input_default() + }; + let url = mysql_url(&input.host, input.port, &input); + // Should be URL-encoded, not raw special chars + assert!(url.contains("%40")); // @ encoded + assert!(!url.contains("p@ss")); +} + +// ── MongoDB — connection string builder ──────────────────────── + +#[test] +fn mongo_connection_string_no_auth() { + let input = ConnectionInput { + user: String::new(), + password: String::new(), + ..mongo_input_default() + }; + let uri = build_mongo_connection_string(&input); + assert_eq!(uri, "mongodb://mongo.example.com:27017/admin"); +} + +#[test] +fn mongo_connection_string_with_auth() { + let input = mongo_input_default(); + let uri = build_mongo_connection_string(&input); + assert!(uri.starts_with("mongodb://")); + assert!(uri.contains("admin:secret@mongo.example.com:27017/admin")); +} + +#[test] +fn mongo_connection_string_defaults() { + let input = ConnectionInput { + host: String::new(), + port: 0, + database: String::new(), + user: String::new(), + password: String::new(), + ..mongo_input_default() + }; + let uri = build_mongo_connection_string(&input); + assert_eq!(uri, "mongodb://localhost:27017/admin"); +} + +#[test] +fn mongo_connection_string_with_extra_params() { + let mut params = HashMap::new(); + params.insert("replicaSet".to_string(), "rs0".to_string()); + params.insert("authSource".to_string(), "admin".to_string()); + let input = ConnectionInput { + extra_params: Some(params), + ..mongo_input_default() + }; + let uri = build_mongo_connection_string(&input); + assert!(uri.contains("?")); + assert!(uri.contains("replicaSet=rs0")); + assert!(uri.contains("authSource=admin")); +} + +#[test] +fn mongo_connection_string_custom_port() { + let input = ConnectionInput { + port: 27018, + ..mongo_input_default() + }; + let uri = build_mongo_connection_string(&input); + assert!(uri.contains(":27018/")); +} + +// ── Redis — URL builder ──────────────────────────────────────── + +#[test] +fn redis_url_no_auth() { + let input = redis_input_default(); + let url = build_redis_url(&input); + assert_eq!(url, "redis://redis.example.com:6379"); +} + +#[test] +fn redis_url_with_auth() { + let input = ConnectionInput { + user: "default".into(), + password: "secret".into(), + ..redis_input_default() + }; + let url = build_redis_url(&input); + assert_eq!(url, "redis://default:secret@redis.example.com:6379"); +} + +#[test] +fn redis_url_defaults_to_localhost() { + let input = ConnectionInput { + host: String::new(), + port: 0, + ..redis_input_default() + }; + let url = build_redis_url(&input); + assert_eq!(url, "redis://127.0.0.1:6379"); +} + +// ── Identifier safety ────────────────────────────────────────── + +#[test] +fn accepts_valid_identifiers() { + assert!(is_safe_identifier("users")); + assert!(is_safe_identifier("_internal")); + assert!(is_safe_identifier("Table123")); + assert!(is_safe_identifier("a")); + assert!(is_safe_identifier("snake_case_name")); + assert_eq!(require_safe_identifier("users", "table"), Ok("users")); +} + +#[test] +fn rejects_sql_injection_identifiers() { + assert!(!is_safe_identifier("\"; DROP TABLE users; --")); + assert!(!is_safe_identifier("foo; DELETE FROM bar")); + assert!(!is_safe_identifier("a\0b")); + assert!(!is_safe_identifier("")); + assert!(!is_safe_identifier("foo bar")); + assert!(!is_safe_identifier("table-name")); + assert!(require_safe_identifier("foo);", "table name").is_err()); + assert!(require_safe_identifier("", "column name").is_err()); +} + +#[test] +fn rejects_overly_long_identifiers() { + let long = "a".repeat(129); + assert!(!is_safe_identifier(&long)); +} + +#[test] +fn quote_identifier_escapes_double_quotes() { + assert_eq!(quote_identifier("hello"), "hello"); + assert_eq!(quote_identifier("it\"s"), "it\"\"s"); + assert_eq!(quote_identifier("\"\""), "\"\"\"\""); + assert_eq!(quote_identifier("normal_name"), "normal_name"); +} + +// ── SQL statement splitting ──────────────────────────────────── + +#[test] +fn splits_simple_statements() { + let parts = split_sql_statements("select 1; select 2;"); + assert_eq!(parts, vec!["select 1", "select 2"]); +} + +#[test] +fn splits_trailing_semicolon() { + let parts = split_sql_statements("select 1;"); + assert_eq!(parts, vec!["select 1"]); +} + +#[test] +fn ignores_semicolons_in_single_quoted_strings() { + let parts = split_sql_statements("select 'a;b'; select 'c';"); + assert_eq!(parts, vec!["select 'a;b'", "select 'c'"]); +} + +#[test] +fn ignores_semicolons_in_double_quoted_strings() { + let parts = split_sql_statements(r#"select "a;b"; select "c";"#); + assert_eq!(parts, vec![r#"select "a;b""#, r#"select "c""#]); +} + +#[test] +fn ignores_semicolons_in_block_comments() { + let parts = split_sql_statements("select 1 /* this; is; a; comment */; select 2;"); + assert_eq!(parts, vec!["select 1 /* this; is; a; comment */", "select 2"]); +} + +#[test] +fn ignores_semicolons_in_line_comments() { + let parts = split_sql_statements("select 1 -- ; still comment\n; select 2;"); + // line comment absorbs the first ';' and 'still comment' up to \n, + // then the ';' after \n splits, producing two statements + assert_eq!(parts.len(), 2); + assert!(parts[0].contains("select 1")); + assert!(parts[0].contains("--")); + assert_eq!(parts[1], "select 2"); +} + +#[test] +fn handles_escaped_quotes() { + let parts = split_sql_statements("select 'it''s'; select 2;"); + assert_eq!(parts, vec!["select 'it''s'", "select 2"]); +} + +#[test] +fn handles_escaped_double_quotes() { + let parts = split_sql_statements(r#"select "it""s"; select 2;"#); + assert_eq!(parts, vec![r#"select "it""s""#, "select 2"]); +} + +#[test] +fn handles_dollar_quoting() { + let parts = split_sql_statements("select $$hello; world$$; select 2;"); + assert_eq!(parts, vec!["select $$hello; world$$", "select 2"]); +} + +#[test] +fn handles_named_dollar_quoting() { + let parts = split_sql_statements("select $tag$hello; world$tag$; select 2;"); + assert_eq!(parts, vec!["select $tag$hello; world$tag$", "select 2"]); +} + +#[test] +fn handles_nested_block_comments() { + let parts = split_sql_statements( + "select 1 /* outer /* inner */ still comment */; select 2;", + ); + assert_eq!(parts.len(), 2); + assert!(parts[0].contains("select 1")); + assert!(parts[1].contains("select 2")); +} + +#[test] +fn handles_empty_input() { + let parts = split_sql_statements(""); + assert!(parts.is_empty()); +} + +#[test] +fn handles_whitespace_only() { + let parts = split_sql_statements(" \n \t "); + assert!(parts.is_empty()); +} + +#[test] +fn single_statement_no_semicolon() { + let parts = split_sql_statements("select 1"); + assert_eq!(parts, vec!["select 1"]); +} + +// ── Model type conversions ───────────────────────────────────── + +#[test] +fn connection_input_to_stored_connection() { + let input = postgres_input_default(); + let stored = StoredConnection::from_input("conn-1".into(), input.clone()); + + assert_eq!(stored.id, "conn-1"); + assert_eq!(stored.name, "test-pg"); + assert_eq!(stored.engine, DatabaseEngine::Postgres); + assert_eq!(stored.host, "pg.example.com"); + assert_eq!(stored.port, 5432); + assert_eq!(stored.database, "app"); + assert_eq!(stored.user, "postgres"); + assert!(stored.password.is_none()); // password stripped in from_input + assert!(stored.connected_at.parse::().is_ok()); +} + +#[test] +fn stored_connection_to_summary() { + let input = postgres_input_default(); + let stored = StoredConnection::from_input("conn-1".into(), input); + let summary = stored.summary(); + + assert_eq!(summary.id, "conn-1"); + assert_eq!(summary.name, "test-pg"); + assert_eq!(summary.engine, DatabaseEngine::Postgres); + assert!(summary.table_property_editing_supported); +} + +#[test] +fn stored_connection_to_input_roundtrip() { + let input = postgres_input_default(); + let stored = StoredConnection::from_input("conn-1".into(), input.clone()); + let roundtripped = stored.to_input(); + + assert_eq!(roundtripped.id, Some("conn-1".into())); + assert_eq!(roundtripped.host, input.host); + assert_eq!(roundtripped.port, input.port); + assert_eq!(roundtripped.engine, input.engine); +} + +#[test] +fn table_property_editing_only_for_postgres() { + let pg = StoredConnection::from_input("pg".into(), postgres_input_default()); + let mysql = StoredConnection::from_input("my".into(), mysql_input_default()); + let sqlite = StoredConnection::from_input("sl".into(), sqlite_input_memory()); + let mongo = StoredConnection::from_input("mg".into(), mongo_input_default()); + + assert!(pg.summary().table_property_editing_supported); + assert!(!mysql.summary().table_property_editing_supported); + assert!(!sqlite.summary().table_property_editing_supported); + assert!(!mongo.summary().table_property_editing_supported); +} + +// ── DatabaseEngine enum variants ──────────────────────────────── + +#[test] +fn database_engine_default_is_postgres() { + let engine = DatabaseEngine::default(); + assert_eq!(engine, DatabaseEngine::Postgres); +} + +#[test] +fn database_engine_all_variants_exist() { + // Verify each variant can be constructed and compared + let engines = vec![ + DatabaseEngine::Postgres, + DatabaseEngine::Mysql, + DatabaseEngine::Sqlite, + DatabaseEngine::Mongo, + DatabaseEngine::Duckdb, + DatabaseEngine::Redis, + ]; + assert_eq!(engines.len(), 6); + for engine in &engines { + assert_eq!(engine, engine); + } +} + +// ── ConnectionSslMode enum ───────────────────────────────────── + +#[test] +fn ssl_mode_default_is_prefer() { + assert_eq!(ConnectionSslMode::default(), ConnectionSslMode::Prefer); +} + +#[test] +fn ssl_mode_all_variants() { + let modes = [ + ConnectionSslMode::Disable, + ConnectionSslMode::Prefer, + ConnectionSslMode::Require, + ]; + assert_eq!(modes.len(), 3); +} + +// ── Retryable error classification ───────────────────────────── + +#[test] +fn classifies_broken_pipe_as_retryable() { + assert!(is_retryable_connection_error("Broken pipe (os error 32)")); + assert!(is_retryable_connection_error("broken PIPE")); +} + +#[test] +fn classifies_connection_reset_as_retryable() { + assert!(is_retryable_connection_error("Connection reset by peer")); + assert!(is_retryable_connection_error("connection RESET")); +} + +#[test] +fn classifies_connection_refused_as_retryable() { + assert!(is_retryable_connection_error("Connection refused")); +} + +#[test] +fn classifies_eof_as_retryable() { + assert!(is_retryable_connection_error("unexpected EOF")); + assert!(is_retryable_connection_error("EOF has been reached")); + assert!(is_retryable_connection_error("unexpected end of file")); +} + +#[test] +fn classifies_server_closed_as_retryable() { + assert!(is_retryable_connection_error("server closed the connection")); + assert!(is_retryable_connection_error("connection closed")); + assert!(is_retryable_connection_error("closed the connection unexpectedly")); +} + +#[test] +fn classifies_communication_error_as_retryable() { + assert!(is_retryable_connection_error( + "error communicating with the server" + )); + assert!(is_retryable_connection_error("could not receive data from server")); + assert!(is_retryable_connection_error("could not send data to server")); +} + +#[test] +fn classifies_timeout_as_retryable() { + assert!(is_retryable_connection_error( + "timeout occurred while waiting" + )); + assert!(is_retryable_connection_error( + "timeout occurred while creating" + )); + assert!(is_retryable_connection_error( + "timeout occurred while recycling" + )); +} + +#[test] +fn does_not_classify_syntax_error_as_retryable() { + assert!(!is_retryable_connection_error("syntax error at or near SELECT")); + assert!(!is_retryable_connection_error("column does not exist")); + assert!(!is_retryable_connection_error("permission denied")); +} + +#[test] +fn retryable_error_is_case_insensitive() { + assert!(is_retryable_connection_error("BROKEN PIPE")); + assert!(is_retryable_connection_error("Connection Refused")); + assert!(is_retryable_connection_error("Unexpected EOF")); +} + +// ── SSH config validation ────────────────────────────────────── + +#[test] +fn ssh_config_inactive_when_disabled() { + let config = veloxdb_lib::models::SshConfig { + enabled: false, + host: "jump.example.com".into(), + port: 22, + user: "deploy".into(), + auth_method: veloxdb_lib::models::SshAuthMethod::KeyFile, + password: None, + private_key_path: Some("~/.ssh/id_rsa".into()), + passphrase: None, + }; + assert!(!config.is_active()); +} + +#[test] +fn ssh_config_inactive_when_host_empty() { + let config = veloxdb_lib::models::SshConfig { + enabled: true, + host: String::new(), + port: 22, + user: "deploy".into(), + auth_method: veloxdb_lib::models::SshAuthMethod::KeyFile, + password: None, + private_key_path: None, + passphrase: None, + }; + assert!(!config.is_active()); +} + +#[test] +fn ssh_config_active_with_key_file() { + let config = veloxdb_lib::models::SshConfig { + enabled: true, + host: "jump.example.com".into(), + port: 22, + user: "deploy".into(), + auth_method: veloxdb_lib::models::SshAuthMethod::KeyFile, + password: None, + private_key_path: None, + passphrase: None, + }; + assert!(config.is_active()); +} + +#[test] +fn ssh_config_active_with_password() { + let config = veloxdb_lib::models::SshConfig { + enabled: true, + host: "jump.example.com".into(), + port: 22, + user: "deploy".into(), + auth_method: veloxdb_lib::models::SshAuthMethod::Password, + password: Some("secret".into()), + private_key_path: None, + passphrase: None, + }; + assert!(config.is_active()); +} + +#[test] +fn ssh_config_inactive_without_password_when_password_method() { + let config = veloxdb_lib::models::SshConfig { + enabled: true, + host: "jump.example.com".into(), + port: 22, + user: "deploy".into(), + auth_method: veloxdb_lib::models::SshAuthMethod::Password, + password: None, + private_key_path: None, + passphrase: None, + }; + assert!(!config.is_active()); +} + +// ── max_query_rows constant ──────────────────────────────────── + +#[test] +fn max_query_rows_is_positive() { + const _: () = assert!(veloxdb_lib::db::MAX_QUERY_ROWS > 0); +} + +// ── Default port constants ───────────────────────────────────── + +#[test] +fn default_port_constants_are_correct() { + assert_eq!(veloxdb_lib::db::DEFAULT_MYSQL_PORT, 3306); + assert_eq!(veloxdb_lib::db::DEFAULT_MONGO_PORT, 27017); + assert_eq!(veloxdb_lib::db::DEFAULT_REDIS_PORT, 6379); +} diff --git a/src/App.tsx b/src/App.tsx index c06f970..03f2b13 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -1,1251 +1,7 @@ -import { GearIcon, SidebarSimpleIcon } from "@phosphor-icons/react"; -import { useQueryClient } from "@tanstack/react-query"; -import { - type CSSProperties, - type PointerEvent as ReactPointerEvent, - useCallback, - useEffect, - useMemo, - useRef, - useState, -} from "react"; -import { useTranslation } from "react-i18next"; - -import { ErrorBoundary } from "@/components/ErrorBoundary"; -import { Button } from "@/components/ui/button"; -import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { queryKeys } from "@/data/query-keys"; -import { veloxDbRepository } from "@/data/repositories"; -import type { - AskVeloxyChatResponse, - AskVeloxyConversationResponse, - ConnectionSummary, - TableInfo, -} from "@/data/types"; -import { CommandPalette } from "@/features/commands/components/CommandPalette"; -import { ShortcutSheet } from "@/features/commands/components/ShortcutSheet"; -import { SettingsDialog } from "@/features/commands/components/SettingsDialog"; -import { ConnectionDialog } from "@/features/connections/components/ConnectionDialog"; -import { RenameConnectionDialog } from "@/features/connections/components/RenameConnectionDialog"; -import { ConnectionsSidebarTree } from "@/features/connections/components/ConnectionsSidebarTree"; -import { - useActivateConnectionMutation, - useConnectionsQuery, - useConnectMutation, - useDeleteConnectionMutation, - useRenameConnectionMutation, -} from "@/features/connections/queries"; -import { ModelWorkspace } from "@/features/model/components/ModelWorkspace"; +import { useState } from "react"; +import { VeloxApp } from "@/components/VeloxApp"; import { readOnboardingCompleted } from "@/features/onboarding/constants"; import { OnboardingFlow } from "@/features/onboarding/OnboardingFlow"; -import { - QueryWorkspace, - type QueryWorkspaceHandle, -} from "@/features/queries/components/QueryWorkspace"; -import { - AskVeloxySidebar, - type AskVeloxySubmitResult, -} from "@/features/queries/components/AskVeloxyDialog"; -import { useSaveResultEditsMutation, useDeleteRowsMutation } from "@/features/queries/queries"; -import { notifyError, notifySuccess } from "@/lib/error-notifier"; -import { loadOpenRouterApiKey } from "@/lib/openrouter-credentials"; -import { useSettings, resolveTheme } from "@/lib/settings"; -import { - buildDropTableSql, - buildDeleteTemplateSql, - buildInsertTemplateSql, - buildRenameTableSql, - buildSelectAllSql, - buildSelectCountSql, - buildUpdateTemplateSql, -} from "@/features/queries/sql-templates"; -import type { TableQuickSqlAction } from "@/features/queries/table-quick-actions"; -import { isInsertFormColumn, type ResultEditPatch } from "@/features/queries/result-edits"; -import { quoteIdent } from "@/lib/sql-ident"; -import { TablePropertiesDialog } from "@/features/schema/components/TablePropertiesDialog"; -import { - useTablePropertiesQuery, - useTableSchemaQuery, -} from "@/features/schema/queries"; -import { useTablesQuery } from "@/features/tables/queries"; - -const SIDEBAR_WIDTH_KEY = "veloxdb.sidebarWidth"; -const SIDEBAR_COLLAPSED_KEY = "veloxdb.sidebarCollapsed"; -const RESULTS_HEIGHT_KEY = "veloxdb.resultsHeight"; -const LAST_ACTIVE_CONNECTION_KEY = "veloxdb.lastActiveConnectionId"; -const DEFAULT_SIDEBAR_WIDTH = 280; -const MIN_SIDEBAR_WIDTH = 220; -const MAX_SIDEBAR_WIDTH = 520; -const DEFAULT_RESULTS_HEIGHT = 260; - -function clampSidebarWidth(value: number) { - return Math.min(MAX_SIDEBAR_WIDTH, Math.max(MIN_SIDEBAR_WIDTH, value)); -} - -function readSidebarWidth() { - const value = Number(window.localStorage.getItem(SIDEBAR_WIDTH_KEY)); - return Number.isFinite(value) - ? clampSidebarWidth(value) - : DEFAULT_SIDEBAR_WIDTH; -} - -function readSidebarCollapsed() { - return window.localStorage.getItem(SIDEBAR_COLLAPSED_KEY) === "true"; -} - -function readResultsHeight() { - const value = Number(window.localStorage.getItem(RESULTS_HEIGHT_KEY)); - return Number.isFinite(value) && value > 0 ? value : DEFAULT_RESULTS_HEIGHT; -} - -function persistLastActiveConnectionId(connectionId: string) { - window.localStorage.setItem(LAST_ACTIVE_CONNECTION_KEY, connectionId); -} - -function connectionSecondaryText(connection: ConnectionSummary): string { - if (connection.engine === "sqlite") { - return connection.filePath === ":memory:" - ? "SQLite in-memory database" - : `SQLite file: ${connection.filePath ?? connection.database}`; - } - return `${connection.user}@${connection.host}:${connection.port}${connection.sshConfig ? " (via SSH)" : ""}`; -} - -function connectionHeadline(connection: ConnectionSummary): string { - if (connection.engine === "sqlite") { - return `Connected to SQLite (${connection.filePath ?? connection.database})`; - } - return `Connected to ${connection.database} on ${connection.host}:${connection.port}`; -} - -function engineLabel(engine: ConnectionSummary["engine"]): string { - if (engine === "postgres") return "PostgreSQL"; - if (engine === "mysql") return "MySQL"; - return "SQLite"; -} - -function VeloxApp() { - const { t } = useTranslation(); - const [connection, setConnection] = useState(null); - const queryWorkspaceRef = useRef(null); - const [focusedQueryCaps, setFocusedQueryCaps] = useState({ - hasLastQuery: false, - hasResult: false, - }); - const [tableSearch, setTableSearch] = useState(""); - const [selectedTable, setSelectedTable] = useState(null); - const themeSetting = useSettings((s) => s.theme) - const isDark = useMemo(() => resolveTheme(themeSetting) === 'dark', [themeSetting]) - const fontSize = useSettings((s) => s.fontSize) - const [settingsOpen, setSettingsOpen] = useState(false) - const [commandPaletteOpen, setCommandPaletteOpen] = useState(false) - const [connectionDialogOpen, setConnectionDialogOpen] = useState(false); - const [renamingConnection, setRenamingConnection] = useState(null); - const [isSidebarCollapsed, setIsSidebarCollapsed] = - useState(readSidebarCollapsed); - const [sidebarWidth, setSidebarWidth] = useState(readSidebarWidth); - const [resultsHeight, setResultsHeight] = useState(readResultsHeight); - const [tablePropertiesDialogOpen, setTablePropertiesDialogOpen] = - useState(false); - const [tablePropertiesTarget, setTablePropertiesTarget] = useState<{ - connectionId: string; - table: TableInfo; - } | null>(null); - const [insertRowTrigger, setInsertRowTrigger] = useState(0); - const [mainWorkspace, setMainWorkspace] = useState<"query" | "model">( - "query", - ); - const [askVeloxyPending, setAskVeloxyPending] = useState(false); - const [askVeloxyError, setAskVeloxyError] = useState(null); - const veloxyOpenRouterApiKey = useSettings((s) => s.veloxyOpenRouterApiKey); - const veloxyModel = useSettings((s) => s.veloxyModel); - const veloxyBaseUrl = useSettings((s) => s.veloxyBaseUrl); - - const queryClient = useQueryClient(); - - const requestInsertRow = useCallback(() => { - setInsertRowTrigger((n) => n + 1); - }, []); - - const handleInsertRowSuccess = useCallback(() => { - void queryClient.invalidateQueries({ - queryKey: queryKeys.tableProperties(connection?.id, selectedTable), - }); - }, [connection?.id, queryClient, selectedTable]); - - useEffect(() => { - document.documentElement.classList.toggle("dark", isDark); - }, [isDark]); - - useEffect(() => { - const sizes = { sm: 12, md: 14, lg: 16 } - document.documentElement.style.fontSize = `${sizes[fontSize]}px` - }, [fontSize]) - - useEffect(() => { - void loadOpenRouterApiKey() - }, []) - - useEffect(() => { - window.localStorage.setItem( - SIDEBAR_COLLAPSED_KEY, - String(isSidebarCollapsed), - ); - }, [isSidebarCollapsed]); - - useEffect(() => { - window.localStorage.setItem(SIDEBAR_WIDTH_KEY, String(sidebarWidth)); - }, [sidebarWidth]); - - useEffect(() => { - window.localStorage.setItem(RESULTS_HEIGHT_KEY, String(resultsHeight)); - }, [resultsHeight]); - - const connectionsQuery = useConnectionsQuery(); - - useEffect(() => { - if (connectionsQuery.isError && connectionsQuery.error) { - notifyError(connectionsQuery.error, { - title: t("connection.failedToLoad"), - }); - } - }, [connectionsQuery.isError, connectionsQuery.error, t]); - - const connectMutation = useConnectMutation({ - onError: (error) => { - notifyError(error, { category: "connection", force: true }); - }, - onSuccess: (nextConnection) => { - notifySuccess( - t("connection.connected", { database: nextConnection.database }), - connectionSecondaryText(nextConnection), - ); - persistLastActiveConnectionId(nextConnection.id); - setConnection(nextConnection); - setSelectedTable(null); - setTableSearch(""); - setIsSidebarCollapsed(false); - setConnectionDialogOpen(false); - setTablePropertiesDialogOpen(false); - setTablePropertiesTarget(null); - queueMicrotask(() => { - queryWorkspaceRef.current?.setActiveTabConnection(nextConnection.id); - }); - }, - }); - - const activateConnectionMutation = useActivateConnectionMutation({ - onError: (error) => { - notifyError(error, { category: "connection", force: true }); - }, - onSuccess: (nextConnection) => { - persistLastActiveConnectionId(nextConnection.id); - setConnection(nextConnection); - setSelectedTable(null); - setTableSearch(""); - setTablePropertiesDialogOpen(false); - setTablePropertiesTarget(null); - queueMicrotask(() => { - queryWorkspaceRef.current?.setActiveTabConnection(nextConnection.id); - }); - }, - }); - - const deleteConnectionMutation = useDeleteConnectionMutation({ - onError: (error) => { - notifyError(error, { category: "connection", force: true }); - }, - onSuccess: (connectionId) => { - queryWorkspaceRef.current?.detachDeletedConnection(connectionId); - if (connection?.id === connectionId) { - setConnection(null); - setSelectedTable(null); - setTableSearch(""); - } - notifySuccess(t("connection.deleted")); - }, - }); - - const renameConnectionMutation = useRenameConnectionMutation({ - onError: (error) => { - notifyError(error, { category: "connection", force: true }); - }, - }); - - const connectionRestoreAttemptedRef = useRef(false); - const autoReconnect = useSettings((s) => s.autoReconnect) - - useEffect(() => { - if (connectionRestoreAttemptedRef.current) return; - if (!autoReconnect) { connectionRestoreAttemptedRef.current = true; return } - const list = connectionsQuery.data; - if (!list?.length) return; - if (connection) { - connectionRestoreAttemptedRef.current = true; - return; - } - - const savedId = window.localStorage.getItem(LAST_ACTIVE_CONNECTION_KEY); - if (!savedId) { - connectionRestoreAttemptedRef.current = true; - return; - } - - const match = list.find((c) => c.id === savedId); - const target = match ?? list[0]; - if (!target) { - connectionRestoreAttemptedRef.current = true; - return; - } - - connectionRestoreAttemptedRef.current = true; - activateConnectionMutation.mutate(target.id); - }, [connectionsQuery.data, connection, activateConnectionMutation, autoReconnect]); - - const tablesQuery = useTablesQuery(connection?.id); - - const schemaQuery = useTableSchemaQuery({ - connectionId: connection?.id, - table: selectedTable, - enabled: Boolean(connection?.id && selectedTable), - }); - const tablePropertiesQuery = useTablePropertiesQuery({ - connectionId: connection?.id, - table: selectedTable, - enabled: Boolean(connection?.id && selectedTable), - }); - const saveResultEditsMutation = useSaveResultEditsMutation({ - onError: (error) => { - notifyError(error, { - category: "query", - title: t("editor.failedToSave"), - }); - }, - }); - const deleteRowsMutation = useDeleteRowsMutation({ - onError: (error) => { - notifyError(error, { - category: "query", - title: t("editor.failedToDelete"), - }); - }, - }); - const connectionsErrorMessage = - connectionsQuery.error instanceof Error - ? connectionsQuery.error.message - : t("connection.failedToLoad"); - - const tablesErrorMessage = - tablesQuery.error instanceof Error - ? tablesQuery.error.message - : t("table.failedToLoad"); - - const schemaErrorMessage = - schemaQuery.error instanceof Error - ? schemaQuery.error.message - : t("table.failedToLoadSchema"); - const tablePropertiesErrorMessage = - tablePropertiesQuery.error instanceof Error - ? tablePropertiesQuery.error.message - : t("table.failedToLoadProperties"); - - const tablesForUi = tablesQuery.data ?? []; - const activeConnectionEngine = connection?.engine ?? "postgres"; - - const handleSelectTable = (table: TableInfo) => { - setSelectedTable(table); - queryWorkspaceRef.current?.applyTablePreview(table.previewQuery); - }; - - const handleTableQuickAction = useCallback( - async ( - action: TableQuickSqlAction, - connectionId: string, - table: TableInfo, - ) => { - if (action === "tableProperties") { - setTablePropertiesTarget({ connectionId, table }); - setTablePropertiesDialogOpen(true); - return; - } - if (action === "addRow") { - setSelectedTable(table); - setInsertRowTrigger((n) => n + 1); - return; - } - - setSelectedTable(table); - try { - switch (action) { - case "selectAll": - queryWorkspaceRef.current?.openTabWithSql( - buildSelectAllSql(table, 200, activeConnectionEngine), - ); - return; - case "selectCount": - queryWorkspaceRef.current?.openTabWithSql( - buildSelectCountSql(table, activeConnectionEngine), - ); - return; - case "insertTemplate": - case "updateTemplate": - case "deleteTemplate": { - const props = await queryClient.fetchQuery({ - queryKey: queryKeys.tableProperties(connectionId, table), - queryFn: () => - veloxDbRepository.getTableProperties(connectionId, table), - }); - - const pk = props - .filter((c) => c.isPrimaryKey) - .map((c) => c.columnName); - const insertCols = props - .filter(isInsertFormColumn) - .map((c) => c.columnName); - - if (action === "insertTemplate") { - queryWorkspaceRef.current?.openTabWithSql( - buildInsertTemplateSql( - table, - insertCols, - activeConnectionEngine, - ), - ); - return; - } - if (action === "updateTemplate") { - queryWorkspaceRef.current?.openTabWithSql( - buildUpdateTemplateSql( - table, - pk, - activeConnectionEngine, - ), - ); - return; - } - queryWorkspaceRef.current?.openTabWithSql( - buildDeleteTemplateSql( - table, - pk, - activeConnectionEngine, - ), - ); - return; - } - default: - return; - } - } catch (error) { - notifyError(error, { - category: "query", - title: "Table quick action failed", - force: true, - }); - } - }, - [queryClient, activeConnectionEngine], - ); - - const handleSelectConnection = (nextConnection: ConnectionSummary) => { - if (connection?.id === nextConnection.id) { - return; - } - - activateConnectionMutation.mutate(nextConnection.id); - }; - - const handleRefreshConnection = useCallback( - (connectionTarget: ConnectionSummary) => { - void (async () => { - try { - await veloxDbRepository.refreshConnection(connectionTarget.id); - } catch (error) { - notifyError(error, { category: "connection" }); - return; - } - - void queryClient.invalidateQueries({ queryKey: queryKeys.connections() }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.databases(connectionTarget.id), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.databases(connectionTarget.id), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.tables(connectionTarget.id), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.tables(connectionTarget.id), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.queryEditorMetadata(connectionTarget.id), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.queryEditorMetadata(connectionTarget.id), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.foreignKeys(connectionTarget.id), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.foreignKeys(connectionTarget.id), - type: "active", - }); - - if (connection?.id === connectionTarget.id) { - void queryClient.invalidateQueries({ - queryKey: queryKeys.schema(connectionTarget.id, selectedTable), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.schema(connectionTarget.id, selectedTable), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.tableProperties(connectionTarget.id, selectedTable), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.tableProperties(connectionTarget.id, selectedTable), - type: "active", - }); - queryWorkspaceRef.current?.refreshFocusedResults(); - } - })(); - }, - [connection?.id, queryClient, selectedTable], - ); - - const handleRefreshTable = useCallback( - (connectionId: string, table: TableInfo) => { - void queryClient.invalidateQueries({ - queryKey: queryKeys.tables(connectionId), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.tables(connectionId), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.schema(connectionId, table), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.schema(connectionId, table), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.tableProperties(connectionId, table), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.tableProperties(connectionId, table), - type: "active", - }); - void queryClient.invalidateQueries({ - queryKey: queryKeys.tableIndexes(connectionId, table), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.tableIndexes(connectionId, table), - type: "active", - }); - - if ( - connection?.id === connectionId && - selectedTable?.schema === table.schema && - selectedTable?.name === table.name - ) { - queryWorkspaceRef.current?.refreshFocusedResults(); - } - }, - [connection?.id, queryClient, selectedTable?.name, selectedTable?.schema], - ); - - const handleRenameTableRequest = useCallback( - (_connectionId: string, table: TableInfo) => { - setSelectedTable(table); - queryWorkspaceRef.current?.appendQuerySql( - buildRenameTableSql(table, "new_table_name", connection?.engine ?? "postgres"), - ); - }, - [connection?.engine], - ); - - const handleDeleteTableRequest = useCallback( - (_connectionId: string, table: TableInfo) => { - setSelectedTable(table); - queryWorkspaceRef.current?.appendQuerySql( - buildDropTableSql(table, connection?.engine ?? "postgres"), - ); - }, - [connection?.engine], - ); - - const handleRenameConnectionRequest = useCallback( - (connectionTarget: ConnectionSummary) => { - setRenamingConnection(connectionTarget); - }, - [], - ); - - const handleRenameConnectionConfirm = useCallback( - (connectionTarget: ConnectionSummary, newName: string) => { - renameConnectionMutation.mutate( - { connectionId: connectionTarget.id, newName }, - { - onSuccess: (updated) => { - if (connection?.id === updated.id) { - setConnection(updated); - } - }, - }, - ); - setRenamingConnection(null); - }, - [connection?.id, renameConnectionMutation], - ); - - const handleDisconnectConnectionRequest = useCallback( - (connectionTarget: ConnectionSummary) => { - const confirmed = window.confirm( - t("connection.deleteConfirm", { name: connectionTarget.name }), - ); - if (!confirmed) return; - - deleteConnectionMutation.mutate(connectionTarget.id); - - if (connection?.id === connectionTarget.id) { - setConnection(null); - setSelectedTable(null); - setTableSearch(""); - setTablePropertiesDialogOpen(false); - setTablePropertiesTarget(null); - } - }, - [connection?.id, deleteConnectionMutation, t], - ); - - const handleCopyConnectionString = useCallback( - (target: ConnectionSummary) => { - const value = - target.engine === "sqlite" - ? `sqlite://${target.filePath ?? target.database}` - : `${target.engine === "mysql" ? "mysql" : "postgresql"}://${target.user}@${target.host}:${target.port}/${target.database}`; - void navigator.clipboard.writeText(value); - }, - [], - ); - - const handleTruncateTable = useCallback( - (_connectionId: string, table: TableInfo) => { - setSelectedTable(table); - if ((connection?.engine ?? "postgres") === "mysql") { - queryWorkspaceRef.current?.appendQuerySql( - `TRUNCATE TABLE ${quoteIdent(table.schema, "mysql")}.${quoteIdent(table.name, "mysql")};`, - ); - return; - } - if ((connection?.engine ?? "postgres") === "sqlite") { - queryWorkspaceRef.current?.appendQuerySql( - `DELETE FROM ${quoteIdent(table.name, "sqlite")};`, - ); - return; - } - queryWorkspaceRef.current?.appendQuerySql( - `TRUNCATE TABLE ${quoteIdent(table.schema, "postgres")}.${quoteIdent(table.name, "postgres")} RESTART IDENTITY CASCADE;`, - ); - }, - [connection?.engine], - ); - - const handleCopyTableName = useCallback( - (_connectionId: string, table: TableInfo) => { - const engine = connection?.engine ?? "postgres"; - const value = - engine === "sqlite" - ? quoteIdent(table.name, "sqlite") - : `${quoteIdent(table.schema, engine)}.${quoteIdent(table.name, engine)}`; - void navigator.clipboard.writeText(value); - }, - [connection?.engine], - ); - - const handleRefreshDatabases = useCallback( - (connectionId: string) => { - void queryClient.invalidateQueries({ - queryKey: queryKeys.databases(connectionId), - }); - void queryClient.refetchQueries({ - queryKey: queryKeys.databases(connectionId), - type: "active", - }); - }, - [queryClient], - ); - - const handleCopyDatabaseName = useCallback( - (_connectionId: string, database: string) => { - void navigator.clipboard.writeText(database); - }, - [], - ); - - const handleActivateConnectionForTab = useCallback( - (connectionId: string) => { - if (connection?.id === connectionId) { - return; - } - activateConnectionMutation.mutate(connectionId); - }, - [connection?.id, activateConnectionMutation], - ); - - const handleSidebarResizeStart = ( - event: ReactPointerEvent, - ) => { - const startX = event.clientX; - const startWidth = sidebarWidth; - - const handlePointerMove = (moveEvent: PointerEvent) => { - setSidebarWidth( - clampSidebarWidth(startWidth + moveEvent.clientX - startX), - ); - }; - - const handlePointerUp = () => { - window.removeEventListener("pointermove", handlePointerMove); - window.removeEventListener("pointerup", handlePointerUp); - }; - - window.addEventListener("pointermove", handlePointerMove); - window.addEventListener("pointerup", handlePointerUp); - }; - - useEffect(() => { - const onKeyDown = (event: KeyboardEvent) => { - const commandKey = event.metaKey || event.ctrlKey; - - if (commandKey && event.key.toLowerCase() === "p") { - event.preventDefault(); - setCommandPaletteOpen(true); - } - - if (commandKey && event.shiftKey && event.key.toLowerCase() === "c") { - event.preventDefault(); - setConnectionDialogOpen(true); - } - }; - - window.addEventListener("keydown", onKeyDown); - return () => window.removeEventListener("keydown", onKeyDown); - }, []); - - const layoutStyle = { - "--sidebar-width": `${isSidebarCollapsed ? 0 : sidebarWidth}px`, - } as CSSProperties; - - const connectionError = - connectMutation.error ?? activateConnectionMutation.error; - const connectionErrorMessage = - connectionError instanceof Error - ? connectionError.message - : t("connection.failedToConnect"); - const primaryKeyColumns = - tablePropertiesQuery.data - ?.filter((column) => column.isPrimaryKey) - .map((column) => column.columnName) ?? []; - const editableColumns = - tablePropertiesQuery.data - ?.filter((column) => !column.isPrimaryKey) - .map((column) => column.columnName) ?? []; - const hasSelectedTable = Boolean(selectedTable); - const hasQueryResult = focusedQueryCaps.hasResult; - const hasPrimaryKey = primaryKeyColumns.length > 0; - const isResultSingleTableEditable = - hasSelectedTable && - hasQueryResult && - hasPrimaryKey && - !tablePropertiesQuery.isError; - const saveDisabledReason = !hasSelectedTable - ? t("editor.selectTable") - : !hasQueryResult - ? t("editor.runQuery") - : tablePropertiesQuery.isLoading - ? t("editor.loadingMetadata") - : tablePropertiesQuery.isError - ? tablePropertiesErrorMessage - : !hasPrimaryKey - ? t("editor.requiresPrimaryKey") - : undefined; - - const handleSaveResultEdits = async (patches: ResultEditPatch[]) => { - if (!selectedTable || !connection?.id || patches.length === 0) { - return; - } - - await saveResultEditsMutation.mutateAsync({ - connectionId: connection.id, - engine: connection.engine, - table: selectedTable, - patches, - }); - - queryWorkspaceRef.current?.refreshFocusedResults(); - }; - - const handleDeleteRows = async ( - primaryKeys: Record[], - ) => { - if (!selectedTable || !connection?.id || primaryKeys.length === 0) { - return; - } - - await deleteRowsMutation.mutateAsync({ - connectionId: connection.id, - engine: connection.engine, - table: selectedTable, - primaryKeys, - }); - - queryWorkspaceRef.current?.refreshFocusedResults(); - }; - - const handleAskVeloxyChatSubmit = async ( - naturalPrompt: string, - requestId: string, - ): Promise => { - if (!connection?.id) { - const message = t("veloxy.selectConnection"); - setAskVeloxyError(message); - throw new Error(message); - } - if (!veloxyOpenRouterApiKey.trim()) { - const message = t("veloxy.addApiKey"); - setAskVeloxyError(message); - throw new Error(message); - } - if (!veloxyModel.trim()) { - const message = t("veloxy.chooseModel"); - setAskVeloxyError(message); - throw new Error(message); - } - setAskVeloxyPending(true); - setAskVeloxyError(null); - try { - return await veloxDbRepository.chatWithDb({ - connectionId: connection.id, - naturalPrompt, - requestId, - targetTable: selectedTable - ? { schema: selectedTable.schema, name: selectedTable.name } - : undefined, - providerConfig: { - apiKey: veloxyOpenRouterApiKey, - model: veloxyModel, - baseUrl: veloxyBaseUrl, - }, - maxRows: useSettings.getState().maxQueryRows, - }); - } catch (error) { - const message = - error instanceof Error - ? error.message - : t("veloxy.chatFailed"); - setAskVeloxyError(message); - notifyError(error, { category: "query", title: t("veloxy.chatFailed") }); - throw error instanceof Error ? error : new Error(message); - } finally { - setAskVeloxyPending(false); - } - }; - - const handleCancelVeloxyRequest = async () => { - try { - await veloxDbRepository.cancelVeloxyRequest(); - } catch (error) { - const message = - error instanceof Error ? error.message : "Failed to stop Veloxy."; - setAskVeloxyError(message); - } - }; - - const handleAskVeloxyActionSubmit = async ( - naturalPrompt: string, - ): Promise => { - if (!connection?.id) { - const message = t("veloxy.selectConnection"); - setAskVeloxyError(message); - throw new Error(message); - } - if (!veloxyOpenRouterApiKey.trim()) { - const message = t("veloxy.addApiKey"); - setAskVeloxyError(message); - throw new Error(message); - } - if (!veloxyModel.trim()) { - const message = t("veloxy.chooseModel"); - setAskVeloxyError(message); - throw new Error(message); - } - setAskVeloxyPending(true); - setAskVeloxyError(null); - try { - const response = await veloxDbRepository.generateSqlFromNl({ - connectionId: connection.id, - naturalPrompt, - targetTable: selectedTable - ? { schema: selectedTable.schema, name: selectedTable.name } - : undefined, - providerConfig: { - apiKey: veloxyOpenRouterApiKey, - model: veloxyModel, - baseUrl: veloxyBaseUrl, - }, - maxRows: useSettings.getState().maxQueryRows, - }); - const sql = response.sql.trim(); - const lower = sql.toLowerCase(); - const isReadIntent = response.intent === "select"; - const isLikelyLarge = - sql.length > 1800 || - (lower.includes("select") && !lower.includes(" limit ")) || - /\bcross\s+join\b|\bpg_sleep\s*\(/i.test(lower); - const canAutoRun = isReadIntent && !isLikelyLarge; - - if (canAutoRun) { - queryWorkspaceRef.current?.openTabWithSqlAndRun(sql); - notifySuccess(t("veloxy.generatedSql"), t("veloxy.autoRan")); - return { - response, - decision: "auto-ran", - }; - } - - queryWorkspaceRef.current?.openTabWithSql(sql); - return { - response, - decision: "needs-confirmation", - decisionReason: isReadIntent - ? t("veloxy.needsConfirmation") - : t("veloxy.nonReadConfirmation"), - pendingSql: sql, - }; - } catch (error) { - const message = - error instanceof Error - ? error.message - : t("veloxy.generateFailed"); - setAskVeloxyError(message); - notifyError(error, { category: "query", title: t("veloxy.generateFailed") }); - throw error instanceof Error ? error : new Error(message); - } finally { - setAskVeloxyPending(false); - } - }; - - const handleLoadVeloxyConversation = async (): Promise => { - if (!connection?.id) return { messages: [] }; - try { - return await veloxDbRepository.loadVeloxyConversation(connection.id); - } catch (error) { - const message = - error instanceof Error - ? error.message - : t("veloxy.loadFailed"); - setAskVeloxyError(message); - return { messages: [] }; - } - }; - - const handleClearVeloxyConversation = async () => { - if (!connection?.id) return; - try { - await veloxDbRepository.clearVeloxyConversation(connection.id); - } catch (error) { - const message = - error instanceof Error - ? error.message - : t("veloxy.clearFailed"); - setAskVeloxyError(message); - throw error; - } - }; - - - return ( -
- {!isSidebarCollapsed ? ( - <> -
- - Sidebar failed to render. -
- } - > - {connectionsQuery.isError ? ( -
- {connectionsErrorMessage} -
- ) : ( - setConnectionDialogOpen(true)} - onSelectConnection={handleSelectConnection} - onSelectTable={handleSelectTable} - onTableQuickAction={handleTableQuickAction} - onRefreshConnection={handleRefreshConnection} - onRefreshTable={handleRefreshTable} - onRenameConnection={handleRenameConnectionRequest} - onDisconnectConnection={handleDisconnectConnectionRequest} - onRenameTable={handleRenameTableRequest} - onDeleteTable={handleDeleteTableRequest} - onTruncateTable={handleTruncateTable} - onCopyTableName={handleCopyTableName} - onRefreshDatabases={handleRefreshDatabases} - onCopyDatabaseName={handleCopyDatabaseName} - onCopyConnectionString={handleCopyConnectionString} - onDatabaseSwitched={setConnection} - onToggleCollapsed={() => setIsSidebarCollapsed(true)} - /> - )} - -
-
- - ) : null} - -
-
-
-
- {isSidebarCollapsed ? ( - - ) : null} - - - setMainWorkspace(value as "query" | "model") - } - className="shrink-0" - > - - - {t("workspace.query")} - - - {t("workspace.model")} - - - - -
-

- VeloxDB.dev -

-

- {connection - ? connectionHeadline(connection) - : t("workspace.chooseConnection")} -

-
-
- -
- - -
-
-
- - {mainWorkspace === "query" ? ( - setConnectionDialogOpen(true)} - resultsHeight={resultsHeight} - onResultsHeightChange={setResultsHeight} - selectedTable={selectedTable} - schemaLoading={schemaQuery.isLoading} - schemaError={schemaQuery.isError ? schemaErrorMessage : null} - columnCount={schemaQuery.data?.length ?? null} - primaryKeyColumns={primaryKeyColumns} - editableColumns={editableColumns} - saveDisabledReason={saveDisabledReason} - isResultSingleTableEditable={isResultSingleTableEditable} - saveResultEditsMutation={saveResultEditsMutation} - onSaveResultEdits={handleSaveResultEdits} - onDeleteRows={handleDeleteRows} - onFocusedTabCapabilitiesChange={setFocusedQueryCaps} - onActivateConnectionForTab={handleActivateConnectionForTab} - insertRowTrigger={insertRowTrigger} - insertConnectionId={connection?.id ?? null} - insertTable={selectedTable} - canInsertRow={Boolean(connection?.id && selectedTable)} - onInsertRowSuccess={handleInsertRowSuccess} - onOpenAddRow={ - connection?.id && selectedTable - ? requestInsertRow - : undefined - } - askVeloxySidebar={(onClose) => ( - { - setSettingsOpen(true); - }} - onChatSubmit={handleAskVeloxyChatSubmit} - onActionSubmit={handleAskVeloxyActionSubmit} - onLoadConversation={handleLoadVeloxyConversation} - onClearConversation={handleClearVeloxyConversation} - onConfirmRun={async (sql) => { - queryWorkspaceRef.current?.openTabWithSqlAndRun(sql); - notifySuccess(t("veloxy.queryExecuted")); - }} - onInsertSql={(sql) => { - queryWorkspaceRef.current?.appendQuerySql(sql); - }} - onReplaceSql={(sql) => { - queryWorkspaceRef.current?.replaceQuerySql(sql); - }} - onOpenTabWithSql={(sql) => { - queryWorkspaceRef.current?.openTabWithSql(sql); - }} - onCancelRequest={handleCancelVeloxyRequest} - errorMessage={askVeloxyError} - /> - )} - /> - ) : connection?.id ? ( - - {connection.engine === "postgres" ? ( - - ) : ( -
- {t("model.postgresOnly", { engine: engineLabel(connection.engine) })} -
- )} - - ) : ( -
- {t("model.connectToUse")} -
- )} -
- - { - connectMutation.mutate(values); - }} - isPending={connectMutation.isPending} - /> - - setRenamingConnection(null)} - /> - - { - setTablePropertiesDialogOpen(nextOpen); - if (!nextOpen) setTablePropertiesTarget(null); - }} - connectionId={tablePropertiesTarget?.connectionId} - tablePropertyEditingSupported={connection?.tablePropertyEditingSupported} - table={tablePropertiesTarget?.table ?? null} - /> - - setConnectionDialogOpen(true)} - onRunLastQuery={() => { - queryWorkspaceRef.current?.runLastQuery(); - }} - onSelectTable={handleSelectTable} - /> - - -
- ); -} function App() { const [onboardingDone, setOnboardingDone] = useState(() => diff --git a/src/components/VeloxApp.tsx b/src/components/VeloxApp.tsx new file mode 100644 index 0000000..edd5b82 --- /dev/null +++ b/src/components/VeloxApp.tsx @@ -0,0 +1,340 @@ +import { GearIcon, RobotIcon, SidebarSimpleIcon } from "@phosphor-icons/react"; +import { + type CSSProperties, + type PointerEvent as ReactPointerEvent, + useEffect, + useRef, +} from "react"; +import { useTranslation } from "react-i18next"; + +import { ErrorBoundary } from "@/components/ErrorBoundary"; +import { Button } from "@/components/ui/button"; +import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { CommandPalette } from "@/features/commands/components/CommandPalette"; +import { ShortcutSheet } from "@/features/commands/components/ShortcutSheet"; +import { SettingsDialog } from "@/features/commands/components/SettingsDialog"; +import { ConnectionDialog } from "@/features/connections/components/ConnectionDialog"; +import { RenameConnectionDialog } from "@/features/connections/components/RenameConnectionDialog"; +import { ConnectionsSidebarTree } from "@/features/connections/components/ConnectionsSidebarTree"; +import { QueryWorkspace } from "@/features/queries/components/QueryWorkspace"; +import { AskVeloxySidebar } from "@/features/queries/components/AskVeloxyDialog"; +import { ModelWorkspace } from "@/features/model/components/ModelWorkspace"; +import { TablePropertiesDialog } from "@/features/schema/components/TablePropertiesDialog"; + +import { + clampSidebarWidth, + connectionHeadline, + engineLabel, + useAppState, +} from "@/hooks/useAppState"; + +export function VeloxApp() { + const { t } = useTranslation(); + const queryWorkspaceRef = useRef(null); + + const app = useAppState(queryWorkspaceRef); + + const { + connection, + focusedQueryCaps, + tableSearch, setTableSearch, + selectedTable, + settingsOpen, setSettingsOpen, + commandPaletteOpen, setCommandPaletteOpen, + connectionDialogOpen, setConnectionDialogOpen, + renamingConnection, + isSidebarCollapsed, setIsSidebarCollapsed, + sidebarWidth, setSidebarWidth, + resultsHeight, setResultsHeight, + tablePropertiesDialogOpen, setTablePropertiesDialogOpen, + tablePropertiesTarget, + insertRowTrigger, + mainWorkspace, setMainWorkspace, + askVeloxyPending, + askVeloxyError, + isDark, + veloxyModel, + veloxyOpenRouterApiKey, + connectionsQuery, tablesQuery, schemaQuery, + saveResultEditsMutation, + connectMutation, + tablesForUi, + connectionError, connectionErrorMessage, + connectionsErrorMessage, tablesErrorMessage, + schemaErrorMessage, + primaryKeyColumns, editableColumns, + isResultSingleTableEditable, saveDisabledReason, + handleSelectTable, + handleTableQuickAction, + handleSelectConnection, + handleRefreshConnection, + handleRefreshTable, + handleRenameTableRequest, + handleDeleteTableRequest, + handleRenameConnectionRequest, + handleRenameConnectionConfirm, + handleDisconnectConnectionRequest, + handleCopyConnectionString, + handleTruncateTable, + handleCopyTableName, + handleRefreshDatabases, + handleCopyDatabaseName, + handleActivateConnectionForTab, + handleSaveResultEdits, + handleDeleteRows, + requestInsertRow, + handleInsertRowSuccess, + handleAskVeloxyChatSubmit, + handleCancelVeloxyRequest, + handleAskVeloxyActionSubmit, + handleLoadVeloxyConversation, + handleClearVeloxyConversation, + } = app; + + const handleSidebarResizeStart = (event: ReactPointerEvent) => { + const startX = event.clientX; + const startWidth = sidebarWidth; + const handlePointerMove = (moveEvent: PointerEvent) => { + setSidebarWidth(clampSidebarWidth(startWidth + moveEvent.clientX - startX)); + }; + const handlePointerUp = () => { + window.removeEventListener("pointermove", handlePointerMove); + window.removeEventListener("pointerup", handlePointerUp); + }; + window.addEventListener("pointermove", handlePointerMove); + window.addEventListener("pointerup", handlePointerUp); + }; + + useEffect(() => { + const onKeyDown = (event: KeyboardEvent) => { + const commandKey = event.metaKey || event.ctrlKey; + if (commandKey && event.key.toLowerCase() === "p") { + event.preventDefault(); + setCommandPaletteOpen(true); + } + if (commandKey && event.shiftKey && event.key.toLowerCase() === "c") { + event.preventDefault(); + setConnectionDialogOpen(true); + } + }; + window.addEventListener("keydown", onKeyDown); + return () => window.removeEventListener("keydown", onKeyDown); + }, [setCommandPaletteOpen, setConnectionDialogOpen]); + + const layoutStyle = { + "--sidebar-width": `${isSidebarCollapsed ? 0 : sidebarWidth}px`, + } as CSSProperties; + + return ( +
+ {!isSidebarCollapsed ? ( + <> +
+ Sidebar failed to render.
}> + {connectionsQuery.isError ? ( +
{connectionsErrorMessage}
+ ) : ( + setConnectionDialogOpen(true)} + onSelectConnection={handleSelectConnection} + onSelectTable={handleSelectTable} + onTableQuickAction={handleTableQuickAction} + onRefreshConnection={handleRefreshConnection} + onRefreshTable={handleRefreshTable} + onRenameConnection={handleRenameConnectionRequest} + onDisconnectConnection={handleDisconnectConnectionRequest} + onRenameTable={handleRenameTableRequest} + onDeleteTable={handleDeleteTableRequest} + onTruncateTable={handleTruncateTable} + onCopyTableName={handleCopyTableName} + onRefreshDatabases={handleRefreshDatabases} + onCopyDatabaseName={handleCopyDatabaseName} + onCopyConnectionString={handleCopyConnectionString} + onDatabaseSwitched={app.setConnection} + onToggleCollapsed={() => setIsSidebarCollapsed(true)} + /> + )} + +
+
+ + ) : null} + +
+
+
+
+ {isSidebarCollapsed ? ( + + ) : null} + setMainWorkspace(value as "query" | "model")} className="shrink-0"> + + {t("workspace.query")} + + {t("workspace.model")} + + + +
+

VeloxDB.dev

+

+ {connection ? connectionHeadline(connection) : t("workspace.chooseConnection")} +

+
+
+
+ + + +
+
+
+ + {mainWorkspace === "query" ? ( + setConnectionDialogOpen(true)} + resultsHeight={resultsHeight} + onResultsHeightChange={setResultsHeight} + selectedTable={selectedTable} + schemaLoading={schemaQuery?.isLoading ?? false} + schemaError={schemaQuery?.isError ? schemaErrorMessage : null} + columnCount={schemaQuery?.data?.length ?? null} + primaryKeyColumns={primaryKeyColumns} + editableColumns={editableColumns} + saveDisabledReason={saveDisabledReason} + isResultSingleTableEditable={isResultSingleTableEditable} + saveResultEditsMutation={saveResultEditsMutation} + onSaveResultEdits={handleSaveResultEdits} + onDeleteRows={handleDeleteRows} + onFocusedTabCapabilitiesChange={app.setFocusedQueryCaps} + onActivateConnectionForTab={handleActivateConnectionForTab} + insertRowTrigger={insertRowTrigger} + insertConnectionId={connection?.id ?? null} + insertTable={selectedTable} + canInsertRow={Boolean(connection?.id && selectedTable)} + onInsertRowSuccess={handleInsertRowSuccess} + onOpenAddRow={connection?.id && selectedTable ? requestInsertRow : undefined} + askVeloxySidebar={(onClose) => ( + { setSettingsOpen(true); }} + onChatSubmit={handleAskVeloxyChatSubmit} + onActionSubmit={handleAskVeloxyActionSubmit} + onLoadConversation={handleLoadVeloxyConversation} + onClearConversation={handleClearVeloxyConversation} + onConfirmRun={async (sql) => { + queryWorkspaceRef.current?.openTabWithSqlAndRun(sql); + }} + onInsertSql={(sql) => { queryWorkspaceRef.current?.appendQuerySql(sql); }} + onReplaceSql={(sql) => { queryWorkspaceRef.current?.replaceQuerySql(sql); }} + onOpenTabWithSql={(sql) => { queryWorkspaceRef.current?.openTabWithSql(sql); }} + onCancelRequest={handleCancelVeloxyRequest} + errorMessage={askVeloxyError} + /> + )} + /> + ) : connection?.id ? ( + + {connection.engine === "postgres" ? ( + + ) : ( +
+ {t("model.postgresOnly", { engine: engineLabel(connection.engine) })} +
+ )} + + ) : ( +
+ {t("model.connectToUse")} +
+ )} +
+ + { connectMutation.mutate(values); }} + isPending={connectMutation.isPending} + /> + app.setRenamingConnection(null)} + /> + { + setTablePropertiesDialogOpen(nextOpen); + if (!nextOpen) app.setTablePropertiesTarget(null); + }} + connectionId={tablePropertiesTarget?.connectionId} + tablePropertyEditingSupported={connection?.tablePropertyEditingSupported} + table={tablePropertiesTarget?.table ?? null} + /> + setConnectionDialogOpen(true)} + onRunLastQuery={() => { queryWorkspaceRef.current?.runLastQuery(); }} + onSelectTable={handleSelectTable} + /> + + +
+ ); +} diff --git a/src/data/types.ts b/src/data/types.ts index e8610f1..a50f0b9 100644 --- a/src/data/types.ts +++ b/src/data/types.ts @@ -1,6 +1,6 @@ /** PostgreSQL `sslmode`-style TLS (lowercase in JSON for Tauri). */ export type ConnectionSslMode = 'disable' | 'prefer' | 'require' -export type DatabaseEngine = 'postgres' | 'mysql' | 'sqlite' +export type DatabaseEngine = 'postgres' | 'mysql' | 'sqlite' | 'mongo' | 'duckdb' | 'redis' export type SshAuthMethod = 'keyfile' | 'password' @@ -26,6 +26,7 @@ export type ConnectionInput = { user: string password: string sslMode: ConnectionSslMode + srvEnabled?: boolean sshConfig?: SshConfig | null extraParams?: Record | null } diff --git a/src/features/commands/components/SettingsDialog.tsx b/src/features/commands/components/SettingsDialog.tsx index c43f358..86e744e 100644 --- a/src/features/commands/components/SettingsDialog.tsx +++ b/src/features/commands/components/SettingsDialog.tsx @@ -22,7 +22,8 @@ import { Input } from '@/components/ui/input' import { saveOpenRouterApiKey } from '@/lib/openrouter-credentials' import { fetchOpenRouterModels, OPENROUTER_POPULAR_MODELS, type OpenRouterModelOption } from '@/lib/openrouter-models' import { cn } from '@/lib/utils' -import { useSettings, type AppTheme, type FontSize, type NullDisplay } from '@/lib/settings' +import { useSettings, type AppTheme, type FontSize, type NullDisplay, themeLabels } from '@/lib/settings' +import { useUpdateCheck } from '@/hooks/useUpdateCheck' import pkg from '../../../../package.json' const GITHUB_REPO = 'abeni16/veloxdb' @@ -74,25 +75,13 @@ export function SettingsDialog({ open, onOpenChange }: { open: boolean; onOpenCh input.click() }, []) - const [updateStatus, setUpdateStatus] = useState<'idle' | 'checking' | 'up-to-date' | 'available' | 'error'>('idle') - const [latestVersion, setLatestVersion] = useState(null) const [modelOptions, setModelOptions] = useState(OPENROUTER_POPULAR_MODELS) const [modelsStatus, setModelsStatus] = useState<'idle' | 'loading' | 'error'>('idle') const [modelsError, setModelsError] = useState(null) - const checkForUpdates = useCallback(async () => { - setUpdateStatus('checking') - try { - const res = await fetch(`https://api.github.com/repos/${GITHUB_REPO}/releases/latest`) - if (!res.ok) throw new Error(`HTTP ${res.status}`) - const release = await res.json() as { tag_name: string } - const latest = release.tag_name.replace(/^v/, '') - setLatestVersion(latest) - setUpdateStatus(latest === pkg.version ? 'up-to-date' : 'available') - } catch { - setUpdateStatus('error') - } - }, []) + const { data: updateInfo, isLoading: isCheckingUpdates, isError: updateCheckFailed, refetch: checkForUpdates } = useUpdateCheck( + { enabled: open }, + ) const refreshOpenRouterModels = useCallback(async () => { setModelsStatus('loading') @@ -161,8 +150,7 @@ export function SettingsDialog({ open, onOpenChange }: { open: boolean; onOpenCh opts={[{ v: 'en', l: 'English' }, { v: 'zh', l: '中文' }]} /> - useSettings.setState({ fontSize: v as FontSize })} @@ -359,26 +347,26 @@ export function SettingsDialog({ open, onOpenChange }: { open: boolean; onOpenCh variant="outline" size="sm" className="h-8 text-xs" - onClick={checkForUpdates} - disabled={updateStatus === 'checking'} + onClick={() => checkForUpdates()} + disabled={isCheckingUpdates} > - {updateStatus === 'checking' ? t("settings.checking") : t("settings.checkNow")} + {isCheckingUpdates ? t("settings.checking") : t("settings.checkNow")} - {updateStatus === 'up-to-date' && ( + {updateInfo && !updateInfo.hasUpdate && ( {t("settings.upToDate")} )} - {updateStatus === 'available' && latestVersion && ( + {updateInfo?.hasUpdate && updateInfo.latestVersion && ( - {t("settings.available", { version: latestVersion })} + {t("settings.available", { version: updateInfo.latestVersion })} )} - {updateStatus === 'error' && ( + {updateCheckFailed && ( {t("settings.failedToCheck")} )} @@ -403,6 +391,22 @@ function Select({ value, onChange, opts }: { value: string; onChange: (v: string return } +const THEME_ORDER: AppTheme[] = ['system', 'light', 'dark', 'sepia', 'ocean', 'forest', 'rose', 'slate', 'amber'] + +function ThemeSelect({ value, onChange }: { value: string; onChange: (v: string) => void }) { + return ( + + ) +} + function Toggle({ value, onChange }: { value: boolean; onChange: (v: boolean) => void }) { return } diff --git a/src/features/connections/components/ConnectionDialog.tsx b/src/features/connections/components/ConnectionDialog.tsx index f4269b5..dd21c27 100644 --- a/src/features/connections/components/ConnectionDialog.tsx +++ b/src/features/connections/components/ConnectionDialog.tsx @@ -1,121 +1,122 @@ -import { useMemo, useState, type ReactNode } from 'react' +import { useMemo, useState, useEffect, useCallback, type ReactNode } from 'react' import { useForm, useWatch } from 'react-hook-form' import { z } from 'zod' import { zodResolver } from '@hookform/resolvers/zod' import { useTranslation } from 'react-i18next' import { open as openFilePicker } from '@tauri-apps/plugin-dialog' -import { FolderOpenIcon } from '@phosphor-icons/react' +import { FolderOpenIcon, PlugsConnectedIcon } from '@phosphor-icons/react' import { Button } from '@/components/ui/button' import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, + Dialog, DialogContent, DialogDescription, + DialogFooter, DialogHeader, DialogTitle, } from '@/components/ui/dialog' import { Input } from '@/components/ui/input' import { - InputGroup, - InputGroupAddon, - InputGroupButton, - InputGroupInput, + InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput, } from '@/components/ui/input-group' -import type { ConnectionInput, DatabaseEngine, SshAuthMethod } from '@/data/types' +import type { ConnectionInput, DatabaseEngine } from '@/data/types' import { cn } from '@/lib/utils' import { parseConnectionString, buildConnectionString } from '@/lib/connection-string' +import { ENGINE_FORM_CONFIG, type EngineFormConfig } from './engine-form-config' -const sslModeSchema = z.enum(['disable', 'prefer', 'require']) -const sshAuthMethodSchema = z.enum(['keyfile', 'password']) - -const connectionSchema = z - .object({ - name: z.string().min(2, 'Enter a connection name.'), - engine: z.enum(['postgres', 'mysql', 'sqlite']), - host: z.string(), - port: z.coerce.number().int().min(1).max(65535), - database: z.string(), - filePath: z.string().optional(), - user: z.string(), - password: z.string(), - sslMode: sslModeSchema, - sshEnabled: z.boolean(), - sshHost: z.string().optional(), - sshPort: z.coerce.number().int().min(1).max(65535).optional(), - sshUser: z.string().optional(), - sshAuthMethod: sshAuthMethodSchema.optional(), - sshPassword: z.string().optional(), - sshPrivateKeyPath: z.string().optional(), - sshPassphrase: z.string().optional(), - }) - .superRefine((values, ctx) => { - if (values.engine !== 'sqlite') { - if (!values.host || values.host.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'Host is required.', - path: ['host'], - }) - } - if (!values.database || values.database.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'Database is required.', - path: ['database'], - }) - } - if (!values.user || values.user.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'User is required.', - path: ['user'], - }) - } - if (!values.password || values.password.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'Password is required.', - path: ['password'], - }) - } - } - if (values.engine === 'sqlite') { - if (!values.filePath || values.filePath.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'SQLite file path is required.', - path: ['filePath'], - }) - } - return - } - if (!values.sshEnabled) return - if (!values.sshHost || values.sshHost.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'SSH host is required.', - path: ['sshHost'], - }) - } - if (!values.sshUser || values.sshUser.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'SSH user is required.', - path: ['sshUser'], - }) - } - if (values.sshAuthMethod === 'password') { - if (!values.sshPassword || values.sshPassword.trim().length === 0) { - ctx.addIssue({ - code: z.ZodIssueCode.custom, - message: 'SSH password is required.', - path: ['sshPassword'], - }) - } +type InputMode = 'string' | 'fields' +type SshAuthMethodForm = 'keyfile' | 'password' + +const engineOptions: { value: DatabaseEngine; label: string }[] = [ + { value: 'postgres', label: 'PostgreSQL' }, + { value: 'mysql', label: 'MySQL' }, + { value: 'sqlite', label: 'SQLite' }, + { value: 'mongo', label: 'MongoDB' }, + { value: 'duckdb', label: 'DuckDB' }, + { value: 'redis', label: 'Redis' }, +] + +// ── Database Flavors ──────────────────────────────────────────── + +type DatabaseFlavor = { + key: string + label: string + engine: DatabaseEngine + defaultHost: string + defaultPort: number + defaultDatabase: string + defaultUser: string + defaultSsl: 'disable' | 'prefer' | 'require' + description: string +} + +const DATABASE_FLAVORS: DatabaseFlavor[] = [ + { key: 'postgres', label: 'PostgreSQL', engine: 'postgres', defaultHost: '127.0.0.1', defaultPort: 5432, defaultDatabase: 'postgres', defaultUser: 'postgres', defaultSsl: 'prefer', description: 'Local or self-hosted PostgreSQL' }, + { key: 'supabase', label: 'Supabase', engine: 'postgres', defaultHost: 'db.xxxxx.supabase.co', defaultPort: 5432, defaultDatabase: 'postgres', defaultUser: 'postgres', defaultSsl: 'require', description: 'Hosted PostgreSQL with real-time, auth, storage' }, + { key: 'neon', label: 'Neon', engine: 'postgres', defaultHost: 'ep-xxxxx.us-east-1.aws.neon.tech', defaultPort: 5432, defaultDatabase: 'neondb', defaultUser: 'neondb_owner', defaultSsl: 'require', description: 'Serverless PostgreSQL with branching' }, + { key: 'cockroachdb', label: 'CockroachDB', engine: 'postgres', defaultHost: '127.0.0.1', defaultPort: 26257,defaultDatabase: 'defaultdb', defaultUser: 'root', defaultSsl: 'require', description: 'Distributed SQL, PostgreSQL-compatible' }, + { key: 'timescaledb', label: 'TimescaleDB', engine: 'postgres', defaultHost: '127.0.0.1', defaultPort: 5432, defaultDatabase: 'postgres', defaultUser: 'postgres', defaultSsl: 'prefer', description: 'Time-series PostgreSQL extension' }, + { key: 'yugabytedb', label: 'YugabyteDB', engine: 'postgres', defaultHost: '127.0.0.1', defaultPort: 5433, defaultDatabase: 'yugabyte', defaultUser: 'yugabyte', defaultSsl: 'prefer', description: 'Distributed PostgreSQL-compatible' }, + { key: 'mysql', label: 'MySQL', engine: 'mysql', defaultHost: '127.0.0.1', defaultPort: 3306, defaultDatabase: '', defaultUser: 'root', defaultSsl: 'prefer', description: 'Local or self-hosted MySQL' }, + { key: 'mariadb', label: 'MariaDB', engine: 'mysql', defaultHost: '127.0.0.1', defaultPort: 3306, defaultDatabase: '', defaultUser: 'root', defaultSsl: 'prefer', description: 'Community fork of MySQL' }, + { key: 'planetscale', label: 'PlanetScale', engine: 'mysql', defaultHost: 'aws.connect.psdb.cloud', defaultPort: 3306, defaultDatabase: '', defaultUser: 'root', defaultSsl: 'require', description: 'Serverless MySQL platform' }, + { key: 'tidb', label: 'TiDB', engine: 'mysql', defaultHost: '127.0.0.1', defaultPort: 4000, defaultDatabase: 'test', defaultUser: 'root', defaultSsl: 'prefer', description: 'Distributed MySQL-compatible (HTAP)' }, + { key: 'aurora', label: 'Aurora MySQL',engine: 'mysql', defaultHost: 'xxxxx.cluster-xxx.us-east-1.rds.amazonaws.com', defaultPort: 3306, defaultDatabase: '', defaultUser: 'admin', defaultSsl: 'require', description: 'AWS MySQL-compatible cloud database' }, +] + +const FLAVORS_BY_ENGINE: Record = {} +for (const f of DATABASE_FLAVORS) { + (FLAVORS_BY_ENGINE[f.engine] ??= []).push(f) +} + +// ── Zod schema ────────────────────────────────────────────────── + +const connectionSchema = z.object({ + name: z.string().min(2, 'Enter a connection name.'), + engine: z.enum(['postgres', 'mysql', 'sqlite', 'mongo', 'duckdb', 'redis'] as const), + host: z.string(), + port: z.coerce.number().int().min(1).max(65535), + database: z.string(), + filePath: z.string().optional(), + user: z.string(), + password: z.string(), + sslMode: z.enum(['disable', 'prefer', 'require'] as const), + srv: z.boolean().optional(), + sshEnabled: z.boolean(), + sshHost: z.string().optional(), + sshPort: z.coerce.number().int().min(1).max(65535).optional(), + sshUser: z.string().optional(), + sshAuthMethod: z.enum(['keyfile', 'password'] as const).optional(), + sshPassword: z.string().optional(), + sshPrivateKeyPath: z.string().optional(), + sshPassphrase: z.string().optional(), +}).superRefine((values, ctx) => { + const cfg = ENGINE_FORM_CONFIG[values.engine] + + if (cfg.showFilePicker) { + if (!values.filePath || values.filePath.trim().length === 0) { + ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'File path is required.', path: ['filePath'] }) } - }) + return + } + + if (cfg.fields.host.label && (!values.host || values.host.trim().length === 0)) { + ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'Host is required.', path: ['host'] }) + } + if (cfg.fields.database.label && (!values.database || values.database.trim().length === 0)) { + ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'Database is required.', path: ['database'] }) + } + if (cfg.requireUser && (!values.user || values.user.trim().length === 0)) { + ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'User is required.', path: ['user'] }) + } + if (cfg.requirePassword && (!values.password || values.password.trim().length === 0)) { + ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'Password is required.', path: ['password'] }) + } + + if (!values.sshEnabled) return + if (!values.sshHost?.trim()) ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'SSH host is required.', path: ['sshHost'] }) + if (!values.sshUser?.trim()) ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'SSH user is required.', path: ['sshUser'] }) + if (values.sshAuthMethod === 'password' && !values.sshPassword?.trim()) { + ctx.addIssue({ code: z.ZodIssueCode.custom, message: 'SSH password is required.', path: ['sshPassword'] }) + } +}) type ConnectionDialogProps = { open: boolean @@ -124,115 +125,100 @@ type ConnectionDialogProps = { isPending?: boolean } -function Field({ - label, - error, - inputId, - children, -}: { - label: string - error?: string - inputId: string - children: ReactNode -}) { +// ── Field helper ───────────────────────────────────────────────── + +function Field({ label, error, inputId, children }: { label: string; error?: string; inputId: string; children: ReactNode }) { return ( -