From 31b0c0f115a9f28a51c1b7f598159de6dac240e8 Mon Sep 17 00:00:00 2001
From: Yugo Nagata <nagata@sraoss.co.jp>
Date: Thu, 5 Jun 2025 09:39:24 +0900
Subject: [PATCH v4 2/3] Add tab completion support for COPY ... TO/FROM STDIN,
 STDOUT, and PROGRAM

Previously, tab completion for COPY only suggested filenames after TO or
FROM, even though STDIN, STDOUT, and PROGRAM are also valid options.
This commit extends the completion to include these keywords. After PROGRAM,
filename suggestions are shown as potential command names.

To support this, a new macro COMPLETE_WITH_FILES_PLUS is introduced, allowing
both literal keywords and filenames to be included in the completion results.
---
 src/bin/psql/tab-complete.in.c | 119 ++++++++++++++++++++++++++-------
 1 file changed, 96 insertions(+), 23 deletions(-)

diff --git a/src/bin/psql/tab-complete.in.c b/src/bin/psql/tab-complete.in.c
index 1b4855edab6..c565c0218e1 100644
--- a/src/bin/psql/tab-complete.in.c
+++ b/src/bin/psql/tab-complete.in.c
@@ -443,13 +443,23 @@ do { \
 	matches = rl_completion_matches(text, complete_from_schema_query); \
 } while (0)
 
-#define COMPLETE_WITH_FILES(escape, force_quote) \
+#define COMPLETE_WITH_FILES_LIST(escape, force_quote, list) \
 do { \
 	completion_charp = escape; \
+	completion_charpp = list; \
 	completion_force_quote = force_quote; \
 	matches = rl_completion_matches(text, complete_from_files); \
 } while (0)
 
+#define COMPLETE_WITH_FILES(escape, force_quote) \
+	COMPLETE_WITH_FILES_LIST(escape, force_quote, NULL)
+
+#define COMPLETE_WITH_FILES_PLUS(escape, force_quote, ...) \
+do { \
+	static const char *const list[] = { __VA_ARGS__, NULL }; \
+	COMPLETE_WITH_FILES_LIST(escape, force_quote, list); \
+} while (0)
+
 #define COMPLETE_WITH_GENERATOR(function) \
 	matches = rl_completion_matches(text, function)
 
@@ -1481,6 +1491,7 @@ static void append_variable_names(char ***varnames, int *nvars,
 static char **complete_from_variables(const char *text,
 									  const char *prefix, const char *suffix, bool need_value);
 static char *complete_from_files(const char *text, int state);
+static char *_complete_from_files(const char *text, int state);
 
 static char *pg_strdup_keyword_case(const char *s, const char *ref);
 static char *escape_string(const char *text);
@@ -3316,42 +3327,60 @@ match_previous_words(int pattern_id,
 	/* Complete COPY <sth> */
 	else if (Matches("COPY|\\copy", MatchAny))
 		COMPLETE_WITH("FROM", "TO");
-	/* Complete COPY <sth> FROM|TO with filename */
-	else if (Matches("COPY", MatchAny, "FROM|TO"))
-		COMPLETE_WITH_FILES("", true);	/* COPY requires quoted filename */
-	else if (Matches("\\copy", MatchAny, "FROM|TO"))
-		COMPLETE_WITH_FILES("", false);
-
-	/* Complete COPY <sth> TO <sth> */
-	else if (Matches("COPY|\\copy", MatchAny, "TO", MatchAny))
+	/* Complete COPY|\copy <sth> FROM|TO with filename or STDIN/STDOUT/PROGRAM */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM|TO"))
+	{
+		/* COPY requires quoted filename */
+		bool force_quote = HeadMatches("COPY");
+
+		if (TailMatches("FROM"))
+			COMPLETE_WITH_FILES_PLUS("", force_quote, "STDIN", "PROGRAM");
+		else
+			COMPLETE_WITH_FILES_PLUS("", force_quote, "STDOUT", "PROGRAM");
+	}
+
+	/* Complete COPY|\copy <sth> FROM|TO PROGRAM command */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM|TO", "PROGRAM"))
+		COMPLETE_WITH_FILES("", HeadMatches("COPY"));	/* COPY requires quoted filename */
+
+	/* Complete COPY <sth> TO [PROGRAM] <sth> */
+	else if (Matches("COPY|\\copy", MatchAny, "TO", MatchAnyExcept("PROGRAM")) ||
+			 Matches("COPY|\\copy", MatchAny, "TO", "PROGRAM", MatchAny))
 		COMPLETE_WITH("WITH (");
 
-	/* Complete COPY <sth> FROM <sth> */
-	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAny))
+	/* Complete COPY <sth> FROM [PROGRAM] <sth> */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAnyExcept("PROGRAM")) ||
+			 Matches("COPY|\\copy", MatchAny, "FROM", "PROGRAM", MatchAny))
 		COMPLETE_WITH("WITH (", "WHERE");
 
