Fix hang on COPY ... FROM STDIN / TO STDOUT

connection.c

@@ -1808,6 +1808,46 @@ CC_internal_rollback(ConnectionClass *self, int rollback_type, BOOL ignore_abort
 	return ret;
 }
 
+/*
+ * Get the connection out of the COPY sub-protocol it was just put into, so
+ * that it stays usable after we reject the statement.  Without this the
+ * connection is stuck: libpq keeps reporting the same COPY status and the
+ * server keeps waiting on the stream.
+ */
+void
+CC_abort_copy(ConnectionClass *self)
+{
+	PGresult	*pgres;
+	PGcancel	*cancel;
+	char		errbuf[256];
+	char		*copybuf = NULL;
+
+	if (!self->pqconn)
+		return;
+
+	/*
+	 * Ends copy-in by failing the COPY, so no half-finished data is left
+	 * behind.  It reports an error instead if we're copying out, which is
+	 * the one case where we can't just say we're done: the server is
+	 * streaming at us, so cancel it and swallow whatever is in flight.
+	 */
+	if (PQputCopyEnd(self->pqconn, NULL) < 0)
+	{
+		/* Without the cancel we'd pull the whole table over the wire. */
+		if (NULL != (cancel = PQgetCancel(self->pqconn)))
+		{
+			PQcancel(cancel, errbuf, sizeof(errbuf));
+			PQfreeCancel(cancel);
+		}
+		while (PQgetCopyData(self->pqconn, &copybuf, 0) >= 0)
+			PQfreemem(copybuf);
+	}
+
+	/* Consume the results the aborted COPY leaves behind. */
+	while ((pgres = PQgetResult(self->pqconn)) != NULL)
+		PQclear(pgres);
+}
+
 /*
  *	The "result_in" is only used by QR_next_tuple() to fetch another group of rows into
  *	the same existing QResultClass (this occurs when the tuple cache is depleted and
@@ -2247,26 +2287,16 @@ MYLOG(DETAIL_LOG_LEVEL, "Discarded a RELEASE result\n");
 				}
 				break;
 			case PGRES_COPY_OUT:
-				/* XXX: We used to read from stdin here. Does that make any sense? */
 			case PGRES_COPY_IN:
-				if (query_completed)
-				{
-					QR_concat(res, QR_Constructor());
-					if (!QR_nextr(res))
-					{
-						CC_set_error(self, CONNECTION_COULD_NOT_RECEIVE, "Could not create result info in send_query.", func);
-						ReadyToReturn = TRUE;
-						retres = NULL;
-						break;
-					}
-					res = QR_nextr(res);
-					nrarg.res = res;
-				}
-				QR_set_rstatus(res, PORES_COPY_IN);
-				ReadyToReturn = TRUE;
-				retres = cmdres;
-				break;
 			case PGRES_COPY_BOTH:
+				CC_abort_copy(self);
+				CC_set_error(self, CONN_NOT_IMPLEMENTED_ERROR, "COPY ... FROM STDIN / TO STDOUT is not supported in ODBC", func);
+				QR_set_rstatus(res, PORES_FATAL_ERROR);
+
+				MYLOG(0, " error - %s\n", CC_get_errormsg(self));
+				ReadyToReturn = TRUE;
+				retres = NULL;
+				goto cleanup;
 			default:
 				/* skip the unexpected response if possible */
 				CC_set_error(self, CONNECTION_BACKEND_CRAZY, "Unexpected result status (send_query)", func);

connection.h

@@ -437,6 +437,7 @@ void		ProcessRollback(ConnectionClass *conn, BOOL undo, BOOL partial);
 const char	*CC_get_current_schema(ConnectionClass *conn);
 int             CC_mark_a_object_to_discard(ConnectionClass *conn, int type, const char *plan);
 int             CC_discard_marked_objects(ConnectionClass *conn);
+void		CC_abort_copy(ConnectionClass *self);
 
 int		CC_get_max_idlen(ConnectionClass *self);
 char	CC_get_escape(const ConnectionClass *self);

statement.c

@@ -2250,6 +2250,14 @@ MYLOG(DETAIL_LOG_LEVEL, "!!%p->miscinfo=%x res=%p\n", self, self->miscinfo, firs
 			 * SQL_SUCCESS_WITH_INFO;
 			 */
 		}
+		else if (CONN_NOT_IMPLEMENTED_ERROR == CC_get_errornumber(conn))
+		{
+			/*
+			 * Keep the "optional feature not implemented" SQLSTATE (HYC00).
+			 * Only the number: the message already lives on the connection.
+			 */
+			SC_set_errornumber(self, STMT_NOT_IMPLEMENTED_ERROR);
+		}
 		else
 		{
 			SC_set_error(self, STMT_EXEC_ERROR, CC_get_errormsg(conn), func);
@@ -2780,6 +2788,12 @@ MYLOG(DETAIL_LOG_LEVEL, "get_Result=%p %p\n", res, SC_get_Result(stmt));
 		case PGRES_COPY_OUT:
 		case PGRES_COPY_IN:
 		case PGRES_COPY_BOTH:
+			CC_abort_copy(conn);
+			QR_set_rstatus(res, PORES_FATAL_ERROR);
+			CC_set_error(conn, CONN_NOT_IMPLEMENTED_ERROR, "COPY ... FROM STDIN / TO STDOUT is not supported in ODBC", func);
+
+			QLOG(0, "PQexecXxxx error: - (%d) - %s\n", pgresstatus, CC_get_errormsg(conn));
+			break;
 		default:
 			/* skip the unexpected response if possible */
 			QR_set_rstatus(res, PORES_BAD_RESPONSE);
@@ -3220,6 +3234,17 @@ SC_set_errorinfo(StatementClass *self, QResultClass *res, int errkind)
 		return;
 	}
 
