Implement text reader protocol

Author: HuyNguyen7994

src/include/postgres_text_reader.hpp

@@ -0,0 +1,81 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// postgres_text_reader.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "duckdb.hpp"
+#include "duckdb/common/types.hpp"
+#include "duckdb/common/types/vector.hpp"
+#include "libpq-fe.h"
+#include "postgres_conversion.hpp"
+#include "postgres_connection.hpp"
+#include <cstring>
+
+namespace duckdb {
+
+struct PostgresTextReader {
+	explicit PostgresTextReader(PostgresConnection &con_p) : con(con_p), result(nullptr), col_vec(LogicalType::VARCHAR) {
+	}
+	
+	~PostgresTextReader() {
+		Reset();
+	}
+	
+	PostgresConnection &GetConn() {
+		return con;
+	}
+
+	void ReadTextFrom(const string &query) {
+		Reset();
+		result = PQexec(con.GetConn(), query.c_str());
+		if (!result || PQresultStatus(result) != PGRES_TUPLES_OK ) {
+			throw IOException("Failed to execute query: %s", string(PQerrorMessage(con.GetConn())));
+		}
+	}
+	
+	void Reset() {
+		if (result) {
+			PQclear(result);
+			result = nullptr;
+		}
+	}
+	
+	void ReadColumn(idx_t col_idx) {
+		col_vec.Resize(0, RowCount());
+		for (idx_t row_idx = 0; row_idx < RowCount(); row_idx++) {
+			if (PQgetisnull(result, row_idx, col_idx)) {
+				FlatVector::SetNull(col_vec, row_idx, true);
+				continue;
+			}
+			char *value = PQgetvalue(result, row_idx, col_idx);
+			int value_len = PQgetlength(result, row_idx, col_idx);
+			FlatVector::GetData<string_t>(col_vec)[row_idx] = StringVector::AddStringOrBlob(col_vec, value, value_len);
+		}
+	}
+	
+	void LoadResultTo (idx_t &col_idx, Vector &out_vec) {
+		ReadColumn(col_idx);
+		VectorOperations::DefaultCast(col_vec, out_vec, RowCount());
+	}
+
+	// Get result metadata
+	idx_t ColumnCount() {
+		return PQnfields(result);
+	}
+	
+	idx_t RowCount() {
+		return PQntuples(result);
+	}
+
+private:
+	PostgresConnection &con;
+	PGresult *result;
+	Vector col_vec;
+};
+
+} // namespace duckdb

src/postgres_extension.cpp

@@ -181,6 +181,8 @@ static void LoadInternal(DatabaseInstance &db) {
 	                          LogicalType::VARCHAR, Value(), SetPostgresNullByteReplacement);
 	config.AddExtensionOption("pg_debug_show_queries", "DEBUG SETTING: print all queries sent to Postgres to stdout",
 	                          LogicalType::BOOLEAN, Value::BOOLEAN(false), SetPostgresDebugQueryPrint);
+	config.AddExtensionOption("pg_use_legacy_text_protocol", "Whether or not to use legacy TEXT protocol to read data. Set this to yes will ignore pg_use_binary_copy",
+	                          LogicalType::BOOLEAN, Value::BOOLEAN(true));
 
 	OptimizerExtension postgres_optimizer;
 	postgres_optimizer.optimize_function = PostgresOptimizer::Optimize;

src/postgres_scanner.cpp

@@ -2,6 +2,7 @@
 
 #include <libpq-fe.h>
 
+#include "duckdb/common/unique_ptr.hpp"
 #include "duckdb/main/extension_util.hpp"
 #include "duckdb/common/shared_ptr.hpp"
 #include "duckdb/common/helper.hpp"
@@ -10,6 +11,7 @@
 #include "postgres_scanner.hpp"
 #include "postgres_result.hpp"
 #include "postgres_binary_reader.hpp"
+#include "postgres_text_reader.hpp"
 #include "storage/postgres_catalog.hpp"
 #include "storage/postgres_transaction.hpp"
 #include "storage/postgres_table_set.hpp"
