mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-29 07:06:20 +00:00
feat: COPY FROM STDIN for postgres interface (#9297)
* feat: postgres copy-in support * feat: add postgres copy-in framework * feat(postgres): route COPY FROM STDIN flushes through the pending-rows batcher Follow-up to #9302: begin_copy_in now uses the batching-aware new_query_context, so copy-in batches are submitted to the pending-rows batcher like SQL INSERTs when PostgreSQL batching is enabled for the connection, and keep the direct write path when it is not. The batcher acknowledges synchronously, so the CommandComplete row count stays exact in both modes. Adds test_postgres_copy_in_with_batching asserting the batcher flush metric increments only when batching is enabled. Signed-off-by: Ning Sun <sunning@greptime.com> * refactor(postgres): open up the COPY FROM STDIN WITH clause for format-specific options The WITH (...) option list is now open instead of restricted to sqlparser's COPY grammar: - COPY ... FROM STDIN statements are parsed textually (new copy_in::parse module, tokenizer-based) so arbitrary KEY value entries reach the codecs; sqlparser would reject options it does not know before we ever see them. WITH is optional, matching PostgreSQL. - CopyOptionSet becomes an open ordered map of KEY -> String | Number | Ident | Boolean | Flag with typed take_* accessors (take_string, take_single_byte, take_flag, raw take_value) and case-insensitive keys; duplicate options are rejected. - Each format codec takes the options it understands and rejects the leftovers with one message naming them ('COPY option(s) X are not supported by format Y'), so a typo cannot be silently ignored. The text format now also rejects HEADER/QUOTE/ESCAPE instead of ignoring them, matching PostgreSQL. A future format (e.g. influxdb line protocol with WITH (FORMAT influxdb_line, precision 'ns')) only needs its codec and a registry entry; no parser changes. Signed-off-by: Ning Sun <sunning@greptime.com> * refactor(postgres): reuse arrow-csv decoder for COPY FROM STDIN csv parsing Replace the hand-rolled CSV record parser and per-value conversion with arrow-csv's incremental Decoder — the same engine backing COPY ... FROM 'file.csv' in common-datasource — so both COPY routes accept identical data. The decoder handles record splitting (quoting, escaped quotes, custom delimiters, CRLF, headers, chunk boundaries) and typed conversion against the targeted columns' schema in one pass; decoded batches become insert rows via the existing api::helper::vectors_to_rows and datatypes Helper::try_into_vectors, and the decode schema comes from the greptime Schema's arrow projection. The codec framework now has parsers emit complete insertable rows (api Row) instead of raw byte fields: create_parser takes the targeted columns and is fallible; text converts with the shared text-representation helpers, csv with the arrow decoder. Documented consequences of sharing the engine: JSON and Binary columns and bare epoch-integer timestamps are rejected for csv, and an empty field — quoted or not — is NULL; a custom NULL string is rejected accordingly. Signed-off-by: Ning Sun <sunning@greptime.com> * docs(postgres): make the copy-in framework explicitly format-open Reframe the codec framework docs so nothing assumes csv is the only other text format: the shared pieces (Field, convert_text_field, check_text_representable_type) are described as building blocks for any format whose records carry values as text, while formats with their own typed decoding (csv today) build rows themselves. State explicitly that formats are not limited to PostgreSQL's own — GreptimeDB-specific formats (e.g. a line protocol) register like any other and define their own options through the open WITH clause. No code changes beyond a test rename and a doc clarifier. Signed-off-by: Ning Sun <sunning@greptime.com> * test(postgres): driver-level coverage for CSV copy-in options Add an end-to-end test through tokio-postgres exercising HEADER, custom DELIMITER, quoted fields with embedded newlines, escaped quotes and CRLF record separators, with read-back verification through the real insert path. WITH is omitted, matching PostgreSQL. Signed-off-by: Ning Sun <sunning@greptime.com> * test(postgres): drop the raw-socket copy-in tests The simple-protocol raw client duplicated what tokio-postgres and the integration tests already cover; keep the driver-based tests only. Signed-off-by: Ning Sun <sunning@greptime.com> * refactor(postgres): address review: fix text-format O(n^2) drain, share type checks, hex::decode - text feed: take the longest complete line prefix in one drain and iterate lines over the slice, instead of draining per line (each drain memmoved the whole buffer; ~1300 drains for a 64KB chunk of short lines) - csv check_column_type: delegate to the shared check_text_representable_type and additionally reject Binary/Json - decode_bytea: use the hex crate instead of a hand-rolled decoder - semantic_type_of: use ColumnSchema::is_time_index() like the flow adapter's column_schemas_to_proto, instead of index lookups - parse: drop the unused expect_ident parameter; document the identifier-case decision - PostgresServer:🆕 document that query_handler and copy_in_handler are the same frontend instance Signed-off-by: Ning Sun <sunning@greptime.com> * chore: remove unrelated files * feat(postgres)!: drop the text codec; csv is the only COPY FROM STDIN format Remove the hand-rolled PostgreSQL text-format codec (backslash escapes, octal, \N nulls) that arrow-csv cannot express, keeping arrow-csv-based csv as the only implemented format. - COPY FROM STDIN without FORMAT now fails with guidance ('requires an explicit format in this version; use WITH (FORMAT csv)') rather than silently defaulting to the removed format - FORMAT text/txt and FORMAT binary report 'not supported yet' via registry stubs instead of unknown-format errors - codec.rs shrinks to the trait framework: the text-representation helpers (Field, convert_text_field, bytea hex, timestamp parsing) are removed with the codec; the module docs note where to find their shape should a future text-serializing format need them - csv implements check_column_type directly (no shared default) Breaking for clients relying on the PostgreSQL default text format; explicit WITH (FORMAT csv) is required. Signed-off-by: Ning Sun <sunning@greptime.com> * perf(postgres): raise COPY FROM STDIN flush batch to 32768 rows and harden flush diagnostics - COPY_IN_FLUSH_ROWS 8192 -> 32768: without the pending-rows batcher each flush pays the full insert pipeline (catalog lookup, schema check, partition split, region write round trip), so 4x larger batches cut round trips 4x at linearly-bounded memory; the constant also sizes the arrow decoder batch. Comment documents the trade-off and points high-rate users at the pending-rows batcher, which turns flushes into enqueues. - rows buffer pre-allocates COPY_IN_FLUSH_ROWS. - flush_copy_in warns on unexpected output variants instead of silently counting 0 rows. - batching wire test feeds 70000 rows to stay above the new threshold. Verified against a real psql 18.6 client: \copy ... WITH (FORMAT csv) sends FORMAT csv and succeeds; default \copy (text) gets the 'requires an explicit format' guidance; COPY TO STDOUT round-trips. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): reset copy-in state on successful CopyDone on_copy_done left the CopyInState in place (inert but holding buffers) until the next COPY overwrote it; take it like on_copy_fail does, so buffers are freed eagerly and a stray post-done CopyData gets the protocol-violation error instead of a finished parser. Signed-off-by: Ning Sun <sunning@greptime.com> * chore: bypass typo check * chore: format change * fix(postgres): multi-chunk COPY FROM STDIN corrupted at chunk boundaries The csv feed loop flushed the arrow-csv decoder at every CopyData chunk end — including partial batches at record boundaries — but arrow-csv's contract permits flush() only after decode() returned 0 (a full batch of complete records). Flushing outside that point corrupts subsequent decoding: fields mis-parse shifted by the size of the record prefix buffered so far ('000001,host-001,...' parses as '00001h,ost-001u,...'), either erroring (22P02/22P04) or, worse, silently writing wrong values. Which chunks trigger it depends on where record boundaries land, so it appears intermittent with random-length data (psql 8KB / pgx ~64KB chunking both hit it). Adopt the strict decoder protocol: feed decodes until input is exhausted and flushes only when decode returns 0 (batch full); trailing partial batches surface at the EOF signal in finish(). With batch_size == COPY_IN_FLUSH_ROWS the flush threshold is unchanged; sub-threshold streams now emit their rows at CopyDone instead of per chunk, which is semantically identical. Adds test_csv_chunked_feeds_are_alignment_safe: the same payload fed at chunk sizes 1..65531 (plus shifted starts) with every row's content verified — reproduces the corruption on the old loop at chunk sizes 7, 8, 13, 64 and 1000. Verified with the reporter's Go reproducer (pgx, content-checked): 10/10 clean at -chunk 8192, and 4096/1000/97/65531 plus the previously failing benchmark sizes 500/1000/2000/200/800 all round- trip exactly. Signed-off-by: Ning Sun <sunning@greptime.com> * chore: bypass spell check * chore: bypass spell check * fix(postgres): address PR review — cross-catalog guard, identifier folding, tab delimiter Three review findings (shuiyisong, PR #9297): [P1] Enforce the catalog restriction on COPY targets. copy_in_table now applies validate_catalog_and_schema when QueryOptions:: disallow_cross_catalog_query is set, mirroring SQL INSERT's check_permission path, BEFORE begin_copy_in switches the query context to the target catalog (which previously made the restriction unenforceable downstream). Cross-catalog COPY is now rejected with 'Illegal access to catalog' like INSERT. [P2] Canonicalize unquoted identifiers like the SQL parser (ParserContext::canonicalize_identifier): unquoted table-name parts and columns fold to lowercase, quoted ones keep their case. COPY Metrics (Host, Ts) now addresses the table CREATE TABLE Metrics created, and no longer silently misroutes into a coexisting quoted "Metrics". [P2] Allow a tab as CSV delimiter: take_single_byte accepts a literal tab while still rejecting newlines and other control characters, and the option parser now also accepts E'...' escape-string constants (DELIMITER E'\t'), processing PostgreSQL's backslash escapes. Verified end-to-end with psql 18.6: unquoted/quoted mixed-case round- trip into the right tables, cross-catalog rejection asserted in a new frontend unit test (QueryOptions plugin + second catalog), and both DELIMITER '<tab>' and DELIMITER E'\t' import tab-separated data. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): accept PG boolean spellings in COPY IN; drop copy state on failure Two review findings (shuiyisong, PR #9297): [P2] Boolean columns round-trip between the COPY directions. arrow-csv only parses true/false while PostgreSQL COPY OUT emits t/f, so boolean tables failed to re-import their own export (22P02). Boolean columns now decode as Utf8 in the arrow schema and flush_batch converts them with PostgreSQL's boolean input grammar (true/false, t/f, yes/no, y/n, on/off, 1/0, case-insensitive, whitespace trimmed). Round-trip covered end-to-end: COPY ... TO STDOUT's exact bytes are fed back via COPY FROM STDIN and verified. [P2] Drop the COPY state on every server-side failure. on_copy_data and on_copy_done now take ownership of the state up front, so parse, conversion or insert errors drop the buffered data instead of retaining it on long-lived or pooled connections until the next COPY or connection close (the reviewer measured ~17 MiB retained after a failed 8 MiB row); on_copy_data success puts the state back. An e2e test drives a failed COPY followed by a successful one on the same connection. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): interpret offset-free COPY timestamps in the session timezone Review finding (shuiyisong, PR #9297): the arrow decode schema carried Timestamp fields with no timezone, so the CSV decoder parsed offset-free timestamp strings as UTC while SQL INSERT parses them in the session timezone — the same input string silently stored a different instant depending on the ingestion method (8h apart in Asia/Shanghai). CopyInCodec::create_parser now receives the session timezone and the csv decode schema attaches it to Timestamp fields, so naive strings parse like INSERT; explicit offsets in the input (Z, +08:00) still take precedence, and UTC sessions are unchanged. Covered by test_csv_timestamp_session_timezone (naive vs +08:00 vs Z) and the end-to-end test_postgres_copy_in_session_timezone_parity, which SETs TIMEZONE TO 'Asia/Shanghai', inserts and COPY-loads the same string, and asserts both CAST(ts AS BIGINT) readbacks agree on 1704038400000. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): don't unescape E-string COPY options twice Review finding (shuiyisong, PR #9297): sqlparser's tokenizer already decodes PostgreSQL escape-string constants before returning Token::EscapedStringLiteral, so the second unescape_pg_string pass altered the decoded value — DELIMITER E'\\t' (a literal backslash followed by t, two characters) was turned into a tab and wrongly accepted as a single-byte delimiter. Use the EscapedStringLiteral payload as-is, like SingleQuotedString, and delete unescape_pg_string. Regression coverage distinguishes the three cases: E'\t' -> tab, E'\\' -> single backslash, and E'\\t' -> rejected as a multi-character delimiter. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): COPY OUT timestamps carry session-tz offset; quoted empty strings survive COPY IN Two review findings (shuiyisong, PR #9297): [P2] Preserve timestamp instants across COPY OUT/IN in non-UTC sessions. COPY OUT rendered naive UTC text while COPY IN (since the session-tz fix) interprets offset-free strings as local time, so re-importing an export shifted instants by the zone offset. COPY OUT now renders timestamps with an explicit session-timezone offset (2024-01-01 08:00:00.000000+08): the copy-out arrow schema carries the session tz (arrow-pg converts to a tz-aware DateTime) and the pg field type becomes TIMESTAMPTZ so the offset-bearing text format applies. Embedded offsets win on re-import regardless of session, keeping the OUT -> IN round trip exact; UTC sessions are unchanged. [P2] Preserve quoted empty strings. The shared arrow-csv decoder nulls every empty field, so a quoted empty string ("" — an empty STRING in PostgreSQL) imported as NULL. A shadow quote-mask scanner runs over the same bytes with the decoder's own splitting rules and records per field whether it was quoted; flush_batch restores quoted empties to empty strings for string columns, while unquoted empty fields stay NULL. Masks track records across chunk boundaries and align with decoder batches. Covered end-to-end: test_postgres_copy_out_in_timestamp_round_trip (Asia/Shanghai export asserted to carry +08, exact-epoch re-import comparison) and test_postgres_copy_in_quoted_empty_string (NULL, quoted empty and non-empty in one import), plus unit tests including a chunk-boundary sweep for the quoted-empty mask alignment. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): keep COPY IN quote masks aligned after doubled quotes; allow quoted empties in NOT NULL columns Two review findings (fengjiachun, PR #9297): [P2] Doubled quotes desynchronized the shadow quote-mask scanner. In QuoteInQuoted the doubled-quote branch skipped the second quote AND the following byte, so a closing quote or delimiter directly after an escaped quote was swallowed and later fields lost their masks -- import semantics depended on CopyData chunk boundaries. The branch now consumes only the escaped quote and returns to the in-quoted state, mirroring the decoder. Regression: "a""b","" ... verified across chunk sizes 1-64. [P2] A NOT NULL string column rejected quoted empty strings before restoration could run: the decode schema preserved the target nullability, and arrow-csv's batch validation rejects the temporary NULL a quoted empty decodes to. String fields now decode as nullable (safe: quoted-vs-unquoted is exactly the restoration's job), and the target NOT NULL constraint is enforced in flush_batch after restoration, so "" imports as an empty string while an unquoted empty field fails with a not-null-constraint error. Regressions at unit and e2e level: NOT NULL column accepting "" and rejecting unquoted NULL. Signed-off-by: Ning Sun <sunning@greptime.com> * fix(postgres): align COPY IN quote-mask scanner with the CSV decoder on blank lines and CR terminators Two review findings (fengjiachun, PR #9297): [P2] Blank lines emitted a spurious empty mask. csv-core skips blank lines, but the scanner queued one mask per blank record, so every subsequent decoded row consumed the wrong mask (a leading blank line turned a quoted empty into NULL and an unquoted empty into an empty string, and tripped NOT NULL on the wrong row). A blank line before the header also shifted the header expectation. end_record now returns no mask for a completely blank record without consuming the header expectation. [P2] Standalone CR was not a record terminator. The decoder accepts CR, CRLF and LF alike; the scanner only recognized CR before LF or at EOF, so CR-separated input desynchronized masks and flush_batch failed with 'quote masks out of sync'. is_record_end now treats CR alone as a terminator, consuming one byte (two for CRLF), while a CR at a chunk boundary still defers via carry so CRLF split across CopyData chunks stays one terminator. Regressions: blank lines before/inside records and before a header (NULL/empty preservation plus NOT NULL behavior), and mixed CR/CRLF/LF terminators across chunk sizes 1-64 with a trailing CR-only record. Signed-off-by: Ning Sun <sunning@greptime.com> * refactor(postgres)!: drop the quote-mask scanner; COPY IN CSV behaves exactly like file COPY The shadow quote-mask scanner — kept to preserve PostgreSQL's quoted-empty-string vs unquoted-NULL distinction — mirrored the decoder's record-splitting grammar by hand and required a fix for every rule it drifted on (doubled quotes, blank lines, CR terminators, NOT NULL interplay). Per maintainers' decision, that complexity is not worth the divergence: COPY FROM STDIN now behaves EXACTLY like the file-based COPY ... FROM 'file.csv' — an empty field, quoted or not, imports as NULL (arrow-csv's engine semantics, identical in both routes). Removed: QuoteMaskScanner and helpers, the always-nullable string decode schema, the restoration + NOT NULL enforcement in flush_batch, set_string_at, and the scanner-era unit tests. The decode schema again preserves target nullability, so the decoder enforces NOT NULL itself, same as the file route. Kept (small, schema-level, no second grammar): boolean spellings (t/f) and session-timezone timestamp decoding, both addressing separate review findings. E2E: test_postgres_copy_in_empty_field_is_null pins the parity — quoted empty, unquoted empty and non-empty in one import. Signed-off-by: Ning Sun <sunning@greptime.com> * chore: remove empty file * chore: lint issues --------- Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
Ning Sun
Generated
+1
@@ -15394,6 +15394,7 @@ dependencies = [
|
||||
"auth",
|
||||
"axum 0.8.4",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"cache",
|
||||
"catalog",
|
||||
"chrono",
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
pub mod builder;
|
||||
mod copy_in;
|
||||
mod dashboard;
|
||||
mod entity_graph;
|
||||
mod export_database;
|
||||
@@ -4699,4 +4700,69 @@ mod tests {
|
||||
assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// COPY FROM STDIN must enforce the same cross-catalog restriction as
|
||||
/// SQL INSERT when `disallow_cross_catalog_query` is configured.
|
||||
#[tokio::test]
|
||||
async fn test_copy_in_cross_catalog_restriction() -> TestResult<()> {
|
||||
use catalog::memory::MemoryCatalogManager;
|
||||
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
|
||||
use servers::query_handler::CopyInHandler;
|
||||
|
||||
let catalog = MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
|
||||
// A table in another catalog.
|
||||
catalog.register_catalog_sync("other").unwrap();
|
||||
catalog
|
||||
.register_schema_sync(catalog::RegisterSchemaRequest {
|
||||
catalog: "other".to_string(),
|
||||
schema: DEFAULT_SCHEMA_NAME.to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
catalog
|
||||
.register_table_sync(catalog::RegisterTableRequest {
|
||||
catalog: "other".to_string(),
|
||||
schema: DEFAULT_SCHEMA_NAME.to_string(),
|
||||
table_name: "copy_guard".to_string(),
|
||||
table_id: 2048,
|
||||
table: test_table(2048, "copy_guard")?,
|
||||
})
|
||||
.unwrap();
|
||||
let executor = Arc::new(MockProcedureExecutor::new(catalog.clone()));
|
||||
let plugins = Plugins::new();
|
||||
plugins.insert(QueryOptions {
|
||||
disallow_cross_catalog_query: true,
|
||||
});
|
||||
let instance = test_instance_with_catalog_manager(
|
||||
catalog,
|
||||
test_table(1025, "target")?,
|
||||
plugins,
|
||||
None,
|
||||
executor,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let ctx = Arc::new(QueryContext::with(
|
||||
DEFAULT_CATALOG_NAME,
|
||||
DEFAULT_SCHEMA_NAME,
|
||||
));
|
||||
|
||||
// Cross-catalog COPY is rejected before any data is streamed.
|
||||
let err = instance
|
||||
.copy_in_table("other", DEFAULT_SCHEMA_NAME, "copy_guard", ctx.clone())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
format!("{err:?}").contains("Illegal access to catalog"),
|
||||
"{err}"
|
||||
);
|
||||
|
||||
// Same-catalog COPY resolves normally.
|
||||
let found = instance
|
||||
.copy_in_table(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "target", ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(found.is_some());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use api::v1::RowInsertRequests;
|
||||
use async_trait::async_trait;
|
||||
use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
|
||||
use common_error::ext::BoxedError;
|
||||
use common_query::Output;
|
||||
use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
|
||||
use servers::error::{AuthSnafu, ExecuteGrpcRequestSnafu, ExecuteQuerySnafu, Result};
|
||||
use servers::query_handler::CopyInHandler;
|
||||
use session::context::QueryContextRef;
|
||||
use snafu::ResultExt;
|
||||
use table::metadata::TableInfoRef;
|
||||
|
||||
use crate::instance::Instance;
|
||||
|
||||
#[async_trait]
|
||||
impl CopyInHandler for Instance {
|
||||
async fn copy_in_table(
|
||||
&self,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Option<TableInfoRef>> {
|
||||
// Mirror SQL INSERT's catalog restriction (see `check_permission`):
|
||||
// when cross-catalog queries are disallowed, a COPY may only
|
||||
// target the current catalog. This must run against the ORIGINAL
|
||||
// query context — `begin_copy_in` switches the context to the
|
||||
// target catalog afterwards, so nothing downstream would catch
|
||||
// it.
|
||||
let disallow_cross_catalog_query = self
|
||||
.plugins
|
||||
.get::<QueryOptions>()
|
||||
.map(|opts| opts.disallow_cross_catalog_query)
|
||||
.unwrap_or_default();
|
||||
if disallow_cross_catalog_query {
|
||||
validate_catalog_and_schema(catalog, schema, &query_ctx)
|
||||
.map_err(BoxedError::new)
|
||||
.context(ExecuteQuerySnafu)?;
|
||||
}
|
||||
|
||||
// A COPY writes rows like a SQL INSERT; check the write permission
|
||||
// before the client starts streaming data.
|
||||
self.plugins
|
||||
.get::<PermissionCheckerRef>()
|
||||
.as_ref()
|
||||
.check_permission(
|
||||
query_ctx.current_user(),
|
||||
PermissionReq::BulkInsert {
|
||||
catalog,
|
||||
schema,
|
||||
table,
|
||||
},
|
||||
)
|
||||
.context(AuthSnafu)?;
|
||||
|
||||
self.catalog_manager()
|
||||
.table(catalog, schema, table, None)
|
||||
.await
|
||||
.map(|table| table.map(|table| table.table_info()))
|
||||
.map_err(servers::error::Error::from)
|
||||
}
|
||||
|
||||
async fn copy_in_insert(
|
||||
&self,
|
||||
requests: RowInsertRequests,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Output> {
|
||||
self.handle_row_inserts(requests, query_ctx, false, false)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(ExecuteGrpcRequestSnafu)
|
||||
}
|
||||
}
|
||||
@@ -420,6 +420,7 @@ where
|
||||
|
||||
let pg_server = Box::new(
|
||||
PostgresServer::new(
|
||||
instance.clone(),
|
||||
instance.clone(),
|
||||
opts.tls.should_force_tls(),
|
||||
tls_server_config,
|
||||
|
||||
@@ -320,6 +320,11 @@ lazy_static! {
|
||||
"servers postgres prepared count"
|
||||
)
|
||||
.unwrap();
|
||||
pub static ref METRIC_POSTGRES_COPY_IN_ROWS: IntCounter = register_int_counter!(
|
||||
"greptime_servers_postgres_copy_in_rows",
|
||||
"servers postgres COPY FROM STDIN rows written"
|
||||
)
|
||||
.unwrap();
|
||||
pub static ref METRIC_SERVER_GRPC_DB_REQUEST_TIMER: HistogramVec = register_histogram_vec!(
|
||||
"greptime_servers_grpc_db_request_elapsed",
|
||||
"servers grpc db request elapsed",
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
mod auth_handler;
|
||||
mod copy_in;
|
||||
mod fixtures;
|
||||
mod handler;
|
||||
mod server;
|
||||
@@ -42,6 +43,7 @@ use session::context::Channel;
|
||||
|
||||
use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator};
|
||||
use self::handler::DefaultQueryParser;
|
||||
use crate::query_handler::CopyInHandlerRef;
|
||||
use crate::query_handler::sql::ServerSqlQueryHandlerRef;
|
||||
|
||||
pub(crate) struct GreptimeDBStartupParameters {
|
||||
@@ -78,6 +80,7 @@ impl ServerParameterProvider for GreptimeDBStartupParameters {
|
||||
pub struct PostgresServerHandlerInner {
|
||||
query_handler: ServerSqlQueryHandlerRef,
|
||||
auth_source: Option<Arc<PgUserProviderSource>>,
|
||||
copy_in_handler: CopyInHandlerRef,
|
||||
force_tls: bool,
|
||||
auth_driver: AuthDriver<Arc<GreptimeDBStartupParameters>, UserInfoRef>,
|
||||
/// SCRAM selected; failures get the uniform metric label.
|
||||
@@ -86,11 +89,16 @@ pub struct PostgresServerHandlerInner {
|
||||
batching_enabled: bool,
|
||||
session: Arc<Session>,
|
||||
query_parser: Arc<DefaultQueryParser>,
|
||||
|
||||
/// State of an in-progress `COPY ... FROM STDIN`; set when the copy
|
||||
/// starts and taken when it finishes, fails or the connection drops.
|
||||
copy_in_state: tokio::sync::Mutex<Option<copy_in::CopyInState>>,
|
||||
}
|
||||
|
||||
#[derive(Builder)]
|
||||
pub(crate) struct MakePostgresServerHandler {
|
||||
query_handler: ServerSqlQueryHandlerRef,
|
||||
copy_in_handler: CopyInHandlerRef,
|
||||
user_provider: Option<UserProviderRef>,
|
||||
#[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
|
||||
param_provider: Arc<GreptimeDBStartupParameters>,
|
||||
@@ -112,6 +120,10 @@ impl PgWireServerHandlers for PostgresServerHandler {
|
||||
self.0.clone()
|
||||
}
|
||||
|
||||
fn copy_handler(&self) -> Arc<impl pgwire::api::copy::CopyHandler> {
|
||||
self.0.clone()
|
||||
}
|
||||
|
||||
fn error_handler(&self) -> Arc<impl ErrorHandler> {
|
||||
self.0.clone()
|
||||
}
|
||||
@@ -136,6 +148,7 @@ impl MakePostgresServerHandler {
|
||||
.user_provider
|
||||
.clone()
|
||||
.map(|p| Arc::new(PgUserProviderSource::new(p))),
|
||||
copy_in_handler: self.copy_in_handler.clone(),
|
||||
force_tls: self.force_tls,
|
||||
auth_driver: AuthDriver::new(self.param_provider.clone())
|
||||
.with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)),
|
||||
@@ -144,6 +157,7 @@ impl MakePostgresServerHandler {
|
||||
batching_enabled,
|
||||
session: session.clone(),
|
||||
query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)),
|
||||
copy_in_state: tokio::sync::Mutex::new(None),
|
||||
};
|
||||
PostgresServerHandler(Arc::new(handler))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//! The COPY FROM STDIN codec framework.
|
||||
//!
|
||||
//! A *codec* binds one COPY data format to its implementation. Codecs are
|
||||
//! stateless strategies: they validate the options they accept at
|
||||
//! statement-parse time, describe which column types they can ingest, and
|
||||
//! create the stateful [`RecordParser`] per stream.
|
||||
//!
|
||||
//! Formats are not limited to PostgreSQL's own — GreptimeDB-specific
|
||||
//! formats (a line protocol, a JSON-lines dialect, ...) register exactly
|
||||
//! like `csv` does, and may define their own options through the open
|
||||
//! `WITH (...)` clause.
|
||||
//!
|
||||
//! # Parsing and conversion
|
||||
//!
|
||||
//! A [`RecordParser`] consumes raw `CopyData` bytes and emits complete,
|
||||
//! insertable rows: one [`ParsedRow`] per record, in the order of the
|
||||
//! targeted columns. Parsers validate record shape (field counts) and
|
||||
//! convert field values themselves — `csv` delegates both splitting and
|
||||
//! typed conversion to arrow-csv's incremental decoder (the engine
|
||||
//! behind `COPY ... FROM 'file.csv'`) and turns decoded batches into
|
||||
//! rows with the shared `api::helper::vectors_to_rows`. A future format
|
||||
//! that carries values as text (e.g. a line protocol) would reintroduce
|
||||
//! shared field-conversion helpers here; see the removed `text` codec in
|
||||
//! git history for the shape they had.
|
||||
//!
|
||||
//! # Options
|
||||
//!
|
||||
//! The `WITH (...)` clause is open: every `KEY value` entry the client
|
||||
//! writes is collected into a [`CopyOptionSet`] (see `options`) regardless
|
||||
//! of whether any format understands it — the statement is parsed
|
||||
//! textually (see `parse`) precisely so options beyond sqlparser's COPY
|
||||
//! grammar can reach the codecs. A codec builder takes the options it
|
||||
//! understands (`take_string`, `take_single_byte`, `take_flag`, or raw
|
||||
//! `take_value` for codec-specific shapes) and rejects the rest with
|
||||
//! `reject_leftovers`, so a future format like
|
||||
//! `COPY t FROM STDIN WITH (FORMAT influxdb_line, precision 'ns')` only
|
||||
//! needs its own codec and registry entry.
|
||||
//!
|
||||
//! To add a format:
|
||||
//!
|
||||
//! 1. implement [`CopyInCodec`] for a config struct holding the format's
|
||||
//! validated options (see `csv::CsvCodec` for the example),
|
||||
//! 2. write a builder `fn(CopyOptionSet) -> PgWireResult<Arc<dyn
|
||||
//! CopyInCodec>>` that applies defaults and rejects options the format
|
||||
//! does not accept,
|
||||
//! 3. register the builder under every accepted `FORMAT` name (including
|
||||
//! aliases) in `REGISTERED_FORMATS` in the parent module.
|
||||
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use pgwire::error::PgWireResult;
|
||||
|
||||
/// A parsed, insertable row: one gRPC value per targeted column, in COPY
|
||||
/// field order. Every parser emits these directly, whatever mechanism it
|
||||
/// used to produce them.
|
||||
pub(super) type ParsedRow = api::v1::Row;
|
||||
|
||||
/// Incremental parser of one COPY FROM STDIN stream.
|
||||
///
|
||||
/// The wire protocol delivers the payload as an arbitrary sequence of
|
||||
/// `CopyData` chunks; implementations must keep state across [`feed`]
|
||||
/// calls because records (and even single fields) may span chunk
|
||||
/// boundaries. [`finish`] is called at `CopyDone`; a final record without
|
||||
/// a trailing record separator is accepted.
|
||||
pub(super) trait RecordParser: Send {
|
||||
/// Feeds a chunk of `CopyData` bytes and returns all complete rows.
|
||||
fn feed(&mut self, data: &[u8]) -> PgWireResult<Vec<ParsedRow>>;
|
||||
/// Parses the buffered tail at `CopyDone`; returns no row when the
|
||||
/// stream ended exactly at a record boundary.
|
||||
fn finish(&mut self) -> PgWireResult<Vec<ParsedRow>>;
|
||||
}
|
||||
|
||||
/// A COPY FROM STDIN data format implementation.
|
||||
///
|
||||
/// Implementations are shared per statement (hence `Send + Sync`) and must
|
||||
/// be stateless; per-stream state belongs to the [`RecordParser`] they
|
||||
/// create. Every method is tailored to the format: `check_column_type`
|
||||
/// declares the types the format can decode (csv is bounded by what
|
||||
/// arrow-csv supports).
|
||||
pub(super) trait CopyInCodec: Send + Sync + std::fmt::Debug {
|
||||
/// The canonical format name, used in error messages.
|
||||
fn name(&self) -> &'static str;
|
||||
/// The format code advertised in the `CopyInResponse` header: 0 when
|
||||
/// the client sends textual data (the default), 1 for binary.
|
||||
fn format_code(&self) -> i8 {
|
||||
0
|
||||
}
|
||||
/// Validates that a column of `data_type` can be ingested by this
|
||||
/// format.
|
||||
fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()>;
|
||||
/// Creates the stateful parser for one copy stream over the targeted
|
||||
/// columns (in COPY field order).
|
||||
///
|
||||
/// `timezone` is the session timezone (an IANA name or a UTC offset):
|
||||
/// formats that parse offset-free timestamp strings from text should
|
||||
/// interpret them in this zone, matching SQL INSERT.
|
||||
fn create_parser(
|
||||
&self,
|
||||
columns: &[ColumnSchema],
|
||||
timezone: &str,
|
||||
) -> PgWireResult<Box<dyn RecordParser + Send>>;
|
||||
}
|
||||
@@ -0,0 +1,703 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//! The `csv` format — currently the only implemented COPY FROM STDIN
|
||||
//! format.
|
||||
//!
|
||||
//! Parsing and typed conversion both reuse arrow-csv's incremental
|
||||
//! [`Decoder`](arrow::csv::reader::Decoder) — the same engine backing
|
||||
//! GreptimeDB's file-based `COPY ... FROM 'file.csv'` (see
|
||||
//! `common-datasource`'s `csv` file format) — so both COPY routes accept
|
||||
//! identical data. The decoder splits records (quoting, escaped quotes,
|
||||
//! custom delimiters, `\n`/`\r\n` separators, optional headers) and
|
||||
//! converts fields against the targeted columns' schema in one pass.
|
||||
//!
|
||||
//! Consequences of sharing the engine: JSON and Binary columns are
|
||||
//! rejected, timestamps must be date/time strings (offset-free ones
|
||||
//! parse in the session timezone), and an empty field — quoted or not —
|
||||
//! is NULL, exactly like file-based COPY.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::helper::vectors_to_rows;
|
||||
use arrow::csv::ReaderBuilder;
|
||||
use arrow::csv::reader::Decoder as CsvDecoder;
|
||||
use arrow_schema::{ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema};
|
||||
use datatypes::data_type::DataType;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::schema::{ColumnSchema, Schema as GreptimeSchema};
|
||||
use datatypes::vectors::{BooleanVector, Helper, StringVector, Vector};
|
||||
use pgwire::error::{PgWireError, PgWireResult};
|
||||
|
||||
use crate::postgres::copy_in::codec::{CopyInCodec, ParsedRow, RecordParser};
|
||||
use crate::postgres::copy_in::options::CopyOptionSet;
|
||||
use crate::postgres::copy_in::{COPY_IN_FLUSH_ROWS, bad_copy_data, invalid_text_repr, unsupported};
|
||||
|
||||
/// The validated `csv` format configuration.
|
||||
#[derive(Debug)]
|
||||
pub(super) struct CsvCodec {
|
||||
delimiter: u8,
|
||||
quote: u8,
|
||||
escape: u8,
|
||||
header: bool,
|
||||
}
|
||||
|
||||
/// Builds the CSV codec from the parsed options.
|
||||
pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult<Arc<dyn CopyInCodec>> {
|
||||
let delimiter = options.take_single_byte("DELIMITER")?.unwrap_or(b',');
|
||||
let quote = options.take_single_byte("QUOTE")?.unwrap_or(b'"');
|
||||
// In PostgreSQL the default escape character for CSV is the quote
|
||||
// character itself (doubled quotes inside quoted fields).
|
||||
let escape = options.take_single_byte("ESCAPE")?.unwrap_or(quote);
|
||||
// An empty string is the default NULL marker (an empty field is NULL);
|
||||
// a custom NULL string is not supported because arrow-csv's null
|
||||
// matching does not distinguish quoted from unquoted fields.
|
||||
if let Some(null) = options.take_string("NULL")?
|
||||
&& !null.is_empty()
|
||||
{
|
||||
return Err(unsupported(
|
||||
"custom NULL string is not supported by format csv; an empty field is NULL",
|
||||
));
|
||||
}
|
||||
let header = options.take_flag("HEADER")?.unwrap_or(false);
|
||||
options.reject_leftovers("csv")?;
|
||||
if delimiter == quote {
|
||||
return Err(bad_copy_data("DELIMITER must not be equal to QUOTE"));
|
||||
}
|
||||
Ok(Arc::new(CsvCodec {
|
||||
delimiter,
|
||||
quote,
|
||||
escape,
|
||||
header,
|
||||
}))
|
||||
}
|
||||
|
||||
impl CopyInCodec for CsvCodec {
|
||||
fn name(&self) -> &'static str {
|
||||
"csv"
|
||||
}
|
||||
|
||||
/// The column types arrow-csv can decode against.
|
||||
fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> {
|
||||
let supported = matches!(
|
||||
data_type,
|
||||
ConcreteDataType::Boolean(_)
|
||||
| ConcreteDataType::Int8(_)
|
||||
| ConcreteDataType::Int16(_)
|
||||
| ConcreteDataType::Int32(_)
|
||||
| ConcreteDataType::Int64(_)
|
||||
| ConcreteDataType::UInt8(_)
|
||||
| ConcreteDataType::UInt16(_)
|
||||
| ConcreteDataType::UInt32(_)
|
||||
| ConcreteDataType::UInt64(_)
|
||||
| ConcreteDataType::Float32(_)
|
||||
| ConcreteDataType::Float64(_)
|
||||
| ConcreteDataType::String(_)
|
||||
| ConcreteDataType::Date(_)
|
||||
| ConcreteDataType::Timestamp(_)
|
||||
| ConcreteDataType::Decimal128(_)
|
||||
);
|
||||
if !supported {
|
||||
return Err(unsupported(format!(
|
||||
"column type \"{}\" is not supported by format csv",
|
||||
data_type.name()
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_parser(
|
||||
&self,
|
||||
columns: &[ColumnSchema],
|
||||
timezone: &str,
|
||||
) -> PgWireResult<Box<dyn RecordParser + Send>> {
|
||||
// Decode schema: the greptime schema's arrow projection, except
|
||||
// Boolean columns decode as Utf8 — arrow-csv only accepts
|
||||
// `true`/`false`, while PostgreSQL COPY IN accepts its full set
|
||||
// of boolean spellings (`t`, `f`, `yes`, `on`, ...), which COPY
|
||||
// OUT emits — and Timestamp fields carry the SESSION timezone so
|
||||
// offset-free strings parse like SQL INSERT does instead of as
|
||||
// UTC. The substitutions are applied in `flush_batch` and by the
|
||||
// decoder's tz-aware parsing respectively; explicit offsets in
|
||||
// the input always win.
|
||||
let boolean_columns: Vec<usize> = columns
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, c)| matches!(c.data_type, ConcreteDataType::Boolean(_)))
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
let arrow_schema = GreptimeSchema::try_new(columns.to_vec())
|
||||
.map_err(|e| unsupported(format!("invalid columns for format csv: {e}")))?
|
||||
.arrow_schema()
|
||||
.clone();
|
||||
let timezone = timezone.to_string();
|
||||
let fields = arrow_schema
|
||||
.fields()
|
||||
.iter()
|
||||
.zip(columns)
|
||||
.map(|(field, column)| match &column.data_type {
|
||||
ConcreteDataType::Boolean(_) => {
|
||||
Field::new(field.name(), ArrowDataType::Utf8, field.is_nullable())
|
||||
}
|
||||
ConcreteDataType::Timestamp(_) => match field.data_type() {
|
||||
ArrowDataType::Timestamp(unit, None) => Field::new(
|
||||
field.name(),
|
||||
ArrowDataType::Timestamp(*unit, Some(Arc::from(timezone.as_str()))),
|
||||
field.is_nullable(),
|
||||
),
|
||||
_ => field.as_ref().clone(),
|
||||
},
|
||||
_ => field.as_ref().clone(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let mut builder = ReaderBuilder::new(Arc::new(ArrowSchema::new(fields)))
|
||||
.with_header(self.header)
|
||||
.with_delimiter(self.delimiter)
|
||||
.with_quote(self.quote)
|
||||
.with_batch_size(COPY_IN_FLUSH_ROWS);
|
||||
// An escape character equal to the quote is the doubled-quote
|
||||
// default; csv-core handles that through quoting alone.
|
||||
if self.escape != self.quote {
|
||||
builder = builder.with_escape(self.escape);
|
||||
}
|
||||
Ok(Box::new(CsvRecordParser {
|
||||
decoder: builder.build_decoder(),
|
||||
boolean_columns,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
/// Incremental parser backed by arrow-csv's decoder.
|
||||
struct CsvRecordParser {
|
||||
decoder: CsvDecoder,
|
||||
/// Indices of Boolean columns, decoded as Utf8 and converted by
|
||||
/// [`flush_batch`] to accept PostgreSQL's boolean spellings.
|
||||
boolean_columns: Vec<usize>,
|
||||
}
|
||||
|
||||
impl RecordParser for CsvRecordParser {
|
||||
fn feed(&mut self, data: &[u8]) -> PgWireResult<Vec<ParsedRow>> {
|
||||
let mut buf = data;
|
||||
let mut rows = Vec::new();
|
||||
while !buf.is_empty() {
|
||||
let consumed = self.decoder.decode(buf).map_err(decode_error)?;
|
||||
if consumed == 0 {
|
||||
while let Some(batch) = self.decoder.flush().map_err(decode_error)? {
|
||||
rows.extend(self.flush_batch(&batch)?);
|
||||
}
|
||||
} else {
|
||||
buf = &buf[consumed..];
|
||||
}
|
||||
}
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> PgWireResult<Vec<ParsedRow>> {
|
||||
// An empty decode signals EOF, finalizing a trailing record that
|
||||
// lacks its record separator.
|
||||
self.decoder.decode(&[]).map_err(decode_error)?;
|
||||
let mut rows = Vec::new();
|
||||
while let Some(batch) = self.decoder.flush().map_err(decode_error)? {
|
||||
rows.extend(self.flush_batch(&batch)?);
|
||||
}
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
impl CsvRecordParser {
|
||||
fn flush_batch(
|
||||
&mut self,
|
||||
batch: &arrow::record_batch::RecordBatch,
|
||||
) -> PgWireResult<Vec<ParsedRow>> {
|
||||
let mut vectors = Helper::try_into_vectors(batch.columns())
|
||||
.map_err(|e| bad_copy_data(format!("failed to convert decoded CSV batch: {e}")))?;
|
||||
for &index in &self.boolean_columns {
|
||||
let Some(vector) = vectors[index].as_any().downcast_ref::<StringVector>() else {
|
||||
return Err(bad_copy_data(
|
||||
"expected boolean column to decode as a string vector",
|
||||
));
|
||||
};
|
||||
let mut values = Vec::with_capacity(batch.num_rows());
|
||||
for row_index in 0..batch.num_rows() {
|
||||
values.push(match vector.get(row_index) {
|
||||
datatypes::value::Value::Null => None,
|
||||
datatypes::value::Value::String(s) => Some(parse_pg_bool(s.as_utf8())?),
|
||||
other => {
|
||||
return Err(bad_copy_data(format!(
|
||||
"expected string values in a boolean column, got {other:?}"
|
||||
)));
|
||||
}
|
||||
});
|
||||
}
|
||||
vectors[index] = Arc::new(BooleanVector::from(values));
|
||||
}
|
||||
Ok(vectors_to_rows(vectors.iter(), batch.num_rows()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Parses PostgreSQL's boolean spellings: `true`/`false`, `t`/`f`,
|
||||
/// `yes`/`no`, `on`/`off` and `1`/`0`, case-insensitive with surrounding
|
||||
/// whitespace allowed (see the BOOLEAN type's input grammar). COPY OUT
|
||||
/// emits `t`/`f`, so accepting them keeps both COPY directions
|
||||
/// round-trippable.
|
||||
fn parse_pg_bool(input: &str) -> PgWireResult<bool> {
|
||||
match input.trim().to_ascii_lowercase().as_str() {
|
||||
"true" | "t" | "yes" | "y" | "on" | "1" => Ok(true),
|
||||
"false" | "f" | "no" | "n" | "off" | "0" => Ok(false),
|
||||
_ => Err(invalid_text_repr(format!(
|
||||
"invalid input syntax for type boolean: \"{input}\""
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_error(error: ArrowError) -> PgWireError {
|
||||
match &error {
|
||||
ArrowError::CsvError(message) => {
|
||||
bad_copy_data(format!("failed to parse CSV data: {message}"))
|
||||
}
|
||||
ArrowError::ParseError(message) | ArrowError::CastError(message) => {
|
||||
invalid_text_repr(format!("failed to parse CSV data: {message}"))
|
||||
}
|
||||
_ => bad_copy_data(format!("failed to parse CSV data: {error}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use datatypes::schema::ColumnSchema;
|
||||
|
||||
use super::*;
|
||||
use crate::postgres::copy_in::options::CopyOptionValue;
|
||||
|
||||
fn test_columns() -> Vec<ColumnSchema> {
|
||||
vec![
|
||||
ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
|
||||
ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true),
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
fn parser() -> Box<dyn RecordParser + Send> {
|
||||
build_csv_codec(CopyOptionSet::default())
|
||||
.unwrap()
|
||||
.create_parser(&test_columns(), "UTC")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn string_at(row: &ParsedRow, index: usize) -> Option<String> {
|
||||
row.values[index].value_data.as_ref().map(|v| match v {
|
||||
api::v1::value::ValueData::StringValue(s) => s.clone(),
|
||||
other => panic!("expected string, got {other:?}"),
|
||||
})
|
||||
}
|
||||
|
||||
fn rows_with(mut set: CopyOptionSet) -> Box<dyn RecordParser + Send> {
|
||||
set.take_format().unwrap();
|
||||
build_csv_codec(set)
|
||||
.unwrap()
|
||||
.create_parser(&test_columns(), "UTC")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Feeds chunks and then the EOF signal, returning all rows: with the
|
||||
/// strict decoder protocol, feeds only emit full
|
||||
/// `COPY_IN_FLUSH_ROWS` batches; a trailing partial batch surfaces at
|
||||
/// `finish`.
|
||||
fn feed_all(parser: &mut Box<dyn RecordParser + Send>, chunks: &[&[u8]]) -> Vec<ParsedRow> {
|
||||
let mut rows = Vec::new();
|
||||
for chunk in chunks {
|
||||
rows.extend(parser.feed(chunk).unwrap());
|
||||
}
|
||||
rows.extend(parser.finish().unwrap());
|
||||
rows
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_defaults() {
|
||||
let codec = build_csv_codec(CopyOptionSet::default()).unwrap();
|
||||
assert_eq!(codec.name(), "csv");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_rejects_bad_options() {
|
||||
let mut set = CopyOptionSet::default();
|
||||
set.insert("delimiter".into(), CopyOptionValue::String("\"".into()))
|
||||
.unwrap();
|
||||
let err = build_csv_codec(set).unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("DELIMITER must not be equal to QUOTE")
|
||||
);
|
||||
|
||||
let mut set = CopyOptionSet::default();
|
||||
set.insert("null".into(), CopyOptionValue::String("NULL".into()))
|
||||
.unwrap();
|
||||
assert!(
|
||||
build_csv_codec(set)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("custom NULL string")
|
||||
);
|
||||
|
||||
let mut set = CopyOptionSet::default();
|
||||
set.take_format().unwrap();
|
||||
set.insert("precision".into(), CopyOptionValue::Ident("ns".into()))
|
||||
.unwrap();
|
||||
assert!(
|
||||
build_csv_codec(set)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("not supported by format csv")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_check_column_type() {
|
||||
let codec = build_csv_codec(CopyOptionSet::default()).unwrap();
|
||||
assert!(
|
||||
codec
|
||||
.check_column_type(&ConcreteDataType::json_datatype())
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
codec
|
||||
.check_column_type(&ConcreteDataType::binary_datatype())
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
codec
|
||||
.check_column_type(&ConcreteDataType::string_datatype())
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_parse_records() {
|
||||
let mut parser = parser();
|
||||
let rows = feed_all(
|
||||
&mut parser,
|
||||
&[b"host1,1.5,2023-11-14 22:13:20\n\
|
||||
,2.5,2023-11-14 22:13:21\n\
|
||||
\"quoted, comma\",3.5,2023-11-14 22:13:22\n"],
|
||||
);
|
||||
assert_eq!(rows.len(), 3);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
|
||||
// An empty field is NULL.
|
||||
assert_eq!(string_at(&rows[1], 0), None);
|
||||
assert_eq!(string_at(&rows[2], 0), Some("quoted, comma".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_escaped_quotes() {
|
||||
let mut parser = parser();
|
||||
let rows = feed_all(
|
||||
&mut parser,
|
||||
&[b"\"say \"\"hi\"\"\",1.0,2023-11-14 22:13:20\n"],
|
||||
);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("say \"hi\"".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_custom_escape() {
|
||||
let mut set = CopyOptionSet::default();
|
||||
set.insert("escape".into(), CopyOptionValue::String("\\".into()))
|
||||
.unwrap();
|
||||
let mut parser = rows_with(set);
|
||||
let rows = feed_all(&mut parser, &[b"\"a\\\"b\",1.0,2023-11-14 22:13:20\n"]);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("a\"b".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_quoted_newline_and_chunk_boundaries() {
|
||||
let mut parser = parser();
|
||||
// Split inside a quoted field containing a newline, and a CRLF
|
||||
// record separator split across chunks. No batch is full, so all
|
||||
// feeds return nothing and both rows surface at the EOF signal.
|
||||
assert!(parser.feed(b"\"multi").unwrap().is_empty());
|
||||
assert!(
|
||||
parser
|
||||
.feed(b"line\nfield\",1.0,2023-11-14 22:13:20")
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
assert!(
|
||||
parser
|
||||
.feed(b"\r\nhost2,2.0,2023-11-14 22:13:21\n")
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
let rows = parser.finish().unwrap();
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("multiline\nfield".into()));
|
||||
assert_eq!(string_at(&rows[1], 0), Some("host2".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_header() {
|
||||
let mut set = CopyOptionSet::default();
|
||||
set.insert("header".into(), CopyOptionValue::Flag).unwrap();
|
||||
let mut parser = rows_with(set);
|
||||
let rows = feed_all(
|
||||
&mut parser,
|
||||
&[b"host,val,ts\nhost1,1.0,2023-11-14 22:13:20\n"],
|
||||
);
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_trailing_record_without_newline() {
|
||||
let mut parser = parser();
|
||||
let fed = parser
|
||||
.feed(b"host1,1.0,2023-11-14 22:13:20\nhost2,2.0,2023-11-14 22:13:21")
|
||||
.unwrap();
|
||||
// While the trailing record is pending, flush cannot emit the
|
||||
// completed one either; both surface at the EOF signal.
|
||||
assert!(fed.is_empty());
|
||||
let rows = parser.finish().unwrap();
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
|
||||
assert_eq!(string_at(&rows[1], 0), Some("host2".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_finish_at_record_boundary() {
|
||||
let mut parser = parser();
|
||||
// The complete record stays buffered (no full batch); the feed
|
||||
// emits nothing and `finish` surfaces it.
|
||||
assert!(
|
||||
parser
|
||||
.feed(b"host1,1.0,2023-11-14 22:13:20\n")
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
let rows = parser.finish().unwrap();
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(string_at(&rows[0], 0), Some("host1".into()));
|
||||
}
|
||||
|
||||
/// Regression test for chunk-boundary corruption: CopyData chunking is
|
||||
/// arbitrary (records may split mid-field), and an earlier feed loop
|
||||
/// flushed partial batches at record-boundary chunk ends, which
|
||||
/// arrow-csv does not permit mid-stream — fields then parse
|
||||
/// ("00001h,ost-001u,..."). Feed the same payload at many chunk sizes
|
||||
/// and verify EVERY cell.
|
||||
#[test]
|
||||
fn test_csv_chunked_feeds_are_alignment_safe() {
|
||||
// (int32 id, timestamp) rows of near-equal length.
|
||||
let sweep_columns = vec![
|
||||
ColumnSchema::new("i", ConcreteDataType::int32_datatype(), true),
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
];
|
||||
// A payload wider than one decoder batch so batch-full flushes
|
||||
// interleave with chunk boundaries too.
|
||||
let total = 2 * COPY_IN_FLUSH_ROWS + 137;
|
||||
let payload: Vec<u8> = (0..total)
|
||||
.map(|i| {
|
||||
format!(
|
||||
"{:06},2024-01-01 00:00:{:02}.{:03}\n",
|
||||
i,
|
||||
(i / 1000) % 60,
|
||||
i % 1000
|
||||
)
|
||||
})
|
||||
.collect::<String>()
|
||||
.into_bytes();
|
||||
|
||||
let check = |chunk: usize| {
|
||||
let mut parser = build_csv_codec(CopyOptionSet::default())
|
||||
.unwrap()
|
||||
.create_parser(&sweep_columns, "UTC")
|
||||
.unwrap();
|
||||
let mut ids: Vec<i32> = Vec::new();
|
||||
for c in payload.chunks(chunk) {
|
||||
for row in parser.feed(c).unwrap() {
|
||||
ids.push(match row.values[0].value_data.as_ref().unwrap() {
|
||||
api::v1::value::ValueData::I32Value(v) => *v,
|
||||
other => panic!("expected i32, got {other:?}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
for row in parser.finish().unwrap() {
|
||||
ids.push(match row.values[0].value_data.as_ref().unwrap() {
|
||||
api::v1::value::ValueData::I32Value(v) => *v,
|
||||
other => panic!("expected i32, got {other:?}"),
|
||||
});
|
||||
}
|
||||
assert_eq!(ids.len(), total, "chunk={chunk}");
|
||||
for (i, v) in ids.iter().enumerate() {
|
||||
assert_eq!(*v, i as i32, "chunk={chunk} row {i} corrupted");
|
||||
}
|
||||
};
|
||||
|
||||
// Small chunk sizes hit every possible boundary alignment,
|
||||
// including multiple decode calls per record.
|
||||
for chunk in [1usize, 2, 3, 7, 8, 13, 31, 64, 97, 128, 1024] {
|
||||
check(chunk);
|
||||
}
|
||||
// Client-natural sizes with a shifted start to vary where
|
||||
// boundaries land relative to records.
|
||||
for chunk in [4096usize, 8192, 65531] {
|
||||
check(chunk);
|
||||
let mut shifted = payload.clone();
|
||||
shifted.drain(..37);
|
||||
let mut parser = build_csv_codec(CopyOptionSet::default())
|
||||
.unwrap()
|
||||
.create_parser(&sweep_columns, "UTC")
|
||||
.unwrap();
|
||||
let mut count = 0usize;
|
||||
for c in shifted.chunks(chunk) {
|
||||
count += parser.feed(c).unwrap().len();
|
||||
}
|
||||
count += parser.finish().unwrap().len();
|
||||
assert_eq!(count, total - 1, "chunk={chunk} shifted");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_boolean_spellings() {
|
||||
// PostgreSQL COPY OUT emits `t`/`f`; COPY IN accepts the full
|
||||
// boolean input grammar.
|
||||
let columns = vec![
|
||||
ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true),
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
];
|
||||
let mut parser = build_csv_codec(CopyOptionSet::default())
|
||||
.unwrap()
|
||||
.create_parser(&columns, "UTC")
|
||||
.unwrap();
|
||||
let rows = feed_all(
|
||||
&mut parser,
|
||||
&[b"t,2024-01-01 00:00:00\n\
|
||||
f,2024-01-01 00:00:01\n\
|
||||
TRUE,2024-01-01 00:00:02\n\
|
||||
no,2024-01-01 00:00:03\n\
|
||||
on,2024-01-01 00:00:04\n\
|
||||
0,2024-01-01 00:00:05\n\
|
||||
,2024-01-01 00:00:06\n"],
|
||||
);
|
||||
let expected = [
|
||||
Some(true),
|
||||
Some(false),
|
||||
Some(true),
|
||||
Some(false),
|
||||
Some(true),
|
||||
Some(false),
|
||||
None,
|
||||
];
|
||||
for (row, want) in rows.iter().zip(expected) {
|
||||
assert_eq!(
|
||||
row.values[0].value_data,
|
||||
want.map(api::v1::value::ValueData::BoolValue),
|
||||
"{:?}",
|
||||
row.values[0].value_data
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_boolean_rejects_garbage() {
|
||||
let columns = vec![
|
||||
ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true),
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
];
|
||||
let mut parser = build_csv_codec(CopyOptionSet::default())
|
||||
.unwrap()
|
||||
.create_parser(&columns, "UTC")
|
||||
.unwrap();
|
||||
parser.feed(b"maybe,2024-01-01 00:00:00\n").unwrap();
|
||||
let err = match parser.finish() {
|
||||
Ok(_) => panic!("expected parse error"),
|
||||
Err(e) => e,
|
||||
};
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("invalid input syntax for type boolean"),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_timestamp_session_timezone() {
|
||||
// Offset-free timestamp strings must parse in the session
|
||||
// timezone (parity with SQL INSERT); explicit offsets win.
|
||||
let columns = vec![
|
||||
ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
];
|
||||
let mut parser = build_csv_codec(CopyOptionSet::default())
|
||||
.unwrap()
|
||||
.create_parser(&columns, "Asia/Shanghai")
|
||||
.unwrap();
|
||||
let rows = feed_all(
|
||||
&mut parser,
|
||||
&[b"host1,2024-01-01 00:00:00\n\
|
||||
host2,2024-01-01T00:00:00+08:00\n\
|
||||
host3,2024-01-01T00:00:00Z\n"],
|
||||
);
|
||||
let millis = |row: &ParsedRow| match row.values[1].value_data.as_ref().unwrap() {
|
||||
api::v1::value::ValueData::TimestampMillisecondValue(v) => *v,
|
||||
other => panic!("expected ts, got {other:?}"),
|
||||
};
|
||||
// 2024-01-01T00:00:00+08:00 == 2023-12-31T16:00:00Z
|
||||
assert_eq!(millis(&rows[0]), 1704038400000);
|
||||
// Explicit +08:00 offset: same instant as the naive parse.
|
||||
assert_eq!(millis(&rows[1]), 1704038400000);
|
||||
// Explicit Z: parsed as UTC regardless of the session zone.
|
||||
assert_eq!(millis(&rows[2]), 1704067200000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_csv_errors() {
|
||||
let mut p = parser();
|
||||
let err = match p.feed(b"host1,1.0\n") {
|
||||
Ok(_) => panic!("expected parse error"),
|
||||
Err(e) => e,
|
||||
};
|
||||
assert!(
|
||||
err.to_string().contains("incorrect number of fields"),
|
||||
"{err}"
|
||||
);
|
||||
|
||||
// Value conversion runs at flush time, so the bad double
|
||||
// surfaces at the EOF signal, not at feed time.
|
||||
let mut p = parser();
|
||||
p.feed(b"host1,not_a_double,2023-11-14 22:13:20\n").unwrap();
|
||||
let err = match p.finish() {
|
||||
Ok(_) => panic!("expected parse error"),
|
||||
Err(e) => e,
|
||||
};
|
||||
assert!(err.to_string().contains("failed to parse"), "{err}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,691 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//! PostgreSQL `COPY ... FROM STDIN` (copy-in) support.
|
||||
//!
|
||||
//! The wire handshake is driven by pgwire: returning [`Response::CopyIn`]
|
||||
//! from a query handler makes the server send a `CopyInResponse` and enter
|
||||
//! the copy-in sub-protocol; subsequent `CopyData` messages from the client
|
||||
//! are dispatched to [`CopyHandler`], whose implementation below parses
|
||||
//! the data stream incrementally and flushes it into the query handler as
|
||||
//! [`RowInsertRequests`] batches.
|
||||
//!
|
||||
//! # Format framework
|
||||
//!
|
||||
//! Each COPY data format is implemented as a [`codec::CopyInCodec`]: a
|
||||
//! stateless strategy that validates its options at statement-parse time,
|
||||
//! describes the column types it can ingest, and creates the stateful
|
||||
//! record parser for one stream. Formats register a builder under every
|
||||
//! accepted `FORMAT` name in [`REGISTERED_FORMATS`] — `csv` is the only
|
||||
//! implemented format today; PostgreSQL's `text`/`binary` and
|
||||
//! GreptimeDB-specific formats alike can be added the same way. See the
|
||||
//! `codec` module docs.
|
||||
|
||||
mod codec;
|
||||
mod csv;
|
||||
mod options;
|
||||
mod parse;
|
||||
|
||||
use std::fmt::Debug;
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::helper::ColumnDataTypeWrapper;
|
||||
use api::v1::column_def::options_from_column_schema;
|
||||
use api::v1::{
|
||||
Row as GrpcRow, RowInsertRequest, RowInsertRequests, Rows as GrpcRows, SemanticType,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use codec::{CopyInCodec, RecordParser};
|
||||
use common_query::OutputData;
|
||||
use common_telemetry::{info, warn};
|
||||
use datafusion::sql::sqlparser::ast::{CopyTarget, Statement as SqlParserStatement};
|
||||
use datatypes::data_type::DataType;
|
||||
use datatypes::schema::{ColumnSchema, SchemaRef};
|
||||
use futures::{Sink, SinkExt, stream};
|
||||
use pgwire::api::copy::CopyHandler;
|
||||
use pgwire::api::results::{CopyResponse, Response, Tag};
|
||||
use pgwire::api::{ClientInfo, PgWireConnectionState};
|
||||
use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
|
||||
use pgwire::messages::PgWireBackendMessage;
|
||||
use pgwire::messages::copy::{CopyData, CopyDone, CopyFail};
|
||||
use session::context::QueryContextRef;
|
||||
use table::metadata::TableInfo;
|
||||
|
||||
use crate::metrics::METRIC_POSTGRES_COPY_IN_ROWS;
|
||||
use crate::postgres::PostgresServerHandlerInner;
|
||||
use crate::postgres::types::PgErrorCode;
|
||||
use crate::postgres::utils::convert_err;
|
||||
|
||||
/// Number of parsed rows buffered before flushing a batch to the handler.
|
||||
///
|
||||
/// Each flush runs the full insert pipeline (catalog lookup, on-demand
|
||||
/// schema check, partition split, region write) and, without the
|
||||
/// pending-rows batcher, waits for it — so larger batches mean fewer
|
||||
/// round trips at the cost of linearly more memory per connection
|
||||
/// (decoder batch + rows buffer). The constant doubles as the arrow
|
||||
/// decoder's batch size. High-rate COPY users should enable the
|
||||
/// pending-rows batcher (`global-table-batch` / `BatcherOptions.protocols`
|
||||
/// containing `Postgres`), which turns flushes into enqueues and makes
|
||||
/// the per-batch fixed costs negligible.
|
||||
pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768;
|
||||
|
||||
/// Builds a codec from the format-neutral option set.
|
||||
type CodecBuilder = fn(options::CopyOptionSet) -> PgWireResult<Arc<dyn CopyInCodec>>;
|
||||
|
||||
/// The registered COPY FROM STDIN formats: accepted `FORMAT` name to
|
||||
/// builder. `csv` is the only implemented format; the others are
|
||||
/// registered so they report a clear "not supported yet" instead of an
|
||||
/// unknown-format error.
|
||||
const REGISTERED_FORMATS: &[(&str, CodecBuilder)] = &[
|
||||
("csv", csv::build_csv_codec),
|
||||
("text", build_unimplemented_text_codec),
|
||||
("txt", build_unimplemented_text_codec),
|
||||
("binary", build_unimplemented_binary_codec),
|
||||
];
|
||||
|
||||
fn build_unimplemented_text_codec(
|
||||
_options: options::CopyOptionSet,
|
||||
) -> PgWireResult<Arc<dyn CopyInCodec>> {
|
||||
Err(unsupported(
|
||||
"the text format is not supported yet; use WITH (FORMAT csv)",
|
||||
))
|
||||
}
|
||||
|
||||
fn build_unimplemented_binary_codec(
|
||||
_options: options::CopyOptionSet,
|
||||
) -> PgWireResult<Arc<dyn CopyInCodec>> {
|
||||
Err(unsupported(
|
||||
"COPY FROM STDIN WITH (FORMAT binary) is not supported yet",
|
||||
))
|
||||
}
|
||||
|
||||
/// Resolves the format codec from the parsed options.
|
||||
fn resolve_codec(options: &mut options::CopyOptionSet) -> PgWireResult<Arc<dyn CopyInCodec>> {
|
||||
let name = options.take_format()?.ok_or_else(|| {
|
||||
unsupported(
|
||||
"COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)",
|
||||
)
|
||||
})?;
|
||||
let builder = REGISTERED_FORMATS
|
||||
.iter()
|
||||
.find(|(registered, _)| *registered == name)
|
||||
.map(|(_, builder)| *builder)
|
||||
.ok_or_else(|| bad_copy_data(format!("unknown COPY format \"{}\"", name)))?;
|
||||
builder(std::mem::take(options))
|
||||
}
|
||||
|
||||
/// The parsed `COPY tbl [(columns)] FROM STDIN [(options)]` statement.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct CopyFromStdin {
|
||||
/// Table name parts as written (`t`, `schema.t`, `catalog.schema.t`);
|
||||
/// quoted names keep their case.
|
||||
table: Vec<String>,
|
||||
/// Requested column list; empty means all columns in table order.
|
||||
columns: Vec<String>,
|
||||
/// The resolved format codec, validated and ready to run.
|
||||
codec: Arc<dyn CopyInCodec>,
|
||||
}
|
||||
|
||||
fn copy_in_error(code: PgErrorCode, message: impl Into<String>) -> PgWireError {
|
||||
PgWireError::UserError(Box::new(code.to_err_info(message.into())))
|
||||
}
|
||||
|
||||
fn unsupported(message: impl Into<String>) -> PgWireError {
|
||||
copy_in_error(PgErrorCode::Ec0A000, message)
|
||||
}
|
||||
|
||||
fn bad_copy_data(message: impl Into<String>) -> PgWireError {
|
||||
copy_in_error(PgErrorCode::Ec22P04, message)
|
||||
}
|
||||
|
||||
fn invalid_text_repr(message: impl Into<String>) -> PgWireError {
|
||||
copy_in_error(PgErrorCode::Ec22P02, message)
|
||||
}
|
||||
|
||||
pub(crate) fn copy_in_must_run_alone_error() -> PgWireError {
|
||||
PgWireError::UserError(Box::new(ErrorInfo::new(
|
||||
"ERROR".to_string(),
|
||||
PgErrorCode::Ec0A000.code(),
|
||||
"COPY FROM STDIN must be executed alone".to_string(),
|
||||
)))
|
||||
}
|
||||
|
||||
/// Parses a leading `COPY ... FROM STDIN` statement from SQL text.
|
||||
///
|
||||
/// The statement is parsed textually — not through sqlparser — because the
|
||||
/// `WITH (...)` option list is open: formats may accept options beyond
|
||||
/// sqlparser's COPY grammar, and unknown options must reach the format
|
||||
/// codec instead of failing statement parsing. Returns `Ok(None)` when
|
||||
/// the string does not lead with a copy-in statement.
|
||||
pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult<Option<CopyFromStdin>> {
|
||||
let raw = match parse::scan_copy_from_stdin(sql)? {
|
||||
parse::CopyInScan::NotCopyIn => return Ok(None),
|
||||
parse::CopyInScan::CopyIn(raw) => raw,
|
||||
};
|
||||
let mut option_set = raw.options;
|
||||
let codec = resolve_codec(&mut option_set)?;
|
||||
Ok(Some(CopyFromStdin {
|
||||
table: raw.table,
|
||||
columns: raw.columns,
|
||||
codec,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Whether an already-parsed statement is a `COPY ... FROM STDIN`; used to
|
||||
/// reject copy-in statements trailing other statements, where the textual
|
||||
/// scanner (which only sees statement-leading COPYs) does not see them.
|
||||
pub(crate) fn is_copy_from_stdin(statement: &SqlParserStatement) -> bool {
|
||||
matches!(
|
||||
statement,
|
||||
SqlParserStatement::Copy {
|
||||
to: false,
|
||||
target: CopyTarget::Stdin,
|
||||
..
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
/// Resolves table name parts against the query context defaults.
|
||||
fn resolve_table_name(
|
||||
query_ctx: &QueryContextRef,
|
||||
parts: &[String],
|
||||
) -> PgWireResult<(String, String, String)> {
|
||||
match parts {
|
||||
[table] => Ok((
|
||||
query_ctx.current_catalog().to_string(),
|
||||
query_ctx.current_schema(),
|
||||
table.clone(),
|
||||
)),
|
||||
[schema, table] => Ok((
|
||||
query_ctx.current_catalog().to_string(),
|
||||
schema.clone(),
|
||||
table.clone(),
|
||||
)),
|
||||
[catalog, schema, table] => Ok((catalog.clone(), schema.clone(), table.clone())),
|
||||
_ => Err(unsupported(format!(
|
||||
"invalid table name for COPY: {}",
|
||||
parts.join(".")
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-connection state of an in-progress COPY FROM STDIN.
|
||||
pub(crate) struct CopyInState {
|
||||
query_ctx: QueryContextRef,
|
||||
/// The gRPC column schemas of the targeted columns. Cloned into each
|
||||
/// flushed batch: `Rows.schema` is an owned `Vec` in the proto
|
||||
/// request, so one copy per request is inherent.
|
||||
grpc_schema: Vec<api::v1::ColumnSchema>,
|
||||
table_name: String,
|
||||
parser: Box<dyn RecordParser + Send>,
|
||||
/// The resolved codec's name, for logging.
|
||||
codec_name: &'static str,
|
||||
/// The wire format code of the resolved codec.
|
||||
codec_format_code: i8,
|
||||
/// Parsed rows not yet flushed to the handler.
|
||||
rows: Vec<GrpcRow>,
|
||||
rows_written: u64,
|
||||
}
|
||||
|
||||
impl CopyInState {
|
||||
fn new(
|
||||
codec: Arc<dyn CopyInCodec>,
|
||||
columns: Vec<String>,
|
||||
table_info: Arc<TableInfo>,
|
||||
query_ctx: QueryContextRef,
|
||||
table_name: String,
|
||||
) -> PgWireResult<Self> {
|
||||
let schema = table_info.meta.schema.clone();
|
||||
let timestamp_index = schema.timestamp_index();
|
||||
|
||||
let columns = resolve_columns(codec.as_ref(), columns, &schema, &table_info)?;
|
||||
if let Some(ts_index) = timestamp_index {
|
||||
let ts_column = &schema.column_schemas()[ts_index];
|
||||
if !columns.iter().any(|column| column.name == ts_column.name) {
|
||||
return Err(unsupported(format!(
|
||||
"COPY FROM STDIN requires the time index column \"{}\" in the column list",
|
||||
ts_column.name
|
||||
)));
|
||||
}
|
||||
}
|
||||
let parser = codec.create_parser(&columns, query_ctx.timezone().to_string().as_str())?;
|
||||
|
||||
let grpc_schema = columns
|
||||
.iter()
|
||||
.map(|column| build_grpc_column_schema(column, &table_info))
|
||||
.collect::<PgWireResult<Vec<_>>>()?;
|
||||
|
||||
Ok(Self {
|
||||
query_ctx,
|
||||
grpc_schema,
|
||||
table_name,
|
||||
parser,
|
||||
codec_name: codec.name(),
|
||||
codec_format_code: codec.format_code(),
|
||||
rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS),
|
||||
rows_written: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// The wire format code advertised to the client.
|
||||
pub(crate) fn format_code(&self) -> i8 {
|
||||
self.codec_format_code
|
||||
}
|
||||
|
||||
/// Feeds a `CopyData` chunk into the parser and buffers parsed rows.
|
||||
fn feed(&mut self, data: &[u8]) -> PgWireResult<()> {
|
||||
self.rows.extend(self.parser.feed(data)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Parses the buffered tail at `CopyDone`; a final record without a
|
||||
/// trailing newline is accepted.
|
||||
fn finish(&mut self) -> PgWireResult<()> {
|
||||
self.rows.extend(self.parser.finish()?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn should_flush(&self) -> bool {
|
||||
self.rows.len() >= COPY_IN_FLUSH_ROWS
|
||||
}
|
||||
|
||||
/// Takes the buffered rows as an insert request, if any.
|
||||
fn take_requests(&mut self) -> Option<RowInsertRequests> {
|
||||
if self.rows.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let rows = GrpcRows {
|
||||
schema: self.grpc_schema.clone(),
|
||||
rows: std::mem::take(&mut self.rows),
|
||||
};
|
||||
Some(RowInsertRequests {
|
||||
inserts: vec![RowInsertRequest {
|
||||
table_name: self.table_name.clone(),
|
||||
rows: Some(rows),
|
||||
}],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_columns(
|
||||
codec: &dyn CopyInCodec,
|
||||
columns: Vec<String>,
|
||||
schema: &SchemaRef,
|
||||
table_info: &Arc<TableInfo>,
|
||||
) -> PgWireResult<Vec<ColumnSchema>> {
|
||||
let table_name = &table_info.name;
|
||||
let selected: Vec<ColumnSchema> = if columns.is_empty() {
|
||||
schema.column_schemas().to_vec()
|
||||
} else {
|
||||
let mut selected = Vec::with_capacity(columns.len());
|
||||
for name in &columns {
|
||||
let column = schema
|
||||
.column_schema_by_name(name)
|
||||
.ok_or_else(|| {
|
||||
copy_in_error(
|
||||
PgErrorCode::Ec42703,
|
||||
format!(
|
||||
"column \"{}\" of relation \"{}\" does not exist",
|
||||
name, table_name
|
||||
),
|
||||
)
|
||||
})?
|
||||
.clone();
|
||||
if selected.iter().any(|c: &ColumnSchema| &c.name == name) {
|
||||
return Err(bad_copy_data(format!(
|
||||
"column \"{}\" specified more than once",
|
||||
name
|
||||
)));
|
||||
}
|
||||
selected.push(column);
|
||||
}
|
||||
selected
|
||||
};
|
||||
|
||||
for column in &selected {
|
||||
codec.check_column_type(&column.data_type)?;
|
||||
}
|
||||
|
||||
Ok(selected)
|
||||
}
|
||||
|
||||
/// Derives a column's semantic type the same way the flow adapter does
|
||||
/// (`flow::adapter::util::column_schemas_to_proto`): the time-index flag
|
||||
/// for timestamps, primary-key membership for tags.
|
||||
fn semantic_type_of(table_info: &TableInfo, column: &ColumnSchema) -> PgWireResult<SemanticType> {
|
||||
let meta = &table_info.meta;
|
||||
if column.is_time_index() {
|
||||
Ok(SemanticType::Timestamp)
|
||||
} else if meta
|
||||
.primary_key_indices
|
||||
.iter()
|
||||
.any(|&index| meta.schema.column_schemas()[index].name == column.name)
|
||||
{
|
||||
Ok(SemanticType::Tag)
|
||||
} else {
|
||||
Ok(SemanticType::Field)
|
||||
}
|
||||
}
|
||||
|
||||
fn build_grpc_column_schema(
|
||||
column: &ColumnSchema,
|
||||
table_info: &Arc<TableInfo>,
|
||||
) -> PgWireResult<api::v1::ColumnSchema> {
|
||||
let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(column.data_type.clone())
|
||||
.map_err(|_| {
|
||||
unsupported(format!(
|
||||
"column type \"{}\" is not supported for COPY FROM STDIN yet",
|
||||
column.data_type.name()
|
||||
))
|
||||
})?
|
||||
.to_parts();
|
||||
Ok(api::v1::ColumnSchema {
|
||||
column_name: column.name.clone(),
|
||||
datatype: datatype as i32,
|
||||
semantic_type: semantic_type_of(table_info, column)? as i32,
|
||||
datatype_extension,
|
||||
options: options_from_column_schema(column),
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// pgwire handlers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
impl PostgresServerHandlerInner {
|
||||
/// Prepares a COPY FROM STDIN: resolves the target table, validates
|
||||
/// columns and stashes the per-connection copy state.
|
||||
pub(crate) async fn begin_copy_in(&self, stmt: CopyFromStdin) -> PgWireResult<Response> {
|
||||
// The batching-aware context routes copy-in flushes through the
|
||||
// pending-rows batcher like SQL INSERT when it is enabled for
|
||||
// this connection.
|
||||
let query_ctx = self.new_query_context();
|
||||
let (catalog, schema, table) = resolve_table_name(&query_ctx, &stmt.table)?;
|
||||
|
||||
let Some(table_info) = self
|
||||
.copy_in_handler
|
||||
.copy_in_table(&catalog, &schema, &table, query_ctx.clone())
|
||||
.await
|
||||
.map_err(convert_err)?
|
||||
else {
|
||||
return Err(copy_in_error(
|
||||
PgErrorCode::Ec42P01,
|
||||
format!("relation \"{}\" does not exist", stmt.table.join(".")),
|
||||
));
|
||||
};
|
||||
|
||||
// The insert path resolves the bare table name against the query
|
||||
// context, so point the context at the resolved catalog and schema
|
||||
// when the COPY target is qualified.
|
||||
let query_ctx = adjust_query_context(query_ctx, &catalog, &schema);
|
||||
|
||||
let state = CopyInState::new(stmt.codec, stmt.columns, table_info, query_ctx, table)?;
|
||||
let num_columns = state.grpc_schema.len();
|
||||
let format_code = state.format_code();
|
||||
|
||||
*self.copy_in_state.lock().await = Some(state);
|
||||
|
||||
Ok(Response::CopyIn(CopyResponse::new(
|
||||
format_code,
|
||||
num_columns,
|
||||
stream::empty(),
|
||||
)))
|
||||
}
|
||||
|
||||
/// Flushes buffered rows of an in-progress copy. Returns the number of
|
||||
/// rows written by the flushed batch.
|
||||
async fn flush_copy_in(&self, state: &mut CopyInState) -> PgWireResult<u64> {
|
||||
let Some(requests) = state.take_requests() else {
|
||||
return Ok(0);
|
||||
};
|
||||
let query_ctx = state.query_ctx.clone();
|
||||
let output = self
|
||||
.copy_in_handler
|
||||
.copy_in_insert(requests, query_ctx)
|
||||
.await
|
||||
.map_err(convert_err)?;
|
||||
match output.data {
|
||||
OutputData::AffectedRows(rows) => Ok(rows as u64),
|
||||
OutputData::RecordBatches(_) | OutputData::Stream(_) => {
|
||||
warn!("unexpected output from COPY FROM STDIN insert");
|
||||
Ok(0)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn adjust_query_context(
|
||||
query_ctx: QueryContextRef,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
) -> QueryContextRef {
|
||||
if query_ctx.current_catalog() == catalog && query_ctx.current_schema() == schema {
|
||||
return query_ctx;
|
||||
}
|
||||
let mut forked = query_ctx.fork();
|
||||
forked.set_current_catalog(catalog);
|
||||
forked.set_current_schema(schema);
|
||||
Arc::new(forked)
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CopyHandler for PostgresServerHandlerInner {
|
||||
async fn on_copy_data<C>(&self, _client: &mut C, copy_data: CopyData) -> PgWireResult<()>
|
||||
where
|
||||
C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
|
||||
C::Error: Debug,
|
||||
PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
|
||||
{
|
||||
let mut guard = self.copy_in_state.lock().await;
|
||||
let Some(mut state) = guard.take() else {
|
||||
return Err(copy_in_error(
|
||||
PgErrorCode::Ec08P01,
|
||||
"CopyData message received outside COPY FROM STDIN",
|
||||
));
|
||||
};
|
||||
// The state is owned here: any parse, conversion or insert error
|
||||
// below aborts the whole copy and drops the buffered data —
|
||||
// otherwise a pooled connection would retain failed-import
|
||||
// buffers until its next COPY. On success it is put back.
|
||||
state.feed(©_data.data)?;
|
||||
if state.should_flush() {
|
||||
let rows = self.flush_copy_in(&mut state).await?;
|
||||
state.rows_written += rows;
|
||||
METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows);
|
||||
}
|
||||
*guard = Some(state);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn on_copy_done<C>(&self, client: &mut C, _done: CopyDone) -> PgWireResult<()>
|
||||
where
|
||||
C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
|
||||
C::Error: Debug,
|
||||
PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
|
||||
{
|
||||
let mut guard = self.copy_in_state.lock().await;
|
||||
// Take ownership up front: even finalization failures (parse
|
||||
// errors in the trailing record, insert errors) drop the state
|
||||
// instead of retaining it on a long-lived connection.
|
||||
let Some(mut state) = guard.take() else {
|
||||
return Err(copy_in_error(
|
||||
PgErrorCode::Ec08P01,
|
||||
"CopyDone message received outside COPY FROM STDIN",
|
||||
));
|
||||
};
|
||||
state.finish()?;
|
||||
let rows = self.flush_copy_in(&mut state).await?;
|
||||
state.rows_written += rows;
|
||||
METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows);
|
||||
|
||||
let rows_written = state.rows_written;
|
||||
let codec_name = state.codec_name;
|
||||
info!(
|
||||
"PostgreSQL COPY FROM STDIN (format {}) finished: {} rows written",
|
||||
codec_name, rows_written
|
||||
);
|
||||
client
|
||||
.send(PgWireBackendMessage::CommandComplete(
|
||||
Tag::new("COPY").with_rows(rows_written as usize).into(),
|
||||
))
|
||||
.await?;
|
||||
|
||||
// After a successful extended-protocol copy-in, pgwire leaves the
|
||||
// connection in `CopyInProgress(true)` and swallows the client's
|
||||
// terminating `Sync` message, so `ReadyForQuery` would never be
|
||||
// sent and the connection would deadlock. Route the state to
|
||||
// `AwaitingSync` so the `Sync` is dispatched to `on_sync`, which
|
||||
// completes the round trip.
|
||||
if let PgWireConnectionState::CopyInProgress(true) = client.state() {
|
||||
client.set_state(PgWireConnectionState::AwaitingSync);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn on_copy_fail<C>(&self, _client: &mut C, fail: CopyFail) -> PgWireError
|
||||
where
|
||||
C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
|
||||
C::Error: Debug,
|
||||
PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
|
||||
{
|
||||
// Drop the buffered state; partial data of a failed copy is not
|
||||
// written.
|
||||
self.copy_in_state.lock().await.take();
|
||||
copy_in_error(
|
||||
PgErrorCode::Ec57014,
|
||||
format!("COPY FROM STDIN mode terminated by user: {}", fail.message),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_copy_from_stdin_requires_format() {
|
||||
let err = copy_from_stdin_from_sql("COPY t FROM STDIN").unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("requires an explicit format in this version"),
|
||||
"{err}"
|
||||
);
|
||||
|
||||
let copy = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv)")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(copy.table, vec!["t"]);
|
||||
assert!(copy.columns.is_empty());
|
||||
assert_eq!(copy.codec.name(), "csv");
|
||||
assert_eq!(copy.codec.format_code(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_copy_from_stdin_csv() {
|
||||
let copy = copy_from_stdin_from_sql(
|
||||
"COPY t (a, b) FROM STDIN WITH (FORMAT csv, HEADER, DELIMITER ';')",
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(copy.columns, vec!["a", "b"]);
|
||||
assert_eq!(copy.codec.name(), "csv");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_copy_from_stdin_text_unsupported() {
|
||||
for sql in [
|
||||
"COPY t FROM STDIN WITH (FORMAT text)",
|
||||
"COPY t FROM STDIN WITH (FORMAT txt)",
|
||||
] {
|
||||
let err = copy_from_stdin_from_sql(sql).unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("text format is not supported yet"),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_copy_from_stdin_binary_unsupported() {
|
||||
let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT binary)").unwrap_err();
|
||||
assert!(err.to_string().contains("not supported"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_copy_from_stdin_unknown_format() {
|
||||
let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT parquet)").unwrap_err();
|
||||
assert!(err.to_string().contains("unknown COPY format"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_copy_from_stdin_not_copy_in() {
|
||||
assert!(
|
||||
copy_from_stdin_from_sql("COPY (SELECT 1) TO STDOUT")
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
copy_from_stdin_from_sql("COPY t TO '/tmp/x.csv'")
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
// GreptimeDB's own file-based COPY FROM.
|
||||
assert!(
|
||||
copy_from_stdin_from_sql("COPY t FROM '/tmp/x.csv' WITH (FORMAT csv)")
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_redundant_options() {
|
||||
let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)")
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("redundant FORMAT"));
|
||||
|
||||
let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (DELIMITER ',', DELIMITER ';')")
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("redundant DELIMITER"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_specific_options_rejected() {
|
||||
// The open WITH section parses anything, but the format rejects
|
||||
// the options it does not understand.
|
||||
let err =
|
||||
copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, PRECISION 'ns', FREEZE)")
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("not supported by format csv"),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multi_statement_rejected() {
|
||||
let err = copy_from_stdin_from_sql("COPY t FROM STDIN; SELECT 1").unwrap_err();
|
||||
assert!(err.to_string().contains("alone"));
|
||||
// Trailing separators and quoted semicolons are fine.
|
||||
assert!(
|
||||
copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv);")
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(
|
||||
copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';')")
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//! Format-neutral collection of COPY `WITH (...)` options.
|
||||
//!
|
||||
//! The option list is **open**: every `KEY value` entry the client writes
|
||||
//! is collected here, regardless of whether any format understands it.
|
||||
//! Each format codec then takes the options it knows ([`CopyOptionSet::take_string`]
|
||||
//! and friends) and rejects everything else via
|
||||
//! [`CopyOptionSet::reject_leftovers`]. This way a format can introduce
|
||||
//! its own options — say a line-protocol format accepting
|
||||
//! `WITH (FORMAT influxdb_line, precision 'ns')` — without touching the
|
||||
//! statement parser or any other codec.
|
||||
|
||||
use pgwire::error::{PgWireError, PgWireResult};
|
||||
|
||||
use crate::postgres::copy_in::bad_copy_data;
|
||||
use crate::postgres::types::PgErrorCode;
|
||||
|
||||
/// The value of one COPY option.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) enum CopyOptionValue {
|
||||
/// A single-quoted string literal.
|
||||
String(String),
|
||||
/// A numeric literal, kept as written.
|
||||
Number(String),
|
||||
/// A bare identifier.
|
||||
Ident(String),
|
||||
/// `TRUE` / `FALSE`.
|
||||
Boolean(bool),
|
||||
/// The option appears without a value.
|
||||
Flag,
|
||||
}
|
||||
|
||||
impl CopyOptionValue {
|
||||
/// The value as a descriptive word for error messages.
|
||||
pub(crate) fn kind(&self) -> &'static str {
|
||||
match self {
|
||||
CopyOptionValue::String(_) => "a string value",
|
||||
CopyOptionValue::Number(_) => "a numeric value",
|
||||
CopyOptionValue::Ident(_) => "an identifier value",
|
||||
CopyOptionValue::Boolean(_) => "a boolean value",
|
||||
CopyOptionValue::Flag => "no value",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The options of a `COPY ... FROM STDIN WITH (...)` statement.
|
||||
///
|
||||
/// Entries keep statement order; keys are lowercased for lookup and
|
||||
/// displayed uppercased in error messages.
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct CopyOptionSet {
|
||||
entries: Vec<(String, CopyOptionValue)>,
|
||||
}
|
||||
|
||||
impl CopyOptionSet {
|
||||
/// Adds an option; duplicate keys are rejected like PostgreSQL does.
|
||||
pub(crate) fn insert(&mut self, key: String, value: CopyOptionValue) -> PgWireResult<()> {
|
||||
if self.entries.iter().any(|(k, _)| *k == key) {
|
||||
return Err(bad_copy_data(format!(
|
||||
"conflicting or redundant {} option",
|
||||
key.to_uppercase()
|
||||
)));
|
||||
}
|
||||
self.entries.push((key, value));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Takes the raw value of `key`; for codec-specific options whose
|
||||
/// shape this module does not prescribe.
|
||||
pub(crate) fn take_value(&mut self, key: &str) -> Option<CopyOptionValue> {
|
||||
let key = key.to_lowercase();
|
||||
let index = self.entries.iter().position(|(k, _)| *k == key)?;
|
||||
Some(self.entries.remove(index).1)
|
||||
}
|
||||
|
||||
/// Takes the `FORMAT` option, lowercased. Accepts identifier or string
|
||||
/// values.
|
||||
pub(crate) fn take_format(&mut self) -> PgWireResult<Option<String>> {
|
||||
let Some(value) = self.take_value("format") else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
CopyOptionValue::Ident(name) | CopyOptionValue::String(name) => {
|
||||
Ok(Some(name.to_lowercase()))
|
||||
}
|
||||
other => Err(bad_copy_data(format!(
|
||||
"FORMAT requires an identifier value, got {}",
|
||||
other.kind()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Takes a string-valued option.
|
||||
pub(crate) fn take_string(&mut self, key: &str) -> PgWireResult<Option<String>> {
|
||||
let Some(value) = self.take_value(key) else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
CopyOptionValue::String(s) => Ok(Some(s)),
|
||||
other => Err(bad_copy_data(format!(
|
||||
"COPY option {} requires a string value, got {}",
|
||||
key.to_uppercase(),
|
||||
other.kind()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Takes a single one-byte character option. A tab is a legal (and
|
||||
/// common) CSV delimiter; newlines and other control characters are
|
||||
/// not, matching PostgreSQL's rules.
|
||||
pub(crate) fn take_single_byte(&mut self, key: &str) -> PgWireResult<Option<u8>> {
|
||||
let Some(value) = self.take_string(key)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut chars = value.chars();
|
||||
let (c, rest) = (chars.next(), chars.as_str());
|
||||
match c {
|
||||
Some('\t') if rest.is_empty() => Ok(Some(b'\t')),
|
||||
Some(c) if rest.is_empty() && c.is_ascii() && !c.is_ascii_control() => {
|
||||
Ok(Some(c as u8))
|
||||
}
|
||||
_ => Err(bad_copy_data(format!(
|
||||
"COPY option {} must be a single one-byte character other than a newline or carriage return",
|
||||
key.to_uppercase()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Takes a boolean option; a valueless option means `TRUE`.
|
||||
pub(crate) fn take_flag(&mut self, key: &str) -> PgWireResult<Option<bool>> {
|
||||
let Some(value) = self.take_value(key) else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
CopyOptionValue::Flag => Ok(Some(true)),
|
||||
CopyOptionValue::Boolean(b) => Ok(Some(b)),
|
||||
other => Err(bad_copy_data(format!(
|
||||
"COPY option {} requires a boolean value, got {}",
|
||||
key.to_uppercase(),
|
||||
other.kind()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Rejects every option not consumed by the format, with one message
|
||||
/// naming them all. Codecs call this after their `take_*` calls so a
|
||||
/// typo cannot be silently ignored.
|
||||
pub(crate) fn reject_leftovers(self, format: &str) -> PgWireResult<()> {
|
||||
if self.entries.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let unknown = self
|
||||
.entries
|
||||
.iter()
|
||||
.map(|(key, _)| key.to_uppercase())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
Err(PgWireError::UserError(Box::new(
|
||||
PgErrorCode::Ec0A000.to_err_info(format!(
|
||||
"COPY option(s) {unknown} are not supported by format {format}"
|
||||
)),
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn new_set(entries: &[(&str, CopyOptionValue)]) -> CopyOptionSet {
|
||||
let mut set = CopyOptionSet::default();
|
||||
for (key, value) in entries {
|
||||
set.insert((*key).to_string(), value.clone()).unwrap();
|
||||
}
|
||||
set
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_rejects_duplicates() {
|
||||
let mut set = CopyOptionSet::default();
|
||||
set.insert("format".to_string(), CopyOptionValue::Ident("csv".into()))
|
||||
.unwrap();
|
||||
let err = set
|
||||
.insert("format".to_string(), CopyOptionValue::Ident("text".into()))
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("redundant FORMAT"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_take_single_byte_tab_allowed() {
|
||||
// A literal tab is a legal delimiter (PostgreSQL permits it).
|
||||
let mut set = new_set(&[("delimiter", CopyOptionValue::String("\t".into()))]);
|
||||
assert_eq!(set.take_single_byte("DELIMITER").unwrap(), Some(b'\t'));
|
||||
|
||||
// Newlines and other control characters stay rejected.
|
||||
for bad in ["\n", "\r", "\0"] {
|
||||
let mut set = new_set(&[("delimiter", CopyOptionValue::String(bad.into()))]);
|
||||
assert!(
|
||||
set.take_single_byte("DELIMITER").is_err(),
|
||||
"{bad:?} must be rejected"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_take_string_and_single_byte() {
|
||||
let mut set = new_set(&[
|
||||
("null", CopyOptionValue::String("nil".into())),
|
||||
("delimiter", CopyOptionValue::String("|".into())),
|
||||
]);
|
||||
assert_eq!(set.take_string("null").unwrap().as_deref(), Some("nil"));
|
||||
assert_eq!(set.take_single_byte("delimiter").unwrap(), Some(b'|'));
|
||||
|
||||
let mut set = new_set(&[("delimiter", CopyOptionValue::Number("42".into()))]);
|
||||
let err = set.take_single_byte("delimiter").unwrap_err();
|
||||
assert!(err.to_string().contains("requires a string value"));
|
||||
|
||||
let mut set = new_set(&[("delimiter", CopyOptionValue::String("ab".into()))]);
|
||||
let err = set.take_single_byte("delimiter").unwrap_err();
|
||||
assert!(err.to_string().contains("single one-byte"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_take_flag() {
|
||||
let mut set = new_set(&[
|
||||
("header", CopyOptionValue::Flag),
|
||||
("frozen", CopyOptionValue::Boolean(false)),
|
||||
]);
|
||||
assert_eq!(set.take_flag("header").unwrap(), Some(true));
|
||||
assert_eq!(set.take_flag("frozen").unwrap(), Some(false));
|
||||
assert_eq!(set.take_flag("absent").unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_take_format_normalizes() {
|
||||
let mut set = new_set(&[("format", CopyOptionValue::Ident("CSV".into()))]);
|
||||
assert_eq!(set.take_format().unwrap().as_deref(), Some("csv"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reject_leftovers_lists_all() {
|
||||
let mut set = new_set(&[
|
||||
("precision", CopyOptionValue::String("ns".into())),
|
||||
("weird", CopyOptionValue::Flag),
|
||||
]);
|
||||
set.take_string("precision").unwrap();
|
||||
let err = set.reject_leftovers("text").unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("WEIRD"),
|
||||
"leftover option must be listed: {err}"
|
||||
);
|
||||
assert!(
|
||||
err.to_string().contains("not supported by format text"),
|
||||
"{err}"
|
||||
);
|
||||
// The consumed option is not listed.
|
||||
assert!(!err.to_string().contains("PRECISION"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_raw_value_escape_hatch() {
|
||||
let mut set = new_set(&[("precision", CopyOptionValue::Ident("ns".into()))]);
|
||||
assert_eq!(
|
||||
set.take_value("precision"),
|
||||
Some(CopyOptionValue::Ident("ns".into()))
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//! Textual parser for `COPY ... FROM STDIN` statements.
|
||||
//!
|
||||
//! COPY FROM STDIN is parsed here instead of through sqlparser because the
|
||||
//! `WITH (...)` option list is open: formats may accept options beyond
|
||||
//! sqlparser's COPY grammar (for example a line-protocol format taking
|
||||
//! `precision`), and unknown options must reach the format codec rather
|
||||
//! than fail statement parsing.
|
||||
//!
|
||||
//! The parser commits to error reporting only after `FROM STDIN` has been
|
||||
//! seen; anything else (like `COPY ... TO STDOUT`) is reported as
|
||||
//! [`CopyInScan::NotCopyIn`] so the caller falls back to the regular SQL
|
||||
//! path.
|
||||
|
||||
use datafusion::sql::sqlparser::dialect::GenericDialect;
|
||||
use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer};
|
||||
use pgwire::error::PgWireResult;
|
||||
|
||||
use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue};
|
||||
use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error};
|
||||
|
||||
/// A scanned `COPY tbl [(cols)] FROM STDIN [WITH (...)]` statement.
|
||||
///
|
||||
/// Unquoted identifiers are folded to lowercase and quoted ones keep their
|
||||
/// case (`COPY Metrics` addresses the same table as `CREATE TABLE Metrics`,
|
||||
/// namely `metrics`), matching `ParserContext::canonicalize_identifier` in
|
||||
/// the SQL parser and PostgreSQL's own folding rule.
|
||||
#[derive(Debug)]
|
||||
pub(super) struct RawCopyFromStdin {
|
||||
/// Table name parts, as written (`t`, `schema.t`,
|
||||
/// `catalog.schema.t`); quoted names keep their case.
|
||||
pub(super) table: Vec<String>,
|
||||
/// Requested column list; empty means all columns in table order.
|
||||
pub(super) columns: Vec<String>,
|
||||
/// The open option list, handed to the resolved format codec.
|
||||
pub(super) options: CopyOptionSet,
|
||||
}
|
||||
|
||||
pub(super) enum CopyInScan {
|
||||
/// Not a `COPY ... FROM STDIN` statement.
|
||||
NotCopyIn,
|
||||
/// A complete copy-in statement with only trailing separators left.
|
||||
CopyIn(RawCopyFromStdin),
|
||||
}
|
||||
|
||||
/// Scans a SQL string for a leading `COPY ... FROM STDIN` statement.
|
||||
pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult<CopyInScan> {
|
||||
// Fast path: the statement must start with the COPY keyword, possibly
|
||||
// behind comments; only then is a full tokenization worth its cost.
|
||||
let trimmed = sql.trim_start();
|
||||
let starts_like_copy = trimmed
|
||||
.as_bytes()
|
||||
.get(..4)
|
||||
.is_some_and(|prefix| prefix.eq_ignore_ascii_case(b"copy"));
|
||||
let starts_with_comment = trimmed.starts_with("--") || trimmed.starts_with("/*");
|
||||
if !starts_like_copy && !starts_with_comment {
|
||||
return Ok(CopyInScan::NotCopyIn);
|
||||
}
|
||||
|
||||
let tokens = match Tokenizer::new(&GenericDialect {}, sql).tokenize() {
|
||||
Ok(tokens) => tokens,
|
||||
// Let the regular SQL path produce the parse error.
|
||||
Err(_) => return Ok(CopyInScan::NotCopyIn),
|
||||
};
|
||||
// The tokenizer keeps whitespace; the scanner works on tokens only.
|
||||
let tokens: Vec<Token> = tokens
|
||||
.into_iter()
|
||||
.filter(|token| !matches!(token, Token::Whitespace(_)))
|
||||
.collect();
|
||||
let mut cursor = Tokens {
|
||||
tokens: &tokens,
|
||||
pos: 0,
|
||||
};
|
||||
|
||||
if !cursor.take_keyword("copy") {
|
||||
return Ok(CopyInScan::NotCopyIn);
|
||||
}
|
||||
|
||||
let table = match parse_object_name(&mut cursor) {
|
||||
Some(table) => table,
|
||||
None => return Ok(CopyInScan::NotCopyIn),
|
||||
};
|
||||
let columns = match parse_column_list(&mut cursor) {
|
||||
Some(columns) => columns,
|
||||
None => return Ok(CopyInScan::NotCopyIn),
|
||||
};
|
||||
|
||||
if !cursor.take_keyword("from") || !cursor.take_keyword("stdin") {
|
||||
// A COPY with another direction or source; not ours to handle.
|
||||
return Ok(CopyInScan::NotCopyIn);
|
||||
}
|
||||
|
||||
let mut options = CopyOptionSet::default();
|
||||
// PostgreSQL allows the WITH keyword to be omitted.
|
||||
let has_options = if cursor.take_keyword("with") {
|
||||
if !matches!(cursor.peek(), Some(Token::LParen)) {
|
||||
return Err(bad_copy_data("expected '(' after WITH in COPY FROM STDIN"));
|
||||
}
|
||||
true
|
||||
} else {
|
||||
matches!(cursor.peek(), Some(Token::LParen))
|
||||
};
|
||||
if has_options {
|
||||
// Consume the opening '('.
|
||||
cursor.next();
|
||||
parse_option_list(&mut cursor, &mut options)?;
|
||||
}
|
||||
|
||||
// Only statement separators may follow.
|
||||
while let Some(token) = cursor.next() {
|
||||
if !matches!(token, Token::SemiColon) {
|
||||
return Err(copy_in_must_run_alone_error());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(CopyInScan::CopyIn(RawCopyFromStdin {
|
||||
table,
|
||||
columns,
|
||||
options,
|
||||
}))
|
||||
}
|
||||
|
||||
struct Tokens<'a> {
|
||||
tokens: &'a [Token],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl<'a> Tokens<'a> {
|
||||
fn next(&mut self) -> Option<&'a Token> {
|
||||
let token = self.tokens.get(self.pos)?;
|
||||
self.pos += 1;
|
||||
Some(token)
|
||||
}
|
||||
|
||||
fn peek(&mut self) -> Option<&'a Token> {
|
||||
self.tokens.get(self.pos)
|
||||
}
|
||||
|
||||
/// Consumes the next token if it is the unquoted keyword `keyword`.
|
||||
fn take_keyword(&mut self, keyword: &str) -> bool {
|
||||
match self.peek() {
|
||||
Some(Token::Word(w))
|
||||
if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword) =>
|
||||
{
|
||||
self.pos += 1;
|
||||
true
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the identifier with the SQL parser's canonicalization: quoted
|
||||
/// identifiers keep their case, unquoted ones fold to lowercase.
|
||||
fn expect_ident(cursor: &mut Tokens<'_>) -> Option<String> {
|
||||
match cursor.next() {
|
||||
Some(Token::Word(w)) => Some(if w.quote_style.is_none() {
|
||||
w.value.to_lowercase()
|
||||
} else {
|
||||
w.value.clone()
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_object_name(cursor: &mut Tokens<'_>) -> Option<Vec<String>> {
|
||||
let mut parts = vec![expect_ident(cursor)?];
|
||||
while matches!(cursor.peek(), Some(Token::Period)) {
|
||||
cursor.next();
|
||||
parts.push(expect_ident(cursor)?);
|
||||
}
|
||||
Some(parts)
|
||||
}
|
||||
|
||||
/// Parses an optional `(column, ...)` list. `None` means the statement is
|
||||
/// not a copy-in shape (either no list, or a malformed list better left to
|
||||
/// the regular parser to report).
|
||||
fn parse_column_list(cursor: &mut Tokens<'_>) -> Option<Vec<String>> {
|
||||
if !matches!(cursor.peek(), Some(Token::LParen)) {
|
||||
return Some(Vec::new());
|
||||
}
|
||||
cursor.next();
|
||||
let mut columns = vec![expect_ident(cursor)?];
|
||||
loop {
|
||||
match cursor.next() {
|
||||
Some(Token::Comma) => columns.push(expect_ident(cursor)?),
|
||||
Some(Token::RParen) => break,
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
Some(columns)
|
||||
}
|
||||
|
||||
/// Parses `NAME [value] (, NAME [value])*` until the closing `)`.
|
||||
///
|
||||
/// Option values are open-ended: a string literal, a number, an identifier
|
||||
/// or a boolean word. Whether an option makes sense is decided by the
|
||||
/// format codec, not here.
|
||||
fn parse_option_list(cursor: &mut Tokens<'_>, options: &mut CopyOptionSet) -> PgWireResult<()> {
|
||||
loop {
|
||||
let key = match cursor.next() {
|
||||
Some(Token::Word(w)) if w.quote_style.is_none() => w.value.to_lowercase(),
|
||||
_ => return Err(bad_copy_data("expected an option name after '(' or ','")),
|
||||
};
|
||||
let value = match cursor.peek() {
|
||||
Some(Token::Comma) | Some(Token::RParen) => CopyOptionValue::Flag,
|
||||
Some(Token::SingleQuotedString(s)) => {
|
||||
let value = CopyOptionValue::String(s.clone());
|
||||
cursor.next();
|
||||
value
|
||||
}
|
||||
Some(Token::EscapedStringLiteral(s)) => {
|
||||
// E'...' escape-string constants, e.g. DELIMITER E'\t':
|
||||
// sqlparser's tokenizer already decoded the escapes, so
|
||||
// the payload is used as-is, like SingleQuotedString.
|
||||
let value = CopyOptionValue::String(s.clone());
|
||||
cursor.next();
|
||||
value
|
||||
}
|
||||
Some(Token::Number(n, _)) => {
|
||||
let value = CopyOptionValue::Number(n.clone());
|
||||
cursor.next();
|
||||
value
|
||||
}
|
||||
Some(Token::Word(w)) => {
|
||||
let value = if w.quote_style.is_none() {
|
||||
match w.value.to_ascii_lowercase().as_str() {
|
||||
"true" => CopyOptionValue::Boolean(true),
|
||||
"false" => CopyOptionValue::Boolean(false),
|
||||
_ => CopyOptionValue::Ident(w.value.clone()),
|
||||
}
|
||||
} else {
|
||||
CopyOptionValue::Ident(w.value.clone())
|
||||
};
|
||||
cursor.next();
|
||||
value
|
||||
}
|
||||
_ => return Err(bad_copy_data("expected an option value or ',' or ')'")),
|
||||
};
|
||||
options.insert(key, value)?;
|
||||
match cursor.next() {
|
||||
Some(Token::Comma) => continue,
|
||||
Some(Token::RParen) => break,
|
||||
_ => {
|
||||
return Err(bad_copy_data(
|
||||
"expected ',' or ')' in COPY FROM STDIN options",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn scan(sql: &str) -> PgWireResult<Option<RawCopyFromStdin>> {
|
||||
Ok(match scan_copy_from_stdin(sql)? {
|
||||
CopyInScan::CopyIn(raw) => Some(raw),
|
||||
CopyInScan::NotCopyIn => None,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_basic() {
|
||||
let raw = scan("COPY t FROM STDIN").unwrap().unwrap();
|
||||
assert_eq!(raw.table, vec!["t"]);
|
||||
assert!(raw.columns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_qualified_table_and_columns() {
|
||||
let raw = scan("COPY db.public.Metrics (ts, \"Host\") FROM STDIN")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(raw.table, vec!["db", "public", "metrics"]);
|
||||
assert_eq!(raw.columns, vec!["ts", "Host"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_not_copy_in() {
|
||||
assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none());
|
||||
assert!(
|
||||
scan("COPY t TO '/tmp/x.csv' WITH (FORMAT csv)")
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(scan("SELECT 1").unwrap().is_none());
|
||||
// GreptimeDB's own file-based COPY FROM.
|
||||
assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_leading_comment() {
|
||||
let raw = scan("-- a comment\nCOPY t FROM STDIN").unwrap().unwrap();
|
||||
assert_eq!(raw.table, vec!["t"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_folds_unquoted_identifiers() {
|
||||
// Unquoted identifiers fold to lowercase, like the SQL parser.
|
||||
let raw = scan("copy Metrics (Host, ts) from stdin").unwrap().unwrap();
|
||||
assert_eq!(raw.table, vec!["metrics"]);
|
||||
assert_eq!(raw.columns, vec!["host", "ts"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_preserves_quoted_identifiers() {
|
||||
let raw = scan("COPY \"Metrics\" (\"Host\") FROM STDIN")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(raw.table, vec!["Metrics"]);
|
||||
assert_eq!(raw.columns, vec!["Host"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_options_kinds() {
|
||||
let raw = scan(
|
||||
"COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';', HEADER, \
|
||||
skip_bad_records true, precision 3, mode strict)",
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut options = raw.options;
|
||||
assert_eq!(options.take_format().unwrap().as_deref(), Some("csv"));
|
||||
assert_eq!(options.take_single_byte("DELIMITER").unwrap(), Some(b';'));
|
||||
assert_eq!(options.take_flag("HEADER").unwrap(), Some(true));
|
||||
assert_eq!(
|
||||
options.take_value("skip_bad_records"),
|
||||
Some(CopyOptionValue::Boolean(true))
|
||||
);
|
||||
assert_eq!(
|
||||
options.take_value("precision"),
|
||||
Some(CopyOptionValue::Number("3".into()))
|
||||
);
|
||||
assert_eq!(
|
||||
options.take_value("mode"),
|
||||
Some(CopyOptionValue::Ident("strict".into()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_escaped_string_constants() {
|
||||
// sqlparser's tokenizer already decodes E'...' escapes, so the
|
||||
// payload is used as-is: E'\t' is a tab.
|
||||
let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\t')")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
raw.options.take_single_byte("DELIMITER").unwrap(),
|
||||
Some(b'\t')
|
||||
);
|
||||
|
||||
// E'\\' is a single literal backslash (a legal one-byte
|
||||
// delimiter), NOT an escaped backslash-t.
|
||||
let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\')")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
raw.options.take_single_byte("DELIMITER").unwrap(),
|
||||
Some(b'\\')
|
||||
);
|
||||
|
||||
// E'\\t' decodes to TWO characters — backslash followed by t —
|
||||
// and must be rejected as a multi-character delimiter instead of
|
||||
// being unescaped a second time into a tab.
|
||||
// The scanner leaves option validation to the codec; use the
|
||||
// full resolution path to assert the delimiter rejection.
|
||||
let err = match crate::postgres::copy_in::copy_from_stdin_from_sql(
|
||||
"COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\t')",
|
||||
) {
|
||||
Ok(Some(_)) => panic!("expected the multi-character delimiter to be rejected"),
|
||||
Ok(None) => panic!("expected the statement to scan as copy-in"),
|
||||
Err(e) => e,
|
||||
};
|
||||
assert!(
|
||||
err.to_string().contains("single one-byte character"),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_semicolon_inside_option_value() {
|
||||
// A quoted semicolon is not a statement separator.
|
||||
let mut raw = scan("COPY t FROM STDIN WITH (NULL ';')").unwrap().unwrap();
|
||||
assert_eq!(
|
||||
raw.options.take_string("NULL").unwrap().as_deref(),
|
||||
Some(";")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_trailing_statements_rejected() {
|
||||
let err = scan("COPY t FROM STDIN; SELECT 1").unwrap_err();
|
||||
assert!(err.to_string().contains("alone"));
|
||||
|
||||
// Trailing separators alone are fine.
|
||||
assert!(scan("COPY t FROM STDIN;").unwrap().is_some());
|
||||
assert!(scan("COPY t FROM STDIN ; ;").unwrap().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_options_without_with_keyword() {
|
||||
let raw = scan("COPY t FROM STDIN (FORMAT csv, HEADER)")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut options = raw.options;
|
||||
assert_eq!(options.take_format().unwrap().as_deref(), Some("csv"));
|
||||
assert_eq!(options.take_flag("HEADER").unwrap(), Some(true));
|
||||
assert!(options.reject_leftovers("csv").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_malformed_after_stdin() {
|
||||
let err = scan("COPY t FROM STDIN WITH").unwrap_err();
|
||||
assert!(err.to_string().contains("'('"), "{err}");
|
||||
let err = scan("COPY t FROM STDIN WITH (").unwrap_err();
|
||||
assert!(err.to_string().contains("option"), "{err}");
|
||||
// A malformed column list never commits to the copy-in shape and
|
||||
// falls through to the regular parser.
|
||||
assert!(scan("COPY t (a,").unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_redundant_option() {
|
||||
let err = scan("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)").unwrap_err();
|
||||
assert!(err.to_string().contains("redundant FORMAT"));
|
||||
}
|
||||
}
|
||||
@@ -56,11 +56,14 @@ use crate::SqlPlan;
|
||||
use crate::error::{DataFusionSnafu, InferParameterTypesSnafu, Result};
|
||||
use crate::postgres::types::*;
|
||||
use crate::postgres::utils::convert_err;
|
||||
use crate::postgres::{PostgresServerHandlerInner, fixtures};
|
||||
use crate::postgres::{PostgresServerHandlerInner, copy_in, fixtures};
|
||||
use crate::query_handler::sql::ServerSqlQueryHandlerRef;
|
||||
|
||||
impl PostgresServerHandlerInner {
|
||||
fn new_query_context(&self) -> QueryContextRef {
|
||||
/// Creates a query context carrying this connection's batching
|
||||
/// selection, so write paths can route through the pending-rows
|
||||
/// batcher when it is enabled.
|
||||
pub(super) fn new_query_context(&self) -> QueryContextRef {
|
||||
let mut ctx = self.session.new_query_context();
|
||||
Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled);
|
||||
ctx
|
||||
@@ -87,8 +90,25 @@ impl SimpleQueryHandler for PostgresServerHandlerInner {
|
||||
return Ok(vec![Response::EmptyQuery]);
|
||||
}
|
||||
|
||||
// A leading COPY FROM STDIN takes over the connection for the copy
|
||||
// sub-protocol; it is parsed textually so the WITH clause may carry
|
||||
// options beyond sqlparser's COPY grammar. It must run alone.
|
||||
if let Some(copy_from_stdin) = copy_in::copy_from_stdin_from_sql(query)? {
|
||||
let response = self.begin_copy_in(copy_from_stdin).await?;
|
||||
send_warning_opt(client, query_ctx).await?;
|
||||
return Ok(vec![response]);
|
||||
}
|
||||
|
||||
let parsed_query = self.query_parser.compatibility_parser.parse(query);
|
||||
|
||||
// A COPY FROM STDIN trailing other statements must be rejected;
|
||||
// the textual scanner above only sees statement-leading COPYs.
|
||||
if let Ok(statements) = &parsed_query
|
||||
&& statements.iter().any(copy_in::is_copy_from_stdin)
|
||||
{
|
||||
return Err(copy_in::copy_in_must_run_alone_error());
|
||||
}
|
||||
|
||||
let query = if let Ok(statements) = &parsed_query {
|
||||
// Comments, whitespace and empty statements also require EmptyQueryResponse.
|
||||
if statements.is_empty() {
|
||||
@@ -282,6 +302,36 @@ where
|
||||
_ => 0,
|
||||
};
|
||||
|
||||
// Render timestamps with an explicit session-timezone offset so the
|
||||
// exported text carries the instant unambiguously: the timezone-aware
|
||||
// COPY IN decoder honors embedded offsets, keeping COPY OUT -> COPY IN
|
||||
// round trips exact in non-UTC sessions (matching SQL INSERT parity).
|
||||
// This touches both halves of the encoder: the arrow schema drives
|
||||
// value conversion into a tz-aware DateTime, and the pg field type
|
||||
// selects the offset-bearing text format.
|
||||
let arrow_schema_override = copy_out_schema_with_timezone(&schema, &query_ctx);
|
||||
let pg_fields = if arrow_schema_override.is_some() {
|
||||
pg_fields
|
||||
.into_iter()
|
||||
.map(|mut field| {
|
||||
if field.datatype() == &Type::TIMESTAMP {
|
||||
field = FieldInfo::new(
|
||||
field.name().to_string(),
|
||||
None,
|
||||
None,
|
||||
Type::TIMESTAMPTZ,
|
||||
field.format(),
|
||||
)
|
||||
.with_type_size(field.type_size())
|
||||
.with_type_modifier(field.type_modifier());
|
||||
}
|
||||
field
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
} else {
|
||||
pg_fields
|
||||
};
|
||||
|
||||
let pg_schema = Arc::new(pg_fields);
|
||||
let num_columns = pg_schema.len();
|
||||
|
||||
@@ -297,7 +347,8 @@ where
|
||||
schema.clone(),
|
||||
recordbatches_stream,
|
||||
copy_encoder,
|
||||
);
|
||||
)
|
||||
.with_arrow_schema_override(arrow_schema_override);
|
||||
|
||||
let copy_stream: RowStream<CopyData> = Box::pin(
|
||||
row_stream
|
||||
@@ -336,6 +387,7 @@ impl DefaultQueryParser {
|
||||
pub struct PgSqlPlan {
|
||||
pub(crate) plan: SqlPlan,
|
||||
pub(crate) copy_to_stdout_format: Option<String>,
|
||||
pub(crate) copy_from_stdin: Option<copy_in::CopyFromStdin>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -360,6 +412,19 @@ impl QueryParser for DefaultQueryParser {
|
||||
return Ok(Some(PgSqlPlan {
|
||||
plan: SqlPlan::Shortcut(sql.to_string()),
|
||||
copy_to_stdout_format: None,
|
||||
copy_from_stdin: None,
|
||||
}));
|
||||
}
|
||||
|
||||
// A COPY FROM STDIN is parsed textually — not through sqlparser —
|
||||
// so the WITH clause may carry options beyond sqlparser's COPY
|
||||
// grammar; the extended query handler starts the copy sub-protocol
|
||||
// when the statement is executed.
|
||||
if let Some(copy_from_stdin) = copy_in::copy_from_stdin_from_sql(sql)? {
|
||||
return Ok(Some(PgSqlPlan {
|
||||
plan: SqlPlan::Shortcut(sql.to_string()),
|
||||
copy_to_stdout_format: None,
|
||||
copy_from_stdin: Some(copy_from_stdin),
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -370,7 +435,7 @@ impl QueryParser for DefaultQueryParser {
|
||||
}
|
||||
let first_stmt = statements.remove(0);
|
||||
let format = check_copy_to_stdout(&first_stmt);
|
||||
(statement_to_query(&first_stmt), format)
|
||||
(first_stmt.to_string(), format)
|
||||
} else {
|
||||
// bypass the error: it can run into error because of different
|
||||
// versions of sqlparser
|
||||
@@ -400,11 +465,13 @@ impl QueryParser for DefaultQueryParser {
|
||||
Ok(Some(PgSqlPlan {
|
||||
plan: SqlPlan::Plan(logical_plan, stmt),
|
||||
copy_to_stdout_format,
|
||||
copy_from_stdin: None,
|
||||
}))
|
||||
} else {
|
||||
Ok(Some(PgSqlPlan {
|
||||
plan: SqlPlan::Statement(stmt, sql),
|
||||
copy_to_stdout_format,
|
||||
copy_from_stdin: None,
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -460,6 +527,14 @@ impl ExtendedQueryHandler for PostgresServerHandlerInner {
|
||||
let pg_sql_plan = &portal.statement.statement;
|
||||
let sql_plan = &pg_sql_plan.plan;
|
||||
|
||||
// A COPY FROM STDIN takes over the connection for the copy
|
||||
// sub-protocol.
|
||||
if let Some(copy_from_stdin) = &pg_sql_plan.copy_from_stdin {
|
||||
let response = self.begin_copy_in(copy_from_stdin.clone()).await?;
|
||||
send_warning_opt(client, query_ctx).await?;
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
let output = match sql_plan {
|
||||
SqlPlan::Empty => {
|
||||
// early return if query is empty
|
||||
@@ -782,6 +857,38 @@ impl ErrorHandler for PostgresServerHandlerInner {
|
||||
}
|
||||
}
|
||||
|
||||
/// Rewrites timestamp fields of the table's arrow schema to carry the
|
||||
/// session timezone, so COPY OUT renders them with an explicit offset
|
||||
/// (`2024-01-01 08:00:00+08:00`). Returns `None` when there is nothing
|
||||
/// to rewrite (no naive timestamps or a UTC session leaving output as-is
|
||||
/// is fine either way — offsets are always unambiguous on re-import).
|
||||
fn copy_out_schema_with_timezone(
|
||||
schema: &SchemaRef,
|
||||
query_ctx: &QueryContextRef,
|
||||
) -> Option<Arc<arrow::datatypes::Schema>> {
|
||||
let timezone = query_ctx.timezone().to_string();
|
||||
let arrow_schema = schema.arrow_schema();
|
||||
let mut changed = false;
|
||||
let fields = arrow_schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|field| match field.data_type() {
|
||||
arrow::datatypes::DataType::Timestamp(unit, None) => {
|
||||
changed = true;
|
||||
field
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_data_type(arrow::datatypes::DataType::Timestamp(
|
||||
*unit,
|
||||
Some(Arc::from(timezone.as_str())),
|
||||
))
|
||||
}
|
||||
_ => field.as_ref().clone(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
changed.then(|| Arc::new(arrow::datatypes::Schema::new(fields)))
|
||||
}
|
||||
|
||||
fn check_copy_to_stdout(statement: &SqlParserStatement) -> Option<String> {
|
||||
if let SqlParserStatement::Copy {
|
||||
target, options, ..
|
||||
|
||||
@@ -29,6 +29,7 @@ use tokio_rustls::TlsAcceptor;
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder};
|
||||
use crate::query_handler::CopyInHandlerRef;
|
||||
use crate::query_handler::sql::ServerSqlQueryHandlerRef;
|
||||
use crate::server::{AbortableStream, BaseTcpServer, Server};
|
||||
use crate::tls::ReloadableTlsServerConfig;
|
||||
@@ -44,9 +45,16 @@ pub struct PostgresServer {
|
||||
}
|
||||
|
||||
impl PostgresServer {
|
||||
/// Creates a new Postgres server with provided query_handler and async runtime
|
||||
/// Creates a new Postgres server with provided query_handler and async runtime.
|
||||
///
|
||||
/// `query_handler` and `copy_in_handler` are expected to be handles to
|
||||
/// the same frontend instance (the SQL and COPY FROM STDIN facets of
|
||||
/// one object); they are separate parameters only because they are
|
||||
/// separate traits.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
query_handler: ServerSqlQueryHandlerRef,
|
||||
copy_in_handler: CopyInHandlerRef,
|
||||
force_tls: bool,
|
||||
tls_server_config: Arc<ReloadableTlsServerConfig>,
|
||||
keep_alive_secs: u64,
|
||||
@@ -57,6 +65,7 @@ impl PostgresServer {
|
||||
let make_handler = Arc::new(
|
||||
MakePostgresServerHandlerBuilder::default()
|
||||
.query_handler(query_handler.clone())
|
||||
.copy_in_handler(copy_in_handler)
|
||||
.user_provider(user_provider.clone())
|
||||
.force_tls(force_tls)
|
||||
.build()
|
||||
|
||||
@@ -285,6 +285,11 @@ where
|
||||
schema: SchemaRef,
|
||||
record_batches: Pin<Box<B>>,
|
||||
encoder: S,
|
||||
/// When set, replaces each batch's arrow schema for value encoding.
|
||||
/// The COPY OUT path uses this to render timestamps with an explicit
|
||||
/// session-timezone offset, so exported text re-imports identically
|
||||
/// by the timezone-aware COPY IN decoder.
|
||||
arrow_schema_override: Option<Arc<arrow::datatypes::Schema>>,
|
||||
}
|
||||
|
||||
impl<S, B> Stream for RecordBatchRowStream<S, B>
|
||||
@@ -304,7 +309,10 @@ where
|
||||
return Poll::Ready(Some(Ok(vec![])));
|
||||
}
|
||||
|
||||
let arrow_schema = record_batch.schema();
|
||||
let arrow_schema = self
|
||||
.arrow_schema_override
|
||||
.clone()
|
||||
.unwrap_or_else(|| record_batch.schema());
|
||||
let query_ctx = self.query_ctx.clone();
|
||||
let pg_schema = self.pg_schema.clone();
|
||||
let schema = self.schema.clone();
|
||||
@@ -352,9 +360,19 @@ where
|
||||
schema,
|
||||
record_batches: Box::pin(record_batches),
|
||||
encoder,
|
||||
arrow_schema_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// See [`RecordBatchRowStream::arrow_schema_override`].
|
||||
pub(crate) fn with_arrow_schema_override(
|
||||
mut self,
|
||||
schema: Option<Arc<arrow::datatypes::Schema>>,
|
||||
) -> Self {
|
||||
self.arrow_schema_override = schema;
|
||||
self
|
||||
}
|
||||
|
||||
fn encode_row(
|
||||
query_ctx: &QueryContextRef,
|
||||
pg_schema: &Arc<Vec<FieldInfo>>,
|
||||
@@ -2464,6 +2482,7 @@ mod test {
|
||||
PgSqlPlan {
|
||||
plan: SqlPlan::Empty,
|
||||
copy_to_stdout_format: None,
|
||||
copy_from_stdin: None,
|
||||
},
|
||||
client_param_types,
|
||||
));
|
||||
|
||||
@@ -67,6 +67,9 @@ pub enum PgErrorCode {
|
||||
/// connection_exception
|
||||
#[snafu(display("connection_exception"))]
|
||||
Ec08000 = 400,
|
||||
/// protocol_violation
|
||||
#[snafu(display("protocol_violation"))]
|
||||
Ec08P01 = 401,
|
||||
// === End of Class 08 — Connection Exception =====
|
||||
|
||||
// === Begin of Class 09 — Triggered Action Exception ===
|
||||
@@ -130,6 +133,12 @@ pub enum PgErrorCode {
|
||||
/// invalid_parameter_value
|
||||
#[snafu(display("invalid_parameter_value"))]
|
||||
Ec22023 = 1401,
|
||||
/// invalid_text_representation
|
||||
#[snafu(display("invalid_text_representation"))]
|
||||
Ec22P02 = 1402,
|
||||
/// bad_copy_file_format
|
||||
#[snafu(display("bad_copy_file_format"))]
|
||||
Ec22P04 = 1403,
|
||||
// === End of Class 22 — Data Exception =====
|
||||
|
||||
// === Begin of Class 23 — Integrity Constraint Violation ===
|
||||
@@ -295,6 +304,9 @@ pub enum PgErrorCode {
|
||||
/// operator_intervention
|
||||
#[snafu(display("operator_intervention"))]
|
||||
Ec57000 = 3600,
|
||||
/// query_canceled
|
||||
#[snafu(display("query_canceled"))]
|
||||
Ec57014 = 3602,
|
||||
|
||||
/// cannot_connect_now
|
||||
#[snafu(display("cannot_connect_now"))]
|
||||
|
||||
@@ -43,6 +43,7 @@ use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetrics
|
||||
use pipeline::{GreptimePipelineParams, Pipeline, PipelineInfo, PipelineVersion, PipelineWay};
|
||||
use serde_json::Value;
|
||||
use session::context::{QueryContext, QueryContextRef};
|
||||
use table::metadata::TableInfoRef;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct DashboardDefinition {
|
||||
@@ -290,3 +291,33 @@ pub trait JaegerQueryHandler {
|
||||
query_params: QueryTraceParams,
|
||||
) -> Result<Output>;
|
||||
}
|
||||
|
||||
/// Handler for the PostgreSQL `COPY ... FROM STDIN` (copy-in) sub-protocol.
|
||||
///
|
||||
/// The servers layer owns the wire protocol and data parsing; this trait is
|
||||
/// the boundary to catalog lookup and ingestion, implemented by the
|
||||
/// frontend instance.
|
||||
#[async_trait]
|
||||
pub trait CopyInHandler: Send + Sync {
|
||||
/// Resolves the target table of a COPY and returns its metadata, or
|
||||
/// `None` when the table does not exist. Implementations should also
|
||||
/// enforce the write permission for the table here so unauthorized
|
||||
/// copies fail before any data is transferred.
|
||||
async fn copy_in_table(
|
||||
&self,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Option<TableInfoRef>>;
|
||||
|
||||
/// Inserts a parsed batch produced by a COPY stream. Returns the
|
||||
/// number of affected rows.
|
||||
async fn copy_in_insert(
|
||||
&self,
|
||||
requests: RowInsertRequests,
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Output>;
|
||||
}
|
||||
|
||||
pub type CopyInHandlerRef = Arc<dyn CopyInHandler + Send + Sync>;
|
||||
|
||||
@@ -14,13 +14,14 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::v1::RowInsertRequests;
|
||||
use api::v1::greptime_request::Request;
|
||||
use api::v1::query_request::Query;
|
||||
use async_trait::async_trait;
|
||||
use catalog::memory::MemoryCatalogManager;
|
||||
use catalog::system_schema::SystemSchemaProvider;
|
||||
use catalog::system_schema::pg_catalog::PGCatalogProvider;
|
||||
use catalog::{CatalogManager, RegisterTableRequest};
|
||||
use catalog::{CatalogManager, CatalogManagerRef, RegisterTableRequest};
|
||||
use common_catalog::consts::{DEFAULT_CATALOG_NAME, PG_CATALOG_NAME};
|
||||
use common_error::ext::BoxedError;
|
||||
use common_grpc::flight::do_put::DoPutResponse;
|
||||
@@ -32,6 +33,7 @@ use query::parser::{PromQuery, QueryStatement};
|
||||
use query::query_engine::DescribeResult;
|
||||
use query::{QueryEngineFactory, QueryEngineRef};
|
||||
use servers::error::{ExecuteQuerySnafu, NotSupportedSnafu, Result};
|
||||
use servers::query_handler::CopyInHandler;
|
||||
use servers::query_handler::grpc::GrpcQueryHandler;
|
||||
use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler};
|
||||
use session::context::QueryContextRef;
|
||||
@@ -39,6 +41,7 @@ use snafu::{ResultExt, ensure};
|
||||
use sql::parser::{ParseOptions, ParserContext};
|
||||
use sql::statements::statement::Statement;
|
||||
use table::TableRef;
|
||||
use table::metadata::TableInfoRef;
|
||||
|
||||
mod http;
|
||||
mod interceptor;
|
||||
@@ -54,6 +57,64 @@ impl DummyInstance {
|
||||
fn new(query_engine: QueryEngineRef) -> Self {
|
||||
Self { query_engine }
|
||||
}
|
||||
|
||||
pub fn catalog_manager(&self) -> CatalogManagerRef {
|
||||
self.query_engine.engine_state().catalog_manager().clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// A [`CopyInHandler`] test double: resolves tables through the catalog
|
||||
/// manager and records the insert requests instead of writing them.
|
||||
pub struct RecordingCopyInHandler {
|
||||
catalog_manager: CatalogManagerRef,
|
||||
pub requests: std::sync::Mutex<Vec<RowInsertRequests>>,
|
||||
}
|
||||
|
||||
impl RecordingCopyInHandler {
|
||||
pub fn new(catalog_manager: CatalogManagerRef) -> Self {
|
||||
Self {
|
||||
catalog_manager,
|
||||
requests: Default::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CopyInHandler for RecordingCopyInHandler {
|
||||
async fn copy_in_table(
|
||||
&self,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
_query_ctx: QueryContextRef,
|
||||
) -> Result<Option<TableInfoRef>> {
|
||||
Ok(self
|
||||
.catalog_manager
|
||||
.table(catalog, schema, table, None)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|table| table.table_info()))
|
||||
}
|
||||
|
||||
async fn copy_in_insert(
|
||||
&self,
|
||||
requests: RowInsertRequests,
|
||||
_query_ctx: QueryContextRef,
|
||||
) -> Result<Output> {
|
||||
let rows: usize = requests
|
||||
.inserts
|
||||
.iter()
|
||||
.map(|insert| {
|
||||
insert
|
||||
.rows
|
||||
.as_ref()
|
||||
.map(|rows| rows.rows.len())
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.sum();
|
||||
self.requests.lock().unwrap().push(requests);
|
||||
Ok(Output::new_with_affected_rows(rows))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -229,7 +229,18 @@ fn create_postgres_server_with_user_provider(
|
||||
tls: TlsOption,
|
||||
user_provider: Option<UserProviderRef>,
|
||||
) -> Result<Box<dyn Server>> {
|
||||
create_postgres_server_inner(table, tls, user_provider).map(|(server, _)| server)
|
||||
}
|
||||
|
||||
fn create_postgres_server_inner(
|
||||
table: TableRef,
|
||||
tls: TlsOption,
|
||||
user_provider: Option<UserProviderRef>,
|
||||
) -> Result<(Box<dyn Server>, Arc<crate::RecordingCopyInHandler>)> {
|
||||
let instance = Arc::new(create_testing_instance(table));
|
||||
let copy_in_handler = Arc::new(crate::RecordingCopyInHandler::new(
|
||||
instance.catalog_manager(),
|
||||
));
|
||||
let io_runtime = RuntimeBuilder::default()
|
||||
.worker_threads(4)
|
||||
.thread_name("postgres-io-handlers")
|
||||
@@ -241,15 +252,19 @@ fn create_postgres_server_with_user_provider(
|
||||
.expect("Failed to load certificates and keys"),
|
||||
);
|
||||
|
||||
Ok(Box::new(PostgresServer::new(
|
||||
instance,
|
||||
tls.should_force_tls(),
|
||||
tls_server_config,
|
||||
0,
|
||||
io_runtime,
|
||||
user_provider,
|
||||
None,
|
||||
)))
|
||||
Ok((
|
||||
Box::new(PostgresServer::new(
|
||||
instance,
|
||||
copy_in_handler.clone(),
|
||||
tls.should_force_tls(),
|
||||
tls_server_config,
|
||||
0,
|
||||
io_runtime,
|
||||
user_provider,
|
||||
None,
|
||||
)),
|
||||
copy_in_handler,
|
||||
))
|
||||
}
|
||||
|
||||
async fn start_test_server_with_user_provider(
|
||||
@@ -1149,3 +1164,237 @@ async fn do_simple_query_with_secure_server(
|
||||
|
||||
do_simple_query(server_tls, client_tls).await
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// COPY FROM STDIN tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
use futures::SinkExt;
|
||||
|
||||
fn copy_in_test_table() -> TableRef {
|
||||
use common_recordbatch::RecordBatch;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::schema::ColumnSchema;
|
||||
use datatypes::types::TimestampType;
|
||||
use datatypes::vectors::{Float64Vector, StringVector, TimestampMillisecondVector, VectorRef};
|
||||
|
||||
let column_schemas = vec![
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::Timestamp(TimestampType::Millisecond(Default::default())),
|
||||
false,
|
||||
)
|
||||
.with_time_index(true),
|
||||
ColumnSchema::new("host", ConcreteDataType::string_datatype(), false),
|
||||
ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true),
|
||||
ColumnSchema::new("note", ConcreteDataType::string_datatype(), true),
|
||||
];
|
||||
let schema = Arc::new(datatypes::schema::Schema::new(column_schemas));
|
||||
let columns: Vec<VectorRef> = vec![
|
||||
Arc::new(TimestampMillisecondVector::from_slice([0i64])),
|
||||
Arc::new(StringVector::from_slice(&["localhost"])),
|
||||
Arc::new(Float64Vector::from_slice([1.0f64])),
|
||||
Arc::new(StringVector::from_slice(&["seed"])),
|
||||
];
|
||||
let recordbatch = RecordBatch::new(schema, columns).unwrap();
|
||||
MemTable::table("metrics", recordbatch)
|
||||
}
|
||||
|
||||
async fn start_copy_in_test_server(
|
||||
table: TableRef,
|
||||
) -> Result<(Box<dyn Server>, u16, Arc<crate::RecordingCopyInHandler>)> {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
let _ = install_default_crypto_provider();
|
||||
let (mut server, handler) = create_postgres_server_inner(table, Default::default(), None)?;
|
||||
let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
|
||||
server.start(listening).await.unwrap();
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
Ok((server, port, handler))
|
||||
}
|
||||
|
||||
fn recorded_rows(handler: &crate::RecordingCopyInHandler) -> Vec<api::v1::Row> {
|
||||
handler
|
||||
.requests
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flat_map(|request| {
|
||||
request.inserts.iter().flat_map(|insert| {
|
||||
insert
|
||||
.rows
|
||||
.clone()
|
||||
.map(|rows| rows.rows)
|
||||
.unwrap_or_default()
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn feed_copy_in(
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
chunks: Vec<bytes::Bytes>,
|
||||
) -> std::result::Result<u64, PgError> {
|
||||
let mut sink = Box::pin(client.copy_in(sql).await?);
|
||||
for chunk in chunks {
|
||||
sink.as_mut().send(chunk).await?;
|
||||
}
|
||||
sink.as_mut().finish().await
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_copy_in_csv() -> Result<()> {
|
||||
let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?;
|
||||
let client = create_plain_connection(port, false).await.unwrap();
|
||||
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY metrics FROM STDIN WITH (FORMAT csv)",
|
||||
vec![bytes::Bytes::from_static(
|
||||
b"2023-11-14 22:13:20,host1,1.5,hello\n2023-11-14 22:13:21,host2,,\"quoted, comma\"\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 2);
|
||||
|
||||
let rows = recorded_rows(&handler);
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(
|
||||
rows[0].values[1].value_data,
|
||||
Some(api::v1::value::ValueData::StringValue("host1".to_string()))
|
||||
);
|
||||
assert_eq!(
|
||||
rows[0].values[2].value_data,
|
||||
Some(api::v1::value::ValueData::F64Value(1.5))
|
||||
);
|
||||
// Unquoted empty field is NULL, quoted field keeps its content.
|
||||
assert_eq!(rows[1].values[2].value_data, None);
|
||||
assert_eq!(
|
||||
rows[1].values[3].value_data,
|
||||
Some(api::v1::value::ValueData::StringValue(
|
||||
"quoted, comma".to_string()
|
||||
))
|
||||
);
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_copy_in_column_list_and_batches() -> Result<()> {
|
||||
let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?;
|
||||
let client = create_plain_connection(port, false).await.unwrap();
|
||||
|
||||
let mut chunks = Vec::new();
|
||||
let mut payload = String::new();
|
||||
for i in 0..70000u64 {
|
||||
payload.push_str(&format!("{i},2023-11-14 22:13:20,host{i},0.5\n"));
|
||||
if (i + 1) % 100 == 0 {
|
||||
chunks.push(bytes::Bytes::from(std::mem::take(&mut payload)));
|
||||
}
|
||||
}
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY metrics (note, ts, host, val) FROM STDIN (FORMAT csv)",
|
||||
chunks,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 70000);
|
||||
|
||||
let recorded = handler.requests.lock().unwrap().clone();
|
||||
// The 70000 rows must have been flushed in multiple batches of at
|
||||
// most COPY_IN_FLUSH_ROWS.
|
||||
assert!(recorded.len() > 1);
|
||||
let total: usize = recorded
|
||||
.iter()
|
||||
.map(|request| {
|
||||
request.inserts[0]
|
||||
.rows
|
||||
.as_ref()
|
||||
.map(|rows| rows.rows.len())
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.sum();
|
||||
assert_eq!(total, 70000);
|
||||
|
||||
let rows = recorded_rows(&handler);
|
||||
// The column list reorders the fields: first field lands in `note`.
|
||||
assert_eq!(
|
||||
rows[0].values[0].value_data,
|
||||
Some(api::v1::value::ValueData::StringValue("0".to_string()))
|
||||
);
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_copy_in_unknown_table() -> Result<()> {
|
||||
let (server, port, _handler) = start_copy_in_test_server(copy_in_test_table()).await?;
|
||||
let client = create_plain_connection(port, false).await.unwrap();
|
||||
|
||||
let err = match client
|
||||
.copy_in::<_, bytes::Bytes>("COPY not_exist FROM STDIN WITH (FORMAT csv)")
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("expected COPY to unknown table to fail"),
|
||||
Err(e) => e,
|
||||
};
|
||||
let message = err.as_db_error().expect("expected db error").message();
|
||||
assert!(message.contains("does not exist"), "{message}");
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_copy_in_bad_value() -> Result<()> {
|
||||
let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?;
|
||||
let client = create_plain_connection(port, false).await.unwrap();
|
||||
|
||||
let err = match feed_copy_in(
|
||||
&client,
|
||||
"COPY metrics FROM STDIN WITH (FORMAT csv)",
|
||||
vec![bytes::Bytes::from_static(b"not_a_timestamp,host1,1.5,x\n")],
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("expected COPY to fail"),
|
||||
Err(e) => e,
|
||||
};
|
||||
let message = err.as_db_error().expect("expected db error").message();
|
||||
assert!(message.contains("failed to parse"), "{message}");
|
||||
// The failed batch must not be written.
|
||||
assert!(handler.requests.lock().unwrap().is_empty());
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_copy_in_multi_statement_rejected() -> Result<()> {
|
||||
let (server, port, _handler) = start_copy_in_test_server(copy_in_test_table()).await?;
|
||||
let client = create_plain_connection(port, false).await.unwrap();
|
||||
|
||||
let err = match feed_copy_in(
|
||||
&client,
|
||||
"COPY metrics FROM STDIN WITH (FORMAT csv); SELECT 1",
|
||||
vec![bytes::Bytes::from_static(b"1,2,3,4\n")],
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("expected multi-statement COPY to fail"),
|
||||
Err(e) => e,
|
||||
};
|
||||
let message = err
|
||||
.as_db_error()
|
||||
.expect("expected db error")
|
||||
.message()
|
||||
.to_string();
|
||||
assert!(message.contains("alone"), "{message}");
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -77,6 +77,7 @@ uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
catalog = { workspace = true, features = ["testing"] }
|
||||
chrono.workspace = true
|
||||
clap.workspace = true
|
||||
|
||||
@@ -1205,7 +1205,8 @@ pub async fn setup_pg_server_with_slow_query_threshold(
|
||||
);
|
||||
|
||||
let mut pg_server = Box::new(PostgresServer::new(
|
||||
fe_instance_ref,
|
||||
fe_instance_ref.clone(),
|
||||
fe_instance_ref.clone(),
|
||||
opts.tls.should_force_tls(),
|
||||
tls_server_config,
|
||||
0,
|
||||
@@ -1248,7 +1249,8 @@ pub async fn setup_pg_server_with_user_provider(
|
||||
);
|
||||
|
||||
let mut pg_server = Box::new(PostgresServer::new(
|
||||
fe_instance_ref,
|
||||
fe_instance_ref.clone(),
|
||||
fe_instance_ref.clone(),
|
||||
opts.tls.should_force_tls(),
|
||||
tls_server_config,
|
||||
0,
|
||||
@@ -1303,6 +1305,7 @@ pub async fn setup_pg_server_with_prom_native_histogram(
|
||||
);
|
||||
|
||||
let mut pg_server = Box::new(PostgresServer::new(
|
||||
instance.fe_instance().clone(),
|
||||
instance.fe_instance().clone(),
|
||||
opts.tls.should_force_tls(),
|
||||
tls_server_config,
|
||||
|
||||
@@ -17,6 +17,7 @@ mod instance_kafka_wal_test;
|
||||
mod instance_noop_wal_test;
|
||||
mod instance_test;
|
||||
mod nested_merge_scan_capability_test;
|
||||
mod postgres_copy_in_test;
|
||||
mod promql_test;
|
||||
mod reconcile_table;
|
||||
mod remote_dyn_filter_test;
|
||||
|
||||
@@ -0,0 +1,521 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//! End-to-end tests for `COPY ... FROM STDIN` over the PostgreSQL protocol,
|
||||
//! covering the full ingestion path (parser -> insert -> storage).
|
||||
|
||||
use bytes::Bytes;
|
||||
use common_telemetry::init_default_ut_logging;
|
||||
use futures::SinkExt;
|
||||
use tokio_postgres::NoTls;
|
||||
|
||||
use crate::test_util::{StorageType, setup_pg_server};
|
||||
|
||||
async fn connect(port: u16) -> tokio_postgres::Client {
|
||||
let url = format!("host=127.0.0.1 port={port} dbname=public connect_timeout=5");
|
||||
let (client, conn) = tokio_postgres::connect(&url, NoTls).await.unwrap();
|
||||
let _handle = tokio::spawn(conn);
|
||||
client
|
||||
}
|
||||
|
||||
async fn feed_copy_in(
|
||||
client: &tokio_postgres::Client,
|
||||
sql: &str,
|
||||
chunks: Vec<Bytes>,
|
||||
) -> std::result::Result<u64, tokio_postgres::Error> {
|
||||
let mut sink = Box::pin(client.copy_in(sql).await?);
|
||||
for chunk in chunks {
|
||||
sink.as_mut().send(chunk).await?;
|
||||
}
|
||||
sink.as_mut().finish().await
|
||||
}
|
||||
|
||||
fn host_of(message: &tokio_postgres::SimpleQueryMessage, column: usize) -> String {
|
||||
let tokio_postgres::SimpleQueryMessage::Row(row) = message else {
|
||||
panic!("expected row, got {message:?}");
|
||||
};
|
||||
// simple_query renders NULL as `None`, not empty string.
|
||||
row.get(column).map(str::to_string).unwrap_or_default()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_end_to_end() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_e2e").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE copy_demo (host STRING, val DOUBLE, note STRING, ts TIMESTAMP, TIME INDEX(ts))",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// CSV format with quoted field and NULL (unquoted empty) value.
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY copy_demo FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(
|
||||
b"host1,1.5,\"hello, world\",2023-11-14 22:13:20\nhost2,,x,2023-11-14 22:13:21\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 2);
|
||||
|
||||
// A column list targeting a subset.
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY copy_demo (note, ts) FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(
|
||||
b"from_column_list,2023-11-14 22:13:24\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 1);
|
||||
|
||||
// Verify data written through the real insert path.
|
||||
let messages = client
|
||||
.simple_query("SELECT host, val, note FROM copy_demo ORDER BY ts")
|
||||
.await
|
||||
.unwrap();
|
||||
let rows: Vec<_> = messages
|
||||
.iter()
|
||||
.filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_)))
|
||||
.collect();
|
||||
assert_eq!(rows.len(), 3);
|
||||
assert_eq!(host_of(rows[0], 0), "host1");
|
||||
assert_eq!(host_of(rows[0], 2), "hello, world");
|
||||
assert_eq!(host_of(rows[1], 0), "host2");
|
||||
// NULL val and note round-trip as NULL.
|
||||
assert_eq!(host_of(rows[2], 0), "");
|
||||
assert_eq!(host_of(rows[2], 2), "from_column_list");
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_error_keeps_table_clean() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_error").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE copy_err (host STRING, val DOUBLE, ts TIMESTAMP, TIME INDEX(ts))",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Bad double value aborts the copy.
|
||||
let err = match feed_copy_in(
|
||||
&client,
|
||||
"COPY copy_err FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(
|
||||
b"host1,not_a_double,2023-11-14 22:13:20\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("expected COPY to fail"),
|
||||
Err(e) => e,
|
||||
};
|
||||
assert!(
|
||||
err.as_db_error()
|
||||
.expect("expected db error")
|
||||
.message()
|
||||
.contains("failed to parse"),
|
||||
"{err}"
|
||||
);
|
||||
|
||||
// A bare COPY (PostgreSQL's text default) fails with guidance.
|
||||
let err = match client.copy_in::<_, Bytes>("COPY copy_err FROM STDIN").await {
|
||||
Ok(_) => panic!("expected bare COPY to fail"),
|
||||
Err(e) => e,
|
||||
};
|
||||
let message = err
|
||||
.as_db_error()
|
||||
.expect("expected db error")
|
||||
.message()
|
||||
.to_string();
|
||||
assert!(message.contains("requires an explicit format"), "{message}");
|
||||
|
||||
// The connection stays usable and no rows were written.
|
||||
let messages = client
|
||||
.simple_query("SELECT host FROM copy_err")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
messages
|
||||
.iter()
|
||||
.all(|m| !matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))),
|
||||
"expected no rows, got {messages:?}"
|
||||
);
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// Exercises the CSV format's options end-to-end through a Postgres
|
||||
/// client driver: HEADER, custom delimiter, quoted fields with embedded
|
||||
/// newlines and escaped quotes, CRLF record separators and NULL fields.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_csv_options() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_csv_options").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE copy_csv_opts (host STRING, val DOUBLE, note STRING, ts TIMESTAMP, TIME INDEX(ts))",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
// WITH omitted, matching PostgreSQL; HEADER skips the first record.
|
||||
"COPY copy_csv_opts FROM STDIN (FORMAT csv, HEADER, DELIMITER ';')",
|
||||
vec![Bytes::from_static(
|
||||
b"host;val;note;ts\r\n\
|
||||
h1;1.5;\"line one\nline two\";2023-11-14 22:13:20\r\n\
|
||||
h2;;\"say \"\"hi\"\"\";2023-11-14 22:13:21\r\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 2);
|
||||
|
||||
// Verify the parsed values came through the real insert path.
|
||||
let messages = client
|
||||
.simple_query("SELECT host, val, note FROM copy_csv_opts ORDER BY ts")
|
||||
.await
|
||||
.unwrap();
|
||||
let rows: Vec<_> = messages
|
||||
.iter()
|
||||
.filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_)))
|
||||
.collect();
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(host_of(rows[0], 0), "h1");
|
||||
assert_eq!(host_of(rows[0], 1), "1.5");
|
||||
// Quoted field with an embedded newline survives the round trip.
|
||||
assert_eq!(host_of(rows[0], 2), "line one\nline two");
|
||||
assert_eq!(host_of(rows[1], 0), "h2");
|
||||
// Unquoted empty field is NULL; escaped quotes decode to one quote.
|
||||
assert_eq!(host_of(rows[1], 1), "");
|
||||
assert_eq!(host_of(rows[1], 2), "say \"hi\"");
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// COPY OUT → COPY IN round trip: PostgreSQL's COPY OUT emits `t`/`f`
|
||||
/// for booleans, which COPY IN must accept back.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_bool_round_trip() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_bool_rt").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE bool_src (flag BOOLEAN, ts TIMESTAMP, TIME INDEX(ts));
|
||||
CREATE TABLE bool_dst (flag BOOLEAN, ts TIMESTAMP, TIME INDEX(ts));
|
||||
INSERT INTO bool_src VALUES (true, '2024-01-01 00:00:00'), (false, '2024-01-01 00:00:01'), (true, '2024-01-01 00:00:02');",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// COPY OUT: collect the exported bytes.
|
||||
let mut exported = Vec::new();
|
||||
let mut out = Box::pin(
|
||||
client
|
||||
.copy_out(
|
||||
"COPY (SELECT flag, ts FROM bool_src ORDER BY ts) TO STDOUT WITH (FORMAT csv)",
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk) = out.as_mut().next().await {
|
||||
exported.extend_from_slice(&chunk.unwrap());
|
||||
}
|
||||
// The wire encoder emits PostgreSQL boolean spellings.
|
||||
let text = String::from_utf8(exported.clone()).unwrap();
|
||||
assert!(text.contains("t,"), "COPY OUT should emit t/f, got: {text}");
|
||||
assert!(text.contains("f,"), "COPY OUT should emit t/f, got: {text}");
|
||||
|
||||
// COPY IN: the exact bytes must import back.
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY bool_dst FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from(exported)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 3);
|
||||
|
||||
let messages = client
|
||||
.simple_query("SELECT flag FROM bool_dst ORDER BY ts")
|
||||
.await
|
||||
.unwrap();
|
||||
let flags: Vec<&str> = messages
|
||||
.iter()
|
||||
.filter_map(|m| match m {
|
||||
tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
// The simple-query path also renders booleans in the t/f spelling.
|
||||
assert_eq!(flags, vec!["t", "f", "t"]);
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// A failed COPY must release its state so a later COPY on the same
|
||||
/// connection works; this guards the error-path state drop.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_failed_copy_then_new_copy() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_fail_reuse").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE copy_reuse (host STRING, val DOUBLE, ts TIMESTAMP, TIME INDEX(ts))",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// A copy that fails at CopyDone (bad double parsed at flush time).
|
||||
let err = match feed_copy_in(
|
||||
&client,
|
||||
"COPY copy_reuse FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(
|
||||
b"host1,not_a_double,2024-01-01 00:00:00\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("expected COPY to fail"),
|
||||
Err(e) => e,
|
||||
};
|
||||
assert!(err.as_db_error().is_some(), "{err}");
|
||||
|
||||
// The connection is usable and a NEW copy succeeds on it.
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY copy_reuse FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(b"host1,1.5,2024-01-01 00:00:01\n")],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 1);
|
||||
|
||||
let messages = client
|
||||
.simple_query("SELECT host, val FROM copy_reuse")
|
||||
.await
|
||||
.unwrap();
|
||||
let found: Vec<_> = messages
|
||||
.iter()
|
||||
.filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_)))
|
||||
.collect();
|
||||
assert_eq!(found.len(), 1);
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// COPY/INSERT parity in a non-UTC session: offset-free timestamp strings
|
||||
/// must be interpreted in the session timezone by both paths.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_session_timezone_parity() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_tz").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE copy_tz (case_id STRING PRIMARY KEY, ts TIMESTAMP, TIME INDEX(ts))",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
client
|
||||
.simple_query("SET TIMEZONE TO 'Asia/Shanghai'")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
client
|
||||
.simple_query("INSERT INTO copy_tz VALUES ('insert', '2024-01-01 00:00:00')")
|
||||
.await
|
||||
.unwrap();
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY copy_tz (case_id, ts) FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(b"copy,2024-01-01 00:00:00\n")],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 1);
|
||||
|
||||
let messages = client
|
||||
.simple_query("SELECT case_id, CAST(ts AS BIGINT) FROM copy_tz ORDER BY case_id")
|
||||
.await
|
||||
.unwrap();
|
||||
let results: Vec<(String, String)> = messages
|
||||
.iter()
|
||||
.filter_map(|m| match m {
|
||||
tokio_postgres::SimpleQueryMessage::Row(row) => {
|
||||
Some((row.get(0)?.to_string(), row.get(1)?.to_string()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(results.len(), 2, "{results:?}");
|
||||
assert_eq!(results[0].0, "copy");
|
||||
assert_eq!(results[1].0, "insert");
|
||||
// The whole point: both paths store the same instant.
|
||||
assert_eq!(
|
||||
results[0].1, results[1].1,
|
||||
"COPY and INSERT must agree on offset-free timestamps: {results:?}"
|
||||
);
|
||||
// And it is the Asia/Shanghai interpretation, not UTC.
|
||||
assert_eq!(results[0].1, "1704038400000");
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// COPY OUT -> COPY IN must preserve timestamp instants in non-UTC
|
||||
/// sessions: COPY OUT renders timestamps with an explicit session-tz
|
||||
/// offset, which the timezone-aware COPY IN decoder honors.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_out_in_timestamp_round_trip() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_out_in_tz").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query(
|
||||
"CREATE TABLE tz_src (host STRING, ts TIMESTAMP, TIME INDEX(ts));
|
||||
CREATE TABLE tz_dst (host STRING, ts TIMESTAMP, TIME INDEX(ts));
|
||||
INSERT INTO tz_src VALUES ('a', '2024-01-01 00:00:00'), ('b', '2024-07-01 12:30:00');",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
client
|
||||
.simple_query("SET TIMEZONE TO 'Asia/Shanghai'")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut exported = Vec::new();
|
||||
let mut out = Box::pin(
|
||||
client
|
||||
.copy_out("COPY (SELECT host, ts FROM tz_src ORDER BY ts) TO STDOUT WITH (FORMAT csv)")
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk) = out.as_mut().next().await {
|
||||
exported.extend_from_slice(&chunk.unwrap());
|
||||
}
|
||||
// The export now carries explicit offsets.
|
||||
let text = String::from_utf8(exported.clone()).unwrap();
|
||||
assert!(
|
||||
text.contains("+08"),
|
||||
"COPY OUT should render timestamps with the session-tz offset, got: {text}"
|
||||
);
|
||||
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY tz_dst (host, ts) FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from(exported)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 2);
|
||||
|
||||
// Epochs must match the source exactly.
|
||||
let compare = client
|
||||
.simple_query(
|
||||
"SELECT (SELECT CAST(ts AS BIGINT) FROM tz_src WHERE host='a') = (SELECT CAST(ts AS BIGINT) FROM tz_dst WHERE host='a') \
|
||||
AND (SELECT CAST(ts AS BIGINT) FROM tz_src WHERE host='b') = (SELECT CAST(ts AS BIGINT) FROM tz_dst WHERE host='b')",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let ok: Vec<_> = compare
|
||||
.iter()
|
||||
.filter_map(|m| match m {
|
||||
tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(ok, vec!["t"], "epochs must round-trip exactly");
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
/// Empty fields — quoted or not — import as NULL, matching the
|
||||
/// file-based COPY ... FROM 'file.csv' behavior exactly.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_postgres_copy_in_empty_field_is_null() {
|
||||
init_default_ut_logging();
|
||||
let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_empty").await;
|
||||
let port = server.bind_addr().unwrap().port();
|
||||
let client = connect(port).await;
|
||||
|
||||
client
|
||||
.simple_query("CREATE TABLE empty_str (note STRING, ts TIMESTAMP, TIME INDEX(ts))")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let rows = feed_copy_in(
|
||||
&client,
|
||||
"COPY empty_str (note, ts) FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(
|
||||
b"plain,2024-01-01 00:00:00\n\"\",2024-01-01 00:00:01\n,2024-01-01 00:00:02\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows, 3);
|
||||
|
||||
let messages = client
|
||||
.simple_query("SELECT note IS NULL AS is_null, note FROM empty_str ORDER BY ts")
|
||||
.await
|
||||
.unwrap();
|
||||
let rows: Vec<(bool, Option<String>)> = messages
|
||||
.iter()
|
||||
.filter_map(|m| match m {
|
||||
tokio_postgres::SimpleQueryMessage::Row(row) => {
|
||||
Some((row.get(0) == Some("t"), row.get(1).map(str::to_string)))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(rows.len(), 3);
|
||||
assert_eq!(rows[0], (false, Some("plain".to_string())));
|
||||
// Quoted empty string: NULL, exactly like the file-based COPY route.
|
||||
assert_eq!(rows[1], (true, None));
|
||||
// Unquoted empty field: NULL.
|
||||
assert_eq!(rows[2], (true, None));
|
||||
|
||||
server.shutdown().await.unwrap();
|
||||
}
|
||||
@@ -2908,3 +2908,112 @@ async fn test_sql_batcher_alignment() {
|
||||
instance.guard.remove_all().await;
|
||||
}
|
||||
}
|
||||
|
||||
/// COPY FROM STDIN must route through the pending-rows batcher (with
|
||||
/// synchronous acknowledgement) when PostgreSQL batching is enabled, and
|
||||
/// keep direct writes when it is not.
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_postgres_copy_in_with_batching() {
|
||||
use bytes::Bytes;
|
||||
use common_telemetry::{dump_metrics, init_default_ut_logging};
|
||||
use frontend::server::Services;
|
||||
use frontend::service_config::BatcherOptions;
|
||||
use futures::SinkExt;
|
||||
use servers::batcher::BatchingProtocol;
|
||||
use tests_integration::standalone::GreptimeDbStandaloneBuilder;
|
||||
|
||||
fn flushes() -> u64 {
|
||||
dump_metrics()
|
||||
.unwrap()
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
line.strip_prefix("greptime_table_batcher_flush_total ")
|
||||
.map(|value| value.parse().unwrap())
|
||||
})
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
async fn feed_copy_in(
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
chunks: Vec<Bytes>,
|
||||
) -> std::result::Result<u64, tokio_postgres::Error> {
|
||||
let mut sink = Box::pin(client.copy_in(sql).await?);
|
||||
for chunk in chunks {
|
||||
sink.as_mut().send(chunk).await?;
|
||||
}
|
||||
sink.as_mut().finish().await
|
||||
}
|
||||
|
||||
init_default_ut_logging();
|
||||
for pg_enabled in [false, true] {
|
||||
let protocols = if pg_enabled {
|
||||
vec![BatchingProtocol::Postgres]
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
let mut instance = GreptimeDbStandaloneBuilder::new("pg_copy_in_batching")
|
||||
.with_table_batcher(BatcherOptions {
|
||||
protocols,
|
||||
pending_rows_flush_interval: Duration::from_millis(10),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.await;
|
||||
let mut opts = instance.opts.clone();
|
||||
opts.http.addr = "127.0.0.1:0".into();
|
||||
opts.grpc.bind_addr = "127.0.0.1:0".into();
|
||||
opts.postgres.addr = "127.0.0.1:0".into();
|
||||
let mut servers = Services::new(opts, instance.fe_instance().clone(), Default::default())
|
||||
.build()
|
||||
.unwrap();
|
||||
servers.start_all().await.unwrap();
|
||||
let (pg, connection) = tokio_postgres::connect(
|
||||
&format!(
|
||||
"postgres://{}/public",
|
||||
servers.addr("POSTGRES_SERVER").unwrap()
|
||||
),
|
||||
NoTls,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let pg_task = tokio::spawn(async move { connection.await.unwrap() });
|
||||
|
||||
pg.simple_query(
|
||||
"CREATE TABLE batch_copy_in (ts TIMESTAMP TIME INDEX, host STRING, val INT)",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let before = flushes();
|
||||
let rows = feed_copy_in(
|
||||
&pg,
|
||||
"COPY batch_copy_in FROM STDIN WITH (FORMAT csv)",
|
||||
vec![Bytes::from_static(
|
||||
b"2023-11-14 22:13:20,host1,1\n2023-11-14 22:13:21,host2,2\n",
|
||||
)],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
// Synchronous acknowledgement: the reported count is exact whether
|
||||
// or not the batcher is active.
|
||||
assert_eq!(rows, 2);
|
||||
// The batch must have been flushed through the pending-rows batcher.
|
||||
assert_eq!(flushes() - before, if pg_enabled { 1 } else { 0 });
|
||||
|
||||
// The flushed rows are visible immediately after the copy.
|
||||
let rows: Vec<String> = pg
|
||||
.query("SELECT host FROM batch_copy_in ORDER BY host", &[])
|
||||
.await
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|row| row.get(0))
|
||||
.collect();
|
||||
assert_eq!(rows, vec!["host1", "host2"]);
|
||||
|
||||
drop(pg);
|
||||
pg_task.await.unwrap();
|
||||
servers.shutdown_all().await.unwrap();
|
||||
instance.guard.remove_all().await;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user