From 095f558ba12d46a1e6ed39f1f705a76acdd7a3de Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Mon, 20 Dec 2021 21:05:37 +0300
Subject: [PATCH] Cover memory tracking of external tables via HTTP queries

Thanks to @den-crane
---
 ..._external_tables_memory_tracking.reference | 16 +++++++
 ...52_http_external_tables_memory_tracking.sh | 48 +++++++++++++++++++
 2 files changed, 64 insertions(+)
 create mode 100644 tests/queries/0_stateless/02152_http_external_tables_memory_tracking.reference
 create mode 100755 tests/queries/0_stateless/02152_http_external_tables_memory_tracking.sh

diff --git a/tests/queries/0_stateless/02152_http_external_tables_memory_tracking.reference b/tests/queries/0_stateless/02152_http_external_tables_memory_tracking.reference
new file mode 100644
index 00000000000..1fc09c8d154
--- /dev/null
+++ b/tests/queries/0_stateless/02152_http_external_tables_memory_tracking.reference
@@ -0,0 +1,16 @@
+Checking input_format_parallel_parsing=false&
+1
+Checking input_format_parallel_parsing=false&cancel_http_readonly_queries_on_client_close=1&readonly=1
+1
+Checking input_format_parallel_parsing=false&send_progress_in_http_headers=true
+1
+Checking input_format_parallel_parsing=false&cancel_http_readonly_queries_on_client_close=1&readonly=1&send_progress_in_http_headers=true
+1
+Checking input_format_parallel_parsing=true&
+1
+Checking input_format_parallel_parsing=true&cancel_http_readonly_queries_on_client_close=1&readonly=1
+1
+Checking input_format_parallel_parsing=true&send_progress_in_http_headers=true
+1
+Checking input_format_parallel_parsing=true&cancel_http_readonly_queries_on_client_close=1&readonly=1&send_progress_in_http_headers=true
+1
diff --git a/tests/queries/0_stateless/02152_http_external_tables_memory_tracking.sh b/tests/queries/0_stateless/02152_http_external_tables_memory_tracking.sh
new file mode 100755
index 00000000000..7fc41a46c33
--- /dev/null
+++ b/tests/queries/0_stateless/02152_http_external_tables_memory_tracking.sh
@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+
+# Regression for proper release of Context,
+# via tracking memory of external tables.
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CURDIR"/../shell_config.sh
+
+tmp_file=$(mktemp "$CURDIR/clickhouse.XXXXXX.csv")
+trap 'rm $tmp_file' EXIT
+
+$CLICKHOUSE_CLIENT -q "SELECT toString(number) FROM numbers(1e6) FORMAT TSV" > "$tmp_file"
+
+function run_and_check()
+{
+    local query_id
+    query_id="$(uuidgen)"
+
+    echo "Checking $*"
+
+    # Run query with external table (implicit StorageMemory user)
+    $CLICKHOUSE_CURL -sS -F "s=@$tmp_file;" "$CLICKHOUSE_URL&s_structure=key+Int&query=SELECT+count()+FROM+s&memory_profiler_sample_probability=1&query_id=$query_id&$*" -o /dev/null
+
+    ${CLICKHOUSE_CURL} -sS "${CLICKHOUSE_URL}" --data-binary @- <<<'SYSTEM FLUSH LOGS'
+
+    # Check that temporary table had been destroyed.
+    ${CLICKHOUSE_CURL} -sS "${CLICKHOUSE_URL}&allow_introspection_functions=1" --data-binary @- <<<"
+    WITH arrayStringConcat(arrayMap(x -> demangle(addressToSymbol(x)), trace), '\n') AS sym
+    SELECT count()>0 FROM system.trace_log
+    WHERE
+        sym LIKE '%DB::StorageMemory::drop%\n%TemporaryTableHolder::~TemporaryTableHolder%' AND
+        query_id = '$query_id'
+    "
+}
+
+for input_format_parallel_parsing in false true; do
+    query_args_variants=(
+        ""
+        "cancel_http_readonly_queries_on_client_close=1&readonly=1"
+        "send_progress_in_http_headers=true"
+        # nested progress callback
+        "cancel_http_readonly_queries_on_client_close=1&readonly=1&send_progress_in_http_headers=true"
+    )
+    for query_args in "${query_args_variants[@]}"; do
+        run_and_check "input_format_parallel_parsing=$input_format_parallel_parsing&$query_args"
+    done
+done