feat(cubesql): Support COPY ... FROM STDIN - #11538
Conversation
Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
Claude finished @MazterQyou's task in 9m 44s —— View job Review: 2 high, 3 medium, 4 low — 6 inline comments postedFull review — findings, verification notes, and todo listTodos
GeneralThis is unusually careful work. The decoder mirrors Findings below; the substantive ones are also inline on the diff. High
Medium
Low / nits
Verification notesI did not build or run the test suite here (the DataFusion fork build is prohibitively long in this environment) — findings 1–5 come from reading the call graph, and the reasoning for each is spelled out in the inline comments so they can be checked quickly. Findings 1 and 2 are deterministic from the code paths shown rather than heuristic. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
|
||
| let schema = Arc::new(Schema::new(fields)); | ||
| let temp_tables = self.state.temp_tables(); | ||
| tokio::task::spawn_blocking(move || { | ||
| temp_tables.save( | ||
| &table_name, | ||
| TempTable::from_arrow_schema(schema, vec![vec![]]), | ||
| ) | ||
| }) | ||
| .await | ||
| .map_err(|err| CompilationError::internal(err.to_string()))??; |
There was a problem hiding this comment.
The table is created during planning, which breaks the extended query protocol.
convert_statement_to_cube_query → plan_query is called from three places that are not query execution:
prepare_statement(shim.rs:1127) — everyParsemessagebind(shim.rs:1006) — everyBindmessageexplain_to_plan(router.rs:318) andDECLAREvalidation (shim.rs:1537), which the comment there explicitly calls "just a verification that we can compile that query"
So for any driver that uses the extended protocol for DDL (JDBC, ODBC, psycopg with server-side prepare, pgx, tokio_postgres::execute), CREATE TEMPORARY TABLE t (n int) is planned at Parse, creating the table, and planned again at Bind, which hits the contains_key check in TempTableManager::save and fails with relation "t" already exists — before Execute ever runs. Same for EXPLAIN CREATE TEMPORARY TABLE …, and a Parse that is never executed still leaves the table behind.
Every test in test_copy.rs uses simple_query, which plans exactly once, so this path is uncovered.
Suggest following the existing CreateTempTable shape: return a plan variant carrying the schema and do the save at execution time (in Portal/write_portal), and add a test that creates the table via client.execute(...)/prepare rather than simple_query.
| let Some(next) = self.peek(self.scanned + 1) else { | ||
| if at_end { | ||
| // A backslash at the very end escapes nothing | ||
| return Ok(Line::Row(self.split_line(0))); | ||
| } | ||
|
|
||
| return Ok(Line::Incomplete); | ||
| }; |
There was a problem hiding this comment.
A trailing lone backslash at end of data drops a field and then spins.
split_line(0) here takes buffer[..scanned] and leaves the backslash in the buffer with scanned reset to 0. Two consequences for text data whose last line ends with \ and has no trailing newline (e.g. 1\tone\):
- The backslash is silently dropped from the row, even though
unescape_text(line 679) correctly treats a trailing backslash as a literal — so the two disagree. finish()'swhile !self.finishedloop re-enterstake_line(true)with the identical state (scanned == 0, buffer =="\\"), getsRow("")again, and loops. For a single-column table each iteration appends an empty row and bumpsbuilt_bytes, so it only terminates ~2M iterations later viacheck_memory_limit, reportingCOPY data exceeds the temporary table memory limitinstead of loading the row.
The fix is to let the ordinary "unterminated last row" branch (line 282) take the whole buffer:
| let Some(next) = self.peek(self.scanned + 1) else { | |
| if at_end { | |
| // A backslash at the very end escapes nothing | |
| return Ok(Line::Row(self.split_line(0))); | |
| } | |
| return Ok(Line::Incomplete); | |
| }; | |
| let Some(next) = self.peek(self.scanned + 1) else { | |
| if at_end { | |
| // A backslash at the very end escapes nothing | |
| self.scanned += 1; | |
| continue; | |
| } | |
| return Ok(Line::Incomplete); | |
| }; |
Worth a unit test for text("1\tone\tt\n2\ttwo\tt\\") and the single-column variant.
| fn parse_decimal(value: &str, precision: usize, scale: usize) -> Option<i128> { | ||
| let (sign, digits) = match value.strip_prefix('-') { | ||
| Some(digits) => (-1, digits), | ||
| None => (1, value.strip_prefix('+').unwrap_or(value)), | ||
| }; | ||
|
|
||
| let mut parts = digits.split('.'); | ||
| let integer = parts.next().unwrap_or(""); | ||
| let fraction = parts.next().unwrap_or(""); | ||
| if parts.next().is_some() || fraction.len() > scale { | ||
| return None; | ||
| } | ||
|
|
||
| let unscaled = format!("{}{:0<width$}", integer, fraction, width = scale); | ||
| if unscaled.len() > precision { | ||
| return None; | ||
| } | ||
|
|
||
| Some(sign * unscaled.parse::<i128>().ok()?) | ||
| } |
There was a problem hiding this comment.
parse_decimal rejects values PostgreSQL accepts.
Both checks are off relative to PG's numeric input:
unscaled.len() > precisioncounts leading zeros, so007intonumeric(2,0)is rejected (PG stores7), and0.05intonumeric(2,2)becomes"005"→ rejected (PG stores0.05, sincenumeric(2,2)holds values in(-1, 1)with 2 decimals).fraction.len() > scaleerrors, but PG rounds to the declared scale:1.234intonumeric(4,2)is1.23, not an error. With theExactNumberInfo::Nonedefault ofDecimal(38, 10)(router.rs:1048) that means bareNUMERICrefuses anything with more than 10 decimals, andNUMERIC(10)(scale 0) refuses1.5.
Roughly: round/truncate the fraction to scale first, then strip leading zeros from unscaled before comparing against precision.
| fn column_names(columns: &[ast::Ident]) -> Vec<String> { | ||
| columns.iter().map(|column| column.value.clone()).collect() | ||
| } |
There was a problem hiding this comment.
Column names here keep the identifier exactly as written, while the COPY column list goes through normalize_ident (router.rs:972) and the schema stores folded names. So COPY t (A) FROM STDIN works, but COPY t FROM STDIN WITH (FORMAT csv, FORCE_NOT_NULL (A)) fails the schema.index_of validation in copy_from_plan with column "A" of relation "t" does not exist, and even if it passed, resolve_value compares against the folded field name and would never match.
column_names should apply the same folding (lowercase unquoted idents, keep quoted ones verbatim).
| // Columns which the COPY statement did not list stay NULL | ||
| let mut values: Vec<Option<String>> = vec![None; self.builders.len()]; |
There was a problem hiding this comment.
NOT NULL is accepted by create_empty_table_to_plan and stored in the Arrow field's nullable flag, but nothing enforces it here: unlisted columns are pre-filled with None, and an explicit \N for a non-nullable column is appended as a null too. test_copy_from_stdin_column_types declares s varchar(10) NOT NULL but always supplies a value, so the gap isn't covered.
Depending on whether this fork's RecordBatch::try_new validates nullability, the result is either a silently violated constraint or an opaque Unable to build COPY data: … internal error at finish(). Either way it'd be better to check at append time and raise PG's null value in column "s" violates not-null constraint (23502) — or, if enforcing is out of scope, reject NOT NULL in the DDL like the other unsupported column options.
| ast::DataType::Timestamp(_, _) | ast::DataType::Datetime(_) => { | ||
| DataType::Timestamp(TimeUnit::Nanosecond, None) | ||
| } |
There was a problem hiding this comment.
Minor: the timezone flag of ast::DataType::Timestamp(_, tz) is discarded, so timestamptz silently becomes a naive Timestamp(Nanosecond, None), and parse_timestamp (copy.rs:949) has no format with an offset — a perfectly ordinary 2024-03-01 10:20:30+03 then fails with invalid input syntax for type timestamp. Either reject WITH TIME ZONE explicitly at DDL time or accept an offset and normalize to UTC.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #11538 +/- ##
===========================================
+ Coverage 59.35% 84.26% +24.91%
===========================================
Files 223 260 +37
Lines 17989 83621 +65632
Branches 3641 0 -3641
===========================================
+ Hits 10677 70463 +59786
- Misses 6793 13158 +6365
+ Partials 519 0 -519
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
1 participant
Check List
Description of Changes Made
This PR adds support for
COPY ... FROM STDINcommand, allowing adding data to temporary tables. Related tests are included.