+	if (CONN_NOT_IMPLEMENTED_ERROR == CC_get_errornumber(conn))
+	{
+		/*
+		 * Keep the "optional feature not implemented" SQLSTATE (HYC00).
+		 * Only the number: the message already lives on the connection.
+		 */
+		if (0 == SC_get_errornumber(self))
+			SC_set_errornumber(self, STMT_NOT_IMPLEMENTED_ERROR);
+		return;
+	}
+
 	switch (QR_get_rstatus(res))
 	{
 		case PORES_NO_MEMORY_ERROR:

test/expected/commands.out

@@ -4,4 +4,20 @@ Testing VACUUM with SQLPrepare/SQLExecute...
 Disabling autocommit...
 Testing VACUUM with SQLExecDirect...
 Testing VACUUM with SQLPrepare/SQLExecute...
+Testing COPY FROM STDIN with SQLExecDirect...
+
+HYC00=Error while executing the query;
+COPY ... FROM STDIN / TO STDOUT is not supported in ODBC
+Testing COPY TO STDOUT with SQLExecDirect...
+
+HYC00=Error while executing the query;
+COPY ... FROM STDIN / TO STDOUT is not supported in ODBC
+Testing COPY FROM STDIN with SQLPrepare/SQLExecute...
+
+HYC00=Error while executing the query;
+COPY ... FROM STDIN / TO STDOUT is not supported in ODBC
+Testing that a column named stdin is not mistaken for COPY FROM STDIN...
+Testing that the connection still works after a rejected COPY...
+Result set:
+1
 disconnecting

test/src/commands-test.c

@@ -73,6 +73,75 @@ int main(int argc, char **argv)
 	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
 	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
 
+	/*
+	 * COPY ... FROM STDIN / TO STDOUT hand the connection over to the COPY
+	 * streaming sub-protocol, which we don't implement. Left alone, the
+	 * server waits forever for CopyData we never send while the driver
+	 * spins on the same COPY result. The driver should instead unwind the
+	 * copy, report it as unimplemented, and leave the connection usable --
+	 * on both the SQLExecDirect and SQLPrepare/SQLExecute paths.
+	 */
+	printf("Testing COPY FROM STDIN with SQLExecDirect...\n");
+	rc = SQLExecDirect(hstmt, (SQLCHAR *) "COPY testtab1 FROM STDIN", SQL_NTS);
+	/* Print error, it is expected */
+	if (!SQL_SUCCEEDED(rc))
+		print_diag("", SQL_HANDLE_STMT, hstmt);
+
+	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
+	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
+
+	printf("Testing COPY TO STDOUT with SQLExecDirect...\n");
+	rc = SQLExecDirect(hstmt, (SQLCHAR *) "COPY testtab1 TO STDOUT", SQL_NTS);
+	/* Print error, it is expected */
+	if (!SQL_SUCCEEDED(rc))
+		print_diag("", SQL_HANDLE_STMT, hstmt);
+
+	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
+	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
+
+	/*
+	 * Same with SQLPrepare/SQLExecute. The statement itself parses fine, so
+	 * the rejection only happens once the server reports copy-in at execute
+	 * time.
+	 */
+	printf("Testing COPY FROM STDIN with SQLPrepare/SQLExecute...\n");
+	rc = SQLPrepare(hstmt, (SQLCHAR *) "COPY testtab1 FROM STDIN", SQL_NTS);
+	CHECK_STMT_RESULT(rc, "SQLPrepare failed", hstmt);
+
+	rc = SQLExecute(hstmt);
+	/* Print error, it is expected */
+	if (!SQL_SUCCEEDED(rc))
+		print_diag("", SQL_HANDLE_STMT, hstmt);
+
+	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
+	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
+
+	/*
+	 * A COPY naming a column called "stdin" is a perfectly ordinary
+	 * file-based COPY and must not be caught by the check above.
+	 */
+	printf("Testing that a column named stdin is not mistaken for COPY FROM STDIN...\n");
+	rc = SQLExecDirect(hstmt, (SQLCHAR *) "CREATE TEMPORARY TABLE copy_stdin_col (stdin text, id int)", SQL_NTS);
+	CHECK_STMT_RESULT(rc, "SQLExecDirect failed", hstmt);
+
+	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
+	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
+
+	rc = SQLExecDirect(hstmt, (SQLCHAR *) "COPY copy_stdin_col (stdin, id) TO '/dev/null'", SQL_NTS);
+	CHECK_STMT_RESULT(rc, "SQLExecDirect failed", hstmt);
+
+	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
+	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
+
+	/* The connection must still be perfectly usable afterwards */
+	printf("Testing that the connection still works after a rejected COPY...\n");
+	rc = SQLExecDirect(hstmt, (SQLCHAR *) "SELECT 1", SQL_NTS);
+	CHECK_STMT_RESULT(rc, "SQLExecDirect failed", hstmt);
+	print_result(hstmt);
+
+	rc = SQLFreeStmt(hstmt, SQL_CLOSE);
+	CHECK_STMT_RESULT(rc, "SQLFreeStmt failed", hstmt);
+
 	/* Clean up */
 	test_disconnect();