From 1b7631eddc6bb40f8de15673510933395c485a85 Mon Sep 17 00:00:00 2001
From: Jon Heinritz <jon.heinritz@protonmail.com>
Date: Fri, 30 Jun 2023 23:59:35 +0200
Subject: [PATCH] Add CLI autocompletion using clap_complete (#2559)

---
 Cargo.lock                  | 10 ++++++++++
 sqlx-cli/Cargo.toml         |  5 ++++-
 sqlx-cli/src/completions.rs | 10 ++++++++++
 sqlx-cli/src/lib.rs         |  5 +++++
 sqlx-cli/src/opt.rs         |  6 ++++++
 5 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 sqlx-cli/src/completions.rs

diff --git a/Cargo.lock b/Cargo.lock
index 3ff716b8ed..29340d3696 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -510,6 +510,15 @@ dependencies = [
  "textwrap",
 ]
 
+[[package]]
+name = "clap_complete"
+version = "3.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f7a2e0a962c45ce25afce14220bc24f9dade0a1787f185cecf96bfba7847cd8"
+dependencies = [
+ "clap",
+]
+
 [[package]]
 name = "clap_derive"
 version = "3.2.7"
@@ -2497,6 +2506,7 @@ dependencies = [
  "cargo_metadata",
  "chrono",
  "clap",
+ "clap_complete",
  "console",
  "dotenvy",
  "filetime",
diff --git a/sqlx-cli/Cargo.toml b/sqlx-cli/Cargo.toml
index ad7c685e63..785d35b32e 100644
--- a/sqlx-cli/Cargo.toml
+++ b/sqlx-cli/Cargo.toml
@@ -34,6 +34,7 @@ sqlx = { workspace = true, default-features = false, features = [
 ] }
 futures = "0.3.19"
 clap = { version = "3.1.0", features = ["derive", "env"] }
+clap_complete = { version = "3.1.0", optional = true }
 chrono = { version = "0.4.19", default-features = false, features = ["clock"] }
 anyhow = "1.0.52"
 url = { version = "2.2.2", default-features = false }
@@ -52,7 +53,7 @@ filetime = "0.2"
 backoff = { version = "0.4.0", features = ["futures", "tokio"] }
 
 [features]
-default = ["postgres", "sqlite", "mysql", "native-tls"]
+default = ["postgres", "sqlite", "mysql", "native-tls", "completions"]
 rustls = ["sqlx/runtime-tokio-rustls"]
 native-tls = ["sqlx/runtime-tokio-native-tls"]
 
@@ -63,3 +64,5 @@ sqlite = ["sqlx/sqlite"]
 
 # workaround for musl + openssl issues
 openssl-vendored = ["openssl/vendored"]
+
+completions = ["dep:clap_complete"]
diff --git a/sqlx-cli/src/completions.rs b/sqlx-cli/src/completions.rs
new file mode 100644
index 0000000000..9e3688e9fb
--- /dev/null
+++ b/sqlx-cli/src/completions.rs
@@ -0,0 +1,10 @@
+use std::io;
+
+use clap::CommandFactory;
+use clap_complete::{generate, Shell};
+
+use crate::opt::Command;
+
+pub fn run(shell: Shell) {
+    generate(shell, &mut Command::command(), "sqlx", &mut io::stdout())
+}
diff --git a/sqlx-cli/src/lib.rs b/sqlx-cli/src/lib.rs
index 8b52d950fa..13786efe90 100644
--- a/sqlx-cli/src/lib.rs
+++ b/sqlx-cli/src/lib.rs
@@ -12,6 +12,8 @@ mod database;
 mod metadata;
 // mod migration;
 // mod migrator;
+#[cfg(feature = "completions")]
+mod completions;
 mod migrate;
 mod opt;
 mod prepare;
@@ -68,6 +70,9 @@ pub async fn run(opt: Opt) -> Result<()> {
             connect_opts,
             args,
         } => prepare::run(check, workspace, connect_opts, args).await?,
+
+        #[cfg(feature = "completions")]
+        Command::Completions { shell } => completions::run(shell),
     };
 
     Ok(())
diff --git a/sqlx-cli/src/opt.rs b/sqlx-cli/src/opt.rs
index 17172e8a48..fb901d565c 100644
--- a/sqlx-cli/src/opt.rs
+++ b/sqlx-cli/src/opt.rs
@@ -1,6 +1,8 @@
 use std::ops::{Deref, Not};
 
 use clap::{Args, Parser};
+#[cfg(feature = "completions")]
+use clap_complete::Shell;
 
 #[derive(Parser, Debug)]
 #[clap(version, about, author)]
@@ -46,6 +48,10 @@ pub enum Command {
 
     #[clap(alias = "mig")]
     Migrate(MigrateOpt),
+
+    #[cfg(feature = "completions")]
+    /// Generate shell completions for the specified shell
+    Completions { shell: Shell },
 }
 
 /// Group of commands for creating and dropping your database.