From f3b3025719815a656d363336e5758572b03effb6 Mon Sep 17 00:00:00 2001
From: alesapin <alesapin@gmail.com>
Date: Mon, 2 Nov 2020 14:39:27 +0300
Subject: [PATCH] Ban non comparable types in primary key

---
 src/Common/ErrorCodes.cpp                                | 1 +
 src/Core/Field.h                                         | 8 +++++---
 src/Storages/KeyDescription.cpp                          | 9 +++++++++
 .../01548_uncomparable_columns_in_keys.reference         | 0
 .../0_stateless/01548_uncomparable_columns_in_keys.sql   | 9 +++++++++
 5 files changed, 24 insertions(+), 3 deletions(-)
 create mode 100644 tests/queries/0_stateless/01548_uncomparable_columns_in_keys.reference
 create mode 100644 tests/queries/0_stateless/01548_uncomparable_columns_in_keys.sql

diff --git a/src/Common/ErrorCodes.cpp b/src/Common/ErrorCodes.cpp
index b14c090c848..371cc6be304 100644
--- a/src/Common/ErrorCodes.cpp
+++ b/src/Common/ErrorCodes.cpp
@@ -512,6 +512,7 @@ namespace ErrorCodes
     extern const int NO_ROW_DELIMITER = 546;
     extern const int INVALID_RAID_TYPE = 547;
     extern const int UNKNOWN_VOLUME = 548;
+    extern const int DATA_TYPE_CANNOT_BE_USED_IN_KEY = 549;
 
     extern const int KEEPER_EXCEPTION = 999;
     extern const int POCO_EXCEPTION = 1000;
diff --git a/src/Core/Field.h b/src/Core/Field.h
index 8973d106c0b..be01352c168 100644
--- a/src/Core/Field.h
+++ b/src/Core/Field.h
@@ -26,6 +26,7 @@ namespace ErrorCodes
     extern const int NOT_IMPLEMENTED;
     extern const int LOGICAL_ERROR;
     extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int LOGICAL_ERROR;
 }
 
 template <typename T, typename SFINAE = void>
@@ -767,9 +768,10 @@ T & Field::get()
 #ifndef NDEBUG
     // Disregard signedness when converting between int64 types.
     constexpr Field::Types::Which target = TypeToEnum<NearestFieldType<ValueType>>::value;
-    assert(target == which
-           || (isInt64FieldType(target) && isInt64FieldType(which))
-           || target == Field::Types::Decimal64 /* DateTime64 fields */);
+    if (target != which
+           && (!isInt64FieldType(target) || !isInt64FieldType(which))
+        && target != Field::Types::Decimal64 /* DateTime64 fields */)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Invalid Field get from type {} to type {}", Types::toString(which), Types::toString(target));
 #endif
 
     ValueType * MAY_ALIAS ptr = reinterpret_cast<ValueType *>(&storage);
diff --git a/src/Storages/KeyDescription.cpp b/src/Storages/KeyDescription.cpp
index 533736d19ed..ebf31f02aab 100644
--- a/src/Storages/KeyDescription.cpp
+++ b/src/Storages/KeyDescription.cpp
@@ -2,10 +2,12 @@
 
 #include <Functions/IFunction.h>
 #include <Parsers/ASTIdentifier.h>
+#include <DataTypes/DataTypeAggregateFunction.h>
 #include <Interpreters/ExpressionActions.h>
 #include <Interpreters/ExpressionAnalyzer.h>
 #include <Interpreters/TreeRewriter.h>
 #include <Storages/extractKeyExpressionList.h>
+#include <Common/quoteString.h>
 
 
 namespace DB
@@ -14,6 +16,7 @@ namespace DB
 namespace ErrorCodes
 {
     extern const int LOGICAL_ERROR;
+    extern const int DATA_TYPE_CANNOT_BE_USED_IN_KEY;
 }
 
 KeyDescription::KeyDescription(const KeyDescription & other)
@@ -115,7 +118,13 @@ KeyDescription KeyDescription::getSortingKeyFromAST(
     }
 
     for (size_t i = 0; i < result.sample_block.columns(); ++i)
+    {
         result.data_types.emplace_back(result.sample_block.getByPosition(i).type);
+        if (!result.data_types.back()->isComparable())
+            throw Exception(ErrorCodes::DATA_TYPE_CANNOT_BE_USED_IN_KEY,
+                            "Column {} with type {} is not allowed in key expression, it's not comparable",
+                            backQuote(result.sample_block.getByPosition(i).name), result.data_types.back()->getName());
+    }
 
     return result;
 }
diff --git a/tests/queries/0_stateless/01548_uncomparable_columns_in_keys.reference b/tests/queries/0_stateless/01548_uncomparable_columns_in_keys.reference
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/queries/0_stateless/01548_uncomparable_columns_in_keys.sql b/tests/queries/0_stateless/01548_uncomparable_columns_in_keys.sql
new file mode 100644
index 00000000000..ff51085f58c
--- /dev/null
+++ b/tests/queries/0_stateless/01548_uncomparable_columns_in_keys.sql
@@ -0,0 +1,9 @@
+DROP TABLE IF EXISTS uncomparable_keys;
+
+CREATE TABLE foo (id UInt64, key AggregateFunction(max, UInt64)) ENGINE MergeTree ORDER BY key; --{serverError 549}
+
+CREATE TABLE foo (id UInt64, key AggregateFunction(max, UInt64)) ENGINE MergeTree PARTITION BY key; --{serverError 549}
+
+CREATE TABLE foo (id UInt64, key AggregateFunction(max, UInt64)) ENGINE MergeTree ORDER BY (key) SAMPLE BY key; --{serverError 549}
+
+DROP TABLE IF EXISTS uncomparable_keys;