feat: implement postgres copy to stdout

[sunng87]

I hereby agree to the terms of the GreptimeDB CLA.

Refer to a related PR or issue link (optional)

#7706

What's changed and what's your intention?

This patch implements COPY (query) TO STDOUT WITH (format binary), which is required by ADBC-postgres for optimizations.

• This copy query will use postgrs PGCOPY protocol to send data to client
• This feature is implemented only for postgres interface
• The COPY TO STDOUT is parsed at protocol layer of postgres interface, to control the output encoding
• I shortcut our SQL parser to treat COPY (query) TO STDOUT just as query.
• Will switch to merged/released pgwire and arrow-pg when implementation confirmed.

PR Checklist

Please convert it to a draft if some of the following conditions are not met.

• I have written the necessary rustdoc comments.
• I have added the necessary unit tests and integration tests.
• This PR requires documentation updates.
• API changes are backward compatible.
• Schema or data changes are backward compatible.

[gemini-code-assist]

Summary of Changes

Hello @sunng87, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant feature by implementing the COPY (query) TO STDOUT command for the PostgreSQL interface. This functionality is vital for performance optimizations, particularly for ADBC-postgres, as it leverages the native PostgreSQL COPY protocol for efficient data export. The changes involve deep integration with custom pgwire and arrow-pg libraries to manage various output formats (binary, CSV, text) and a refined SQL parser that intelligently processes COPY TO STDOUT statements, treating them as direct queries for protocol-level handling.

Highlights

• PostgreSQL COPY TO STDOUT Implementation: Implemented support for the COPY (query) TO STDOUT command specifically for the PostgreSQL interface, enabling optimized data transfer using the PostgreSQL COPY protocol.
• Custom pgwire and arrow-pg Integration: Integrated custom versions of the pgwire and arrow-pg crates from specific git branches to handle binary, CSV, and text encoding for the COPY TO STDOUT functionality.
• SQL Parser Enhancements: Modified the SQL parser to treat COPY (query) TO STDOUT as its inner query, allowing the PostgreSQL protocol layer to manage the data encoding and format options directly.
• Generic Row Encoding: Refactored the RecordBatchRowIterator to be generic over different Encoder types, improving flexibility for various output formats like data rows or copy streams.
• Nullable Array Parameter Handling: Enhanced PostgreSQL parameter parsing to correctly handle nullable elements within array types (e.g., INT2_ARRAY, TIMESTAMP_ARRAY), ensuring NULL values are properly processed.

Changelog

• Cargo.lock
  • Updated various dependencies to newer versions.
  • Switched arrow-pg and pgwire dependencies to custom git branches to incorporate specific features for COPY encoding.
• Cargo.toml
  • Configured pgwire and arrow-pg to use specific git branches for custom feature development.
• src/servers/src/lib.rs
  • Added copy_to_stdout_format: Option<String> to the SqlPlan struct to store the detected format for COPY TO STDOUT statements.
• src/servers/src/mysql/handler.rs
  • Initialized the new copy_to_stdout_format field to None when creating SqlPlan instances, ensuring compatibility for MySQL handlers.
• src/servers/src/postgres/handler.rs
  • Implemented output_to_copy_response and recordbatches_to_copy_response functions to handle COPY TO STDOUT statements and encode data into specified formats (CSV, binary, text).
  • Modified SimpleQueryHandler and ExtendedQueryHandler to detect and route COPY TO STDOUT queries to the new copy response handling logic.
  • Added check_copy_to_stdout utility function to parse COPY statements and extract the output format.
  • Included unit tests for check_copy_to_stdout to verify format detection.
• src/servers/src/postgres/types.rs
  • Refactored RecordBatchRowIterator to be generic over an Encoder trait, allowing it to work with different encoding strategies (e.g., DataRowEncoder, CopyEncoder).
  • Updated encode_struct to accept a generic Encoder and FieldInfo for more flexible data encoding.
  • Modified parameter parsing for array types (e.g., INT2_ARRAY, TIMESTAMP_ARRAY) to correctly handle Option<T> values, supporting nullable elements in arrays.
  • Updated timestamp handling logic to use map and and_then for optional values, improving null safety.
• src/sql/src/parsers/copy_parser.rs
  • Modified parse_copy to identify COPY (query) TO STDOUT statements and return the inner query directly, bypassing the Statement::Copy wrapper for protocol-level handling.
  • Updated parse_copy_query_to to specifically check for the STDOUT keyword and construct a CopyQueryTo with location: 'STDOUT' without parsing additional options.
  • Added unit tests to validate the parsing behavior of COPY (query) TO STDOUT statements.

Activity

• The pull request introduces a new feature to implement PostgreSQL COPY TO STDOUT.
• The author has provided a detailed description of the changes and their intention, including the use of custom pgwire and arrow-pg branches.
• The PR checklist indicates that rustdoc comments, unit/integration tests, documentation updates, and API/schema backward compatibility checks are not yet completed, suggesting the PR is still under active development or requires further review and refinement.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request implements COPY (query) TO STDOUT for the PostgreSQL interface, which is a required optimization for ADBC-postgres. The approach of handling this at the protocol layer while treating it as a standard query in the engine is a good design choice. The code is generally well-structured, with good refactoring to support different encoders.

I've identified a critical issue in how multi-statement queries are handled, which could lead to incorrect output formatting. There's also an issue with the format code used in the CopyOutResponse, which violates the PostgreSQL protocol. Additionally, I've noted a medium-severity concern regarding the use of forked git dependencies, which could impact long-term maintainability. None of the comments contradict the provided rule regarding vector index searches and query optimization.

[gemini-code-assist]

Using git dependencies with specific branches can be risky for long-term maintainability and build reproducibility. Could you please add a note or link to the upstream PRs for these changes? It's important to ensure these changes are merged upstream and the dependencies are updated to released versions from crates.io before this feature is considered stable.