Merge pull request #346 from Mytherin/connectionisbusy

Fix #344: correctly finish binary copy in Postgres result reading to reset the connection back to idle

Author: Mytherin

src/include/postgres_binary_reader.hpp

@@ -24,7 +24,6 @@ struct PostgresBinaryReader : public PostgresResultReader {
 
 protected:
 	bool Next();
-	void CheckResult();
 
 	void Reset();
 	bool Ready();

src/include/postgres_scanner.hpp

@@ -22,6 +22,9 @@ class PostgresTransaction;
 struct PostgresBindData : public FunctionData {
 	static constexpr const idx_t DEFAULT_PAGES_PER_TASK = 1000;
 
+public:
+	PostgresBindData(ClientContext &context);
+
 	PostgresVersion version;
 	string schema_name;
 	string table_name;

src/postgres_binary_reader.cpp

@@ -25,16 +25,16 @@ PostgresReadResult PostgresBinaryReader::Read(DataChunk &output) {
 		while (!Ready()) {
 			if (!Next()) {
 				// finished this batch
-				CheckResult();
 				return PostgresReadResult::FINISHED;
 			}
 		}
 
 		// read a row
 		auto tuple_count = ReadInteger<int16_t>();
-		if (tuple_count <= 0) { // done here, lets try to get more
+		if (tuple_count <= 0) {
+			// tuple_count of -1 signifies the file trailer (i.e. footer) - reset and skip
 			Reset();
-			return PostgresReadResult::FINISHED;
+			continue;
 		}
 
 		D_ASSERT(tuple_count == column_ids.size());
@@ -69,12 +69,17 @@ bool PostgresBinaryReader::Next() {
 
 	// len -1 signals end
 	if (len == -1) {
-		auto final_result = PQgetResult(con.GetConn());
-		if (!final_result || PQresultStatus(final_result) != PGRES_COMMAND_OK) {
-			PQclear(final_result);
-			throw IOException("Failed to fetch header for COPY: %s", string(PQresultErrorMessage(final_result)));
+		// consume all available results
+		while (true) {
+			PostgresResult pg_res(PQgetResult(con.GetConn()));
+			auto final_result = pg_res.res;
+			if (!final_result) {
+				break;
+			}
+			if (PQresultStatus(final_result) != PGRES_COMMAND_OK) {
+				throw IOException("Failed to fetch header for COPY: %s", string(PQresultErrorMessage(final_result)));
+			}
 		}
-		PQclear(final_result);
 		return false;
 	}
 
@@ -89,15 +94,6 @@ bool PostgresBinaryReader::Next() {
 	return true;
 }
 
-void PostgresBinaryReader::CheckResult() {
-	auto result = PQgetResult(con.GetConn());
-	if (!result || PQresultStatus(result) != PGRES_COMMAND_OK) {
-		PQclear(result);
-		throw std::runtime_error("Failed to execute COPY: " + string(PQresultErrorMessage(result)));
-	}
-	PQclear(result);
-}
-
 void PostgresBinaryReader::Reset() {
 	if (buffer) {
 		PQfreemem(buffer);

src/postgres_connection.cpp

@@ -107,16 +107,14 @@ vector<unique_ptr<PostgresResult>> PostgresConnection::ExecuteQueries(const stri
 		if (!res) {
 			break;
 		}
+		auto result = make_uniq<PostgresResult>(res);
 		if (ResultHasError(res)) {
-			PQclear(res);
 			throw std::runtime_error("Failed to execute query \"" + queries +
 			                         "\": " + string(PQresultErrorMessage(res)));
 		}
 		if (PQresultStatus(res) != PGRES_TUPLES_OK) {
-			PQclear(res);
 			continue;
 		}
-		auto result = make_uniq<PostgresResult>(res);
 		results.push_back(std::move(result));
 	}
 	return results;

src/postgres_copy_from.cpp

@@ -4,12 +4,11 @@
 namespace duckdb {
 
 void PostgresConnection::BeginCopyFrom(const string &query, ExecStatusType expected_result) {
-	auto result = PQExecute(query.c_str());
+	PostgresResult pg_res(PQExecute(query.c_str()));
+	auto result = pg_res.res;
 	if (!result || PQresultStatus(result) != expected_result) {
-		PQclear(result);
 		throw std::runtime_error("Failed to prepare COPY \"" + query + "\": " + string(PQresultErrorMessage(result)));
 	}
-	PQclear(result);
 }
 
 } // namespace duckdb

src/postgres_copy_to.cpp

@@ -56,12 +56,11 @@ void PostgresConnection::BeginCopyTo(ClientContext &context, PostgresCopyState &
 	}
 	query += ")";
 
-	auto result = PQExecute(query.c_str());
+	PostgresResult pg_res(PQExecute(query.c_str()));
+	auto result = pg_res.res;
 	if (!result || PQresultStatus(result) != PGRES_COPY_IN) {
-		PQclear(result);
 		throw std::runtime_error("Failed to prepare COPY \"" + query + "\": " + string(PQresultErrorMessage(result)));
 	}
-	PQclear(result);
 	if (state.format == PostgresCopyFormat::BINARY) {
 		// binary copy requires a header
 		PostgresBinaryWriter writer(state);
@@ -106,12 +105,12 @@ void PostgresConnection::FinishCopyTo(PostgresCopyState &state) {
 		throw InternalException("Error during PQputCopyEnd: %s", PQerrorMessage(GetConn()));
 	}
 	// fetch the query result to check for errors
-	auto result = PQgetResult(GetConn());
+	PostgresResult pg_res(PQgetResult(GetConn()));
+	auto result = pg_res.res;
 	if (!result || PQresultStatus(result) != PGRES_COMMAND_OK) {
-		PQclear(result);
-		throw std::runtime_error("Failed to copy data: " + string(PQresultErrorMessage(result)));
+		string error_msg(PQresultErrorMessage(result));
+		throw std::runtime_error("Failed to copy data: " + error_msg);
 	}
-	PQclear(result);
 }
 
 bool NeedsQuotes(const string &to_quote, idx_t size) {

src/postgres_query.cpp

@@ -11,7 +11,7 @@ namespace duckdb {
 
 static unique_ptr<FunctionData> PGQueryBind(ClientContext &context, TableFunctionBindInput &input,
                                             vector<LogicalType> &return_types, vector<string> &names) {
-	auto result = make_uniq<PostgresBindData>();
+	auto result = make_uniq<PostgresBindData>(context);
 
 	if (input.inputs[0].IsNull() || input.inputs[1].IsNull()) {
 		throw BinderException("Parameters to postgres_query cannot be NULL");

src/postgres_scanner.cpp

@@ -115,12 +115,9 @@ void PostgresScanFunction::PrepareBind(PostgresVersion version, ClientContext &c
 	if (context.TryGetCurrentSetting("pg_use_ctid_scan", pg_use_ctid_scan)) {
 		use_ctid_scan = BooleanValue::Get(pg_use_ctid_scan);
 	}
-	Value use_text_protocol;
-	if (context.TryGetCurrentSetting("pg_use_text_protocol", use_text_protocol)) {
-		if (BooleanValue::Get(use_text_protocol)) {
-			bind_data.use_text_protocol = true;
-			use_ctid_scan = false;
-		}
+	if (bind_data.use_text_protocol) {
+		// ctid scan is only supported for binary copy
+		use_ctid_scan = false;
 	}
 
 	if (version.major_v < 14) {
@@ -135,6 +132,15 @@ void PostgresScanFunction::PrepareBind(PostgresVersion version, ClientContext &c
 	bind_data.version = version;
 }
 
+PostgresBindData::PostgresBindData(ClientContext &context) {
+	Value text_protocol;
+	if (context.TryGetCurrentSetting("pg_use_text_protocol", text_protocol)) {
+		if (BooleanValue::Get(text_protocol)) {
+			use_text_protocol = true;
+		}
+	}
+}
+
 void PostgresBindData::SetTablePages(idx_t approx_num_pages) {
 	this->pages_approx = approx_num_pages;
 	if (!read_only || use_text_protocol) {
@@ -166,7 +172,7 @@ void PostgresBindData::SetTable(PostgresTableEntry &table) {
 
 static unique_ptr<FunctionData> PostgresBind(ClientContext &context, TableFunctionBindInput &input,
                                              vector<LogicalType> &return_types, vector<string> &names) {
-	auto bind_data = make_uniq<PostgresBindData>();
+	auto bind_data = make_uniq<PostgresBindData>(context);
 
 	bind_data->dsn = input.inputs[0].GetValue<string>();
 	bind_data->schema_name = input.inputs[1].GetValue<string>();

src/storage/postgres_table_entry.cpp

@@ -39,7 +39,7 @@ TableFunction PostgresTableEntry::GetScanFunction(ClientContext &context, unique
 	auto &pg_catalog = catalog.Cast<PostgresCatalog>();
 	auto &transaction = Transaction::Get(context, catalog).Cast<PostgresTransaction>();
 
-	auto result = make_uniq<PostgresBindData>();
+	auto result = make_uniq<PostgresBindData>(context);
 
 	result->schema_name = schema.name;
 	result->table_name = name;

test/sql/scanner/postgres_query_error.test

@@ -0,0 +1,33 @@
+# name: test/sql/scanner/postgres_query_error.test
+# description: Test running postgres_query
+# group: [scanner]
+
+require postgres_scanner
+
+require-env POSTGRES_TEST_DATABASE_AVAILABLE
+
+statement ok
+ATTACH 'dbname=postgresscanner' AS s1 (TYPE POSTGRES)
+
+statement ok
+BEGIN
+
+statement ok
+CALL pg_clear_cache();
+
+query III
+select * from postgres_query('s1', 'SELECT * FROM cars');
+----
+ferari	testarosa	red
+aston martin	db2	blue
+bentley	mulsanne	gray
+ford	T	black
+
+query III
+SELECT * FROM s1.cars
+----
+ferari	testarosa	red
+aston martin	db2	blue
+bentley	mulsanne	gray
+ford	T	black
+