Skip to content

Commit 300dab3

Browse files
robot-clickhouserobot-clickhouse
committed
Backport #80979 to 25.5: Fix atomic rename with truncate for files which compression is inferred from their file extension
1 parent 6827cb3 commit 300dab3

File tree

2 files changed

+17
-52
lines changed

2 files changed

+17
-52
lines changed

‎src/Client/ClientBase.cpp‎

Lines changed: 14 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -150,51 +150,20 @@ namespace
150150
{
151151
constexpr UInt64 THREAD_GROUP_ID = 0;
152152

153-
bool isSpecialFile(const String & path)
154-
{
155-
return path == "/dev/null" || path == "/dev/stdout" || path == "/dev/stderr";
156-
}
157153

158-
void handleTruncateMode(DB::ASTQueryWithOutput * query_with_output, const String & out_file, String & query)
159-
{
160-
if (!query_with_output->is_outfile_truncate)
161-
return;
162-
163-
/// Skip handling truncate mode of special files
164-
if (isSpecialFile(out_file))
165-
return;
166-
167-
/// Create a temporary file with unique suffix
168-
String tmp_file = out_file + ".tmp." + DB::toString(randomSeed());
169-
170-
/// Update the AST to use the temporary file
171-
auto tmp_file_literal = std::make_shared<DB::ASTLiteral>(tmp_file);
172-
query_with_output->out_file = tmp_file_literal;
173-
174-
/// Update the query string after modifying the AST
175-
query = query_with_output->formatWithSecretsOneLine();
176-
}
177-
178-
void cleanupTempFile(const DB::ASTPtr & parsed_query)
154+
void cleanupTempFile(const DB::ASTPtr & parsed_query, const String & tmp_file)
179155
{
180156
if (const auto * query_with_output = dynamic_cast<const DB::ASTQueryWithOutput *>(parsed_query.get()))
181157
{
182158
if (query_with_output->is_outfile_truncate && query_with_output->out_file)
183159
{
184-
const auto & tmp_file_node = query_with_output->out_file->as<DB::ASTLiteral &>();
185-
String tmp_file = tmp_file_node.value.safeGet<std::string>();
186-
187-
/// Skip rename for special files
188-
if (isSpecialFile(tmp_file))
189-
return;
190-
191160
if (fs::exists(tmp_file))
192161
fs::remove(tmp_file);
193162
}
194163
}
195164
}
196165

197-
void performAtomicRename(const DB::ASTPtr & parsed_query)
166+
void performAtomicRename(const DB::ASTPtr & parsed_query, const String & out_file)
198167
{
199168
if (const auto * query_with_output = dynamic_cast<const DB::ASTQueryWithOutput *>(parsed_query.get()))
200169
{
@@ -203,12 +172,6 @@ void performAtomicRename(const DB::ASTPtr & parsed_query)
203172
const auto & tmp_file_node = query_with_output->out_file->as<DB::ASTLiteral &>();
204173
String tmp_file = tmp_file_node.value.safeGet<std::string>();
205174

206-
/// Skip rename for special files
207-
if (isSpecialFile(tmp_file))
208-
return;
209-
210-
String out_file = tmp_file.substr(0, tmp_file.rfind(".tmp."));
211-
212175
try
213176
{
214177
fs::rename(tmp_file, out_file);
@@ -1169,11 +1132,12 @@ void ClientBase::processOrdinaryQuery(String query, ASTPtr parsed_query)
11691132
}
11701133
}
11711134

1135+
String out_file;
1136+
String out_file_if_truncated;
1137+
11721138
// Run some local checks to make sure queries into output file will work before sending to server.
11731139
if (const auto * query_with_output = dynamic_cast<const ASTQueryWithOutput *>(parsed_query.get()))
11741140
{
1175-
String out_file;
1176-
11771141
if (query_with_output->out_file)
11781142
{
11791143
if (isEmbeeddedClient())
@@ -1182,6 +1146,12 @@ void ClientBase::processOrdinaryQuery(String query, ASTPtr parsed_query)
11821146
const auto & out_file_node = query_with_output->out_file->as<ASTLiteral &>();
11831147
out_file = out_file_node.value.safeGet<std::string>();
11841148

1149+
if (query_with_output->is_outfile_truncate)
1150+
{
1151+
out_file_if_truncated = out_file;
1152+
out_file = fmt::format("tmp_{}.{}", UUIDHelpers::generateV4(), out_file);
1153+
}
1154+
11851155
std::string compression_method_string;
11861156

11871157
if (query_with_output->compression)
@@ -1231,11 +1201,6 @@ void ClientBase::processOrdinaryQuery(String query, ASTPtr parsed_query)
12311201
out_file);
12321202
}
12331203
}
1234-
1235-
if (query_with_output->is_outfile_truncate)
1236-
{
1237-
handleTruncateMode(const_cast<ASTQueryWithOutput *>(query_with_output), out_file, query);
1238-
}
12391204
}
12401205
}
12411206

