Merge pull request #27191 from FArthur-cmd/memory-tracker-26043

Client with flag for memory tracker

programs/client/Client.cpp

@@ -1,3 +1,6 @@
+#include <string>
+#include "Common/MemoryTracker.h"
+#include "Columns/ColumnsNumber.h"
 #include "ConnectionParameters.h"
 #include "QueryFuzzer.h"
 #include "Suggest.h"
@@ -100,6 +103,14 @@
 #pragma GCC optimize("-fno-var-tracking-assignments")
 #endif
 
+namespace CurrentMetrics
+{
+    extern const Metric Revision;
+    extern const Metric VersionInteger;
+    extern const Metric MemoryTracking;
+    extern const Metric MaxDDLEntryID;
+}
+
 namespace fs = std::filesystem;
 
 namespace DB
@@ -524,6 +535,18 @@ private:
     {
         UseSSL use_ssl;
 
+        MainThreadStatus::getInstance();
+
+        /// Limit on total memory usage
+        size_t max_client_memory_usage = config().getInt64("max_memory_usage_in_client", 0 /*default value*/);
+
+        if (max_client_memory_usage != 0)
+        {
+            total_memory_tracker.setHardLimit(max_client_memory_usage);
+            total_memory_tracker.setDescription("(total)");
+            total_memory_tracker.setMetric(CurrentMetrics::MemoryTracking);
+        }
+
         registerFormats();
         registerFunctions();
         registerAggregateFunctions();
@@ -2581,6 +2604,7 @@ public:
             ("opentelemetry-tracestate", po::value<std::string>(), "OpenTelemetry tracestate header as described by W3C Trace Context recommendation")
             ("history_file", po::value<std::string>(), "path to history file")
             ("no-warnings", "disable warnings when client connects to server")
+            ("max_memory_usage_in_client", po::value<int>(), "sets memory limit in client")
         ;
 
         Settings cmd_settings;

src/Core/ya.make

@@ -31,6 +31,10 @@ SRCS(
     MySQL/PacketsProtocolText.cpp
     MySQL/PacketsReplication.cpp
     NamesAndTypes.cpp
+    PostgreSQL/Connection.cpp
+    PostgreSQL/PoolWithFailover.cpp
+    PostgreSQL/Utils.cpp
+    PostgreSQL/insertPostgreSQLValue.cpp
     PostgreSQLProtocol.cpp
     QueryProcessingStage.cpp
     Settings.cpp

src/DataStreams/ya.make

@@ -49,6 +49,7 @@ SRCS(
     TTLUpdateInfoAlgorithm.cpp
     copyData.cpp
     finalizeBlock.cpp
+    formatBlock.cpp
     materializeBlock.cpp
     narrowBlockInputStreams.cpp
 

src/Parsers/ya.make

@@ -21,6 +21,7 @@ SRCS(
     ASTCreateRowPolicyQuery.cpp
     ASTCreateSettingsProfileQuery.cpp
     ASTCreateUserQuery.cpp
+    ASTDatabaseOrNone.cpp
     ASTDictionary.cpp
     ASTDictionaryAttributeDeclaration.cpp
     ASTDropAccessEntityQuery.cpp
@@ -95,6 +96,7 @@ SRCS(
     ParserCreateSettingsProfileQuery.cpp
     ParserCreateUserQuery.cpp
     ParserDataType.cpp
+    ParserDatabaseOrNone.cpp
     ParserDescribeTableQuery.cpp
     ParserDictionary.cpp
     ParserDictionaryAttributeDeclaration.cpp

src/Processors/ya.make

@@ -7,14 +7,8 @@ PEERDIR(
     clickhouse/src/Common
     contrib/libs/msgpack
     contrib/libs/protobuf
-    contrib/libs/arrow
 )
 
-ADDINCL(
-    contrib/libs/arrow/src
-)
-
-CFLAGS(-DUSE_ARROW=1)
 
 SRCS(
     Chunk.cpp
@@ -31,11 +25,6 @@ SRCS(
     Formats/IOutputFormat.cpp
     Formats/IRowInputFormat.cpp
     Formats/IRowOutputFormat.cpp
-    Formats/Impl/ArrowBlockInputFormat.cpp
-    Formats/Impl/ArrowBlockOutputFormat.cpp
-    Formats/Impl/ArrowBufferedStreams.cpp
-    Formats/Impl/ArrowColumnToCHColumn.cpp
-    Formats/Impl/CHColumnToArrowColumn.cpp
     Formats/Impl/BinaryRowInputFormat.cpp
     Formats/Impl/BinaryRowOutputFormat.cpp
     Formats/Impl/CSVRowInputFormat.cpp

src/Storages/ya.make

@@ -141,6 +141,7 @@ SRCS(
     StorageMerge.cpp
     StorageMergeTree.cpp
     StorageMongoDB.cpp
+    StorageMongoDBSocketFactory.cpp
     StorageMySQL.cpp
     StorageNull.cpp
     StorageReplicatedMergeTree.cpp

tests/queries/0_stateless/02003_memory_limit_in_client.expect

@@ -0,0 +1,40 @@
+#!/usr/bin/expect -f
+
+# This is a test for system.warnings. Testing in interactive mode is necessary,
+# as we want to see certain warnings from client
+
+log_user 0
+set timeout 60
+match_max 100000
+
+# A default timeout action is to do nothing, change it to fail
+expect_after {
+    timeout {
+        exit 1
+    }
+}
+
+set basedir [file dirname $argv0]
+spawn bash -c "source $basedir/../shell_config.sh ; \$CLICKHOUSE_CLIENT_BINARY \$CLICKHOUSE_CLIENT_OPT --disable_suggestion --max_memory_usage_in_client=1"
+expect ":) "
+
+send  -- "SELECT arrayMap(x -> range(x), range(number)) FROM numbers(1000)\r"
+expect "Code: 241"
+
+expect ":) "
+
+# Exit.
+send -- "\4"
+expect eof
+
+set basedir [file dirname $argv0]
+spawn bash -c "source $basedir/../shell_config.sh ; \$CLICKHOUSE_CLIENT_BINARY \$CLICKHOUSE_CLIENT_OPT --disable_suggestion --max_memory_usage_in_client=1"
+expect ":) "
+
+send  -- "SELECT * FROM (SELECT * FROM system.numbers LIMIT 600000) as num WHERE num.number=60000\r"
+expect "60000"
+expect ":) "
+
+# Exit.
+send -- "\4"
+expect eof