@@ -34,6 +36,10 @@ struct PostgresLocalState : public LocalTableFunctionState {
 
 	void ScanChunk(ClientContext &context, const PostgresBindData &bind_data, PostgresGlobalState &gstate,
 	               DataChunk &output);
+	void ScanChunkWithBinaryReader(ClientContext &context, const PostgresBindData &bind_data, PostgresGlobalState &gstate,
+	               DataChunk &output);
+	void ScanChunkWithTextReader(ClientContext &context, const PostgresBindData &bind_data, PostgresGlobalState &gstate,
+	               DataChunk &output);
 };
 
 struct PostgresGlobalState : public GlobalTableFunctionState {
@@ -248,6 +254,18 @@ static void PostgresInitInternal(ClientContext &context, const PostgresBindData
 			filter += " AND ";
 		}
 		filter += filter_string;
+	};
+	lstate.exec = false;
+	lstate.done = false;
+	Value use_legacy_text_protocol;
+	if (context.TryGetCurrentSetting("pg_use_legacy_text_protocol", use_legacy_text_protocol)) {
+		if (BooleanValue::Get(use_legacy_text_protocol)) {
+			lstate.sql = StringUtil::Format(
+				R"(SELECT %s FROM %s.%s %s%s;)",
+				col_names, KeywordHelper::WriteQuoted(bind_data->schema_name, '"'),
+				KeywordHelper::WriteQuoted(bind_data->table_name, '"'), filter, bind_data->limit);
+			return;
+		}
 	}
 	if (bind_data->table_name.empty()) {
 		D_ASSERT(!bind_data->sql.empty());
@@ -261,8 +279,7 @@ static void PostgresInitInternal(ClientContext &context, const PostgresBindData
 		    col_names, KeywordHelper::WriteQuoted(bind_data->schema_name, '"'),
 		    KeywordHelper::WriteQuoted(bind_data->table_name, '"'), filter, bind_data->limit);
 	}
-	lstate.exec = false;
-	lstate.done = false;
+
 }
 
 static idx_t PostgresMaxThreads(ClientContext &context, const FunctionData *bind_data_p) {
@@ -417,7 +434,7 @@ static unique_ptr<LocalTableFunctionState> PostgresInitLocalState(ExecutionConte
 	return GetLocalState(context.client, input, gstate);
 }
 
-void PostgresLocalState::ScanChunk(ClientContext &context, const PostgresBindData &bind_data,
+void PostgresLocalState::ScanChunkWithBinaryReader(ClientContext &context, const PostgresBindData &bind_data,
                                    PostgresGlobalState &gstate, DataChunk &output) {
 	idx_t output_offset = 0;
 	PostgresBinaryReader reader(connection);
@@ -472,6 +489,41 @@ void PostgresLocalState::ScanChunk(ClientContext &context, const PostgresBindDat
 	}
 }
 
+void PostgresLocalState::ScanChunkWithTextReader(ClientContext &context, const PostgresBindData &bind_data,
+                                   PostgresGlobalState &gstate, DataChunk &output) {
+	PostgresTextReader reader(connection);
+	if (done && !PostgresParallelStateNext(context, &bind_data, *this, gstate)) {
+		return;
+	}
+
+	if (!exec) {
+		reader.ReadTextFrom(sql);
+		exec = true;
+	}
+
+	output.SetCardinality(reader.RowCount());
+	for (idx_t output_idx = 0; output_idx < output.ColumnCount(); output_idx++) {
+		auto col_idx = column_ids[output_idx];
+		auto &out_vec = output.data[output_idx];
+		reader.LoadResultTo(col_idx, out_vec);
+	}
+	reader.Reset();
+	done = true;
+	return;
+}
+
+void PostgresLocalState::ScanChunk(ClientContext &context, const PostgresBindData &bind_data,
+                                   PostgresGlobalState &gstate, DataChunk &output) {
+	Value use_legacy_text_protocol;
+	if (context.TryGetCurrentSetting("pg_use_legacy_text_protocol", use_legacy_text_protocol)) {
+		if (BooleanValue::Get(use_legacy_text_protocol)) {
+			ScanChunkWithTextReader(context, bind_data, gstate, output);
+			return;
+		}
+	}
+	ScanChunkWithBinaryReader(context, bind_data, gstate, output);
+}
+
 static void PostgresScan(ClientContext &context, TableFunctionInput &data, DataChunk &output) {
 	auto &bind_data = data.bind_data->Cast<PostgresBindData>();
 	auto &gstate = data.global_state->Cast<PostgresGlobalState>();