@@ -1270,7 +1235,7 @@ void ClientBase::processOrdinaryQuery(String query, ASTPtr parsed_query)
12701235
catch (const NetException &)
12711236
{
12721237
// Clean up temporary file if it exists
1273-
cleanupTempFile(parsed_query);
1238+
cleanupTempFile(parsed_query, out_file);
12741239

12751240
// We still want to attempt to process whatever we already received or can receive (socket receive buffer can be not empty)
12761241
receiveResult(parsed_query, signals_before_stop, settings[Setting::partial_result_on_first_cancel]);
@@ -1280,14 +1245,14 @@ void ClientBase::processOrdinaryQuery(String query, ASTPtr parsed_query)
12801245
receiveResult(parsed_query, signals_before_stop, settings[Setting::partial_result_on_first_cancel]);
12811246

12821247
// After successful query execution, perform atomic rename for TRUNCATE mode
1283-
performAtomicRename(parsed_query);
1248+
performAtomicRename(parsed_query, out_file_if_truncated);
12841249

12851250
break;
12861251
}
12871252
catch (const Exception & e)
12881253
{
12891254
// Clean up temporary file if it exists
1290-
cleanupTempFile(parsed_query);
1255+
cleanupTempFile(parsed_query, out_file);
12911256

12921257
/// Retry when the server said "Client should retry" and no rows
12931258
/// has been received yet.

‎tests/queries/0_stateless/03362_into_outfile_atomic_truncate.sh‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,8 @@ ${CLICKHOUSE_CLIENT} --query="SELECT 'old,content' INTO OUTFILE '${CLICKHOUSE_TM
5252
perform "2" "SELECT 'new,content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_2.out' TRUNCATE FORMAT CSV" "\"new,content\"" "0"
5353

5454
# Test 3: Atomic TRUNCATE with compression
55-
${CLICKHOUSE_CLIENT} --query="SELECT 'old content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_3.out.gz' COMPRESSION 'gzip' FORMAT TSV" || { echo "Failed to create initial file for test 3"; exit 1; }
56-
perform "3" "SELECT 'new content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_3.out.gz' TRUNCATE COMPRESSION 'gzip' FORMAT TSV" "new content" "1"
55+
${CLICKHOUSE_CLIENT} --query="SELECT 'old content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_3.out.gz' FORMAT TSV" || { echo "Failed to create initial file for test 3"; exit 1; }
56+
perform "3" "SELECT 'new content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_3.out.gz' TRUNCATE FORMAT TSV" "new content" "1"
5757

5858
# Test 4: Raw text file using RawBLOB format
5959
${CLICKHOUSE_CLIENT} --query="SELECT 'old text' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_4.out' FORMAT RawBLOB" || { echo "Failed to create initial file for test 4"; exit 1; }
@@ -68,4 +68,4 @@ new line2" "0"
6868
${CLICKHOUSE_CLIENT} --query="SELECT 'old content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_6.out.tmp'" || { echo "Failed to create initial file for test 6"; exit 1; }
6969
perform "6" "SELECT 'new content' INTO OUTFILE '${CLICKHOUSE_TMP}/test_atomic_6.out.tmp' TRUNCATE" "new content" "0" ".tmp"
7070

71-
echo "OK"
71+
echo "OK"

0 commit comments

Comments
 (0)