-	/* Complete COPY <sth> FROM filename WITH ( */
-	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAny, "WITH", "("))
+	/* Complete COPY <sth> FROM [PROGRAM] filename WITH ( */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAnyExcept("PROGRAM"), "WITH", "(") ||
+			 Matches("COPY|\\copy", MatchAny, "FROM", "PROGRAM", MatchAny, "WITH", "("))
 		COMPLETE_WITH(Copy_from_options);
 
-	/* Complete COPY <sth> TO filename WITH ( */
-	else if (Matches("COPY|\\copy", MatchAny, "TO", MatchAny, "WITH", "("))
+	/* Complete COPY <sth> TO [PROGRAM] filename WITH ( */
+	else if (Matches("COPY|\\copy", MatchAny, "TO", MatchAnyExcept("PROGRAM"), "WITH", "(") ||
+			 Matches("COPY|\\copy", MatchAny, "TO", "PROGRAM", MatchAny, "WITH", "("))
 		COMPLETE_WITH(Copy_to_options);
 
-	/* Complete COPY <sth> FROM|TO filename WITH (FORMAT */
-	else if (Matches("COPY|\\copy", MatchAny, "FROM|TO", MatchAny, "WITH", "(", "FORMAT"))
+	/* Complete COPY <sth> FROM|TO [PROGRAM] <sth> WITH (FORMAT */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM|TO", MatchAnyExcept("PROGRAM"), "WITH", "(", "FORMAT") ||
+			 Matches("COPY|\\copy", MatchAny, "FROM|TO", "PROGRAM", MatchAny, "WITH", "(", "FORMAT"))
 		COMPLETE_WITH("binary", "csv", "text");
 
-	/* Complete COPY <sth> FROM filename WITH (ON_ERROR */
-	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAny, "WITH", "(", "ON_ERROR"))
+	/* Complete COPY <sth> FROM [PROGRAM] filename WITH (ON_ERROR */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAnyExcept("PROGRAM"), "WITH", "(", "ON_ERROR") ||
+			 Matches("COPY|\\copy", MatchAny, "FROM", "PROGRAM", MatchAny, "WITH", "(", "ON_ERROR"))
 		COMPLETE_WITH("stop", "ignore");
 
-	/* Complete COPY <sth> FROM filename WITH (LOG_VERBOSITY */
-	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAny, "WITH", "(", "LOG_VERBOSITY"))
+	/* Complete COPY <sth> FROM [PROGRAM] filename WITH (LOG_VERBOSITY */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAnyExcept("PROGRAM"), "WITH", "(", "LOG_VERBOSITY") ||
+			 Matches("COPY|\\copy", MatchAny, "FROM", "PROGRAM", MatchAny, "WITH", "(", "LOG_VERBOSITY"))
 		COMPLETE_WITH("silent", "default", "verbose");
 
-	/* Complete COPY <sth> FROM <sth> WITH (<options>) */
-	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAny, "WITH", MatchAny))
+	/* Complete COPY <sth> FROM [PROGRAM] <sth> WITH (<options>) */
+	else if (Matches("COPY|\\copy", MatchAny, "FROM", MatchAnyExcept("PROGRAM"), "WITH", MatchAny) ||
+			 Matches("COPY|\\copy", MatchAny, "FROM", "PROGRAM", MatchAny, "WITH", MatchAny))
 		COMPLETE_WITH("WHERE");
 
 	/* CREATE ACCESS METHOD */
@@ -6243,9 +6272,53 @@ complete_from_variables(const char *text, const char *prefix, const char *suffix
  *
  * Caller must also set completion_force_quote to indicate whether to force
  * quotes around the result.  (The SQL COPY command requires that.)
+ *
+ * If completion_charpp is set to a null-terminated array of literal keywords,
+ * these keywords will be included in the completion results alongside filenames,
+ * as long as they case-insensitively match the current input.
  */
 static char *
 complete_from_files(const char *text, int state)
+{
+	char 	   *result;
+	static int	list_index;
+	static bool	files_done;
+	const char *item;
+
+	/* Initialization */
+	if (state == 0)
+	{
+		list_index = 0;
+		files_done = false;
+	}
+
+	/* Return a filename that matches */
+	if (!files_done && (result = _complete_from_files(text, state)))
+		return result;
+	else if (!completion_charpp)
+		return NULL;
+	else
+		files_done = true;
+
+	/*
+	 * If there are no more matching files, check for hard-wired keywords.
+	 * These will only be returned if they match the input-so-far,
+	 * ignoring case.
+	 */
+	while ((item = completion_charpp[list_index++]))
+	{
+		if (pg_strncasecmp(text, item, strlen(text)) == 0)
+		{
+			completion_force_quote = false;
+			return pg_strdup_keyword_case(item, text);
+		}
+	}
+
+	return NULL;
+}
+
+static char *
+_complete_from_files(const char *text, int state)
 {
 #ifdef USE_FILENAME_QUOTING_FUNCTIONS
 
-- 
2.43.0

