From df49bba655d94b850fc8b2a1e946e8efcd874d37 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 v6 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 | 125 ++++++++++++++++++++++++++-------
 1 file changed, 101 insertions(+), 24 deletions(-)

diff --git a/src/bin/psql/tab-complete.in.c b/src/bin/psql/tab-complete.in.c
index 6176741d20b..a8041401b8b 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(generator) \
 	matches = rl_completion_matches(text, generator)
 
@@ -1484,6 +1494,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);
@@ -3324,42 +3335,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 */
@@ -6241,6 +6270,54 @@ complete_from_variables(const char *text, const char *prefix, const char *suffix
 }
 
 
+/*
+ * This function wraps _complete_from_files() so that both literal keywords
+ * and filenames can be included in the completion results.
+ *
+ * If completion_charpp is set to a null-terminated array of literal keywords,
+ * those keywords are added to 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;
+}
+
 /*
  * This function wraps rl_filename_completion_function() to strip quotes from
  * the input before searching for matches and to quote any matches for which
@@ -6255,7 +6332,7 @@ complete_from_variables(const char *text, const char *prefix, const char *suffix
  * quotes around the result.  (The SQL COPY command requires that.)
  */
 static char *
-complete_from_files(const char *text, int state)
+_complete_from_files(const char *text, int state)
 {
 #ifdef USE_FILENAME_QUOTING_FUNCTIONS
 
-- 
2.43.0

