From 300ec160f40fffb48bf160fe82a11781819ba30e Mon Sep 17 00:00:00 2001
From: Gleb Novikov <gnovikov@iponweb.net>
Date: Sun, 19 May 2019 08:27:00 +0300
Subject: [PATCH] Constraints MVP

---
 .../CheckConstraintsBlockOutputStream.cpp     | 43 +++++++++++++++++
 .../CheckConstraintsBlockOutputStream.h       | 48 +++++++++++++++++++
 .../Interpreters/InterpreterInsertQuery.cpp   |  4 ++
 dbms/src/Storages/ConstraintsDescription.h    | 14 +++++-
 4 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 dbms/src/DataStreams/CheckConstraintsBlockOutputStream.cpp
 create mode 100644 dbms/src/DataStreams/CheckConstraintsBlockOutputStream.h

diff --git a/dbms/src/DataStreams/CheckConstraintsBlockOutputStream.cpp b/dbms/src/DataStreams/CheckConstraintsBlockOutputStream.cpp
new file mode 100644
index 00000000000..99f9f9bc90d
--- /dev/null
+++ b/dbms/src/DataStreams/CheckConstraintsBlockOutputStream.cpp
@@ -0,0 +1,43 @@
+#include <DataStreams/CheckConstraintsBlockOutputStream.h>
+
+
+namespace DB
+{
+
+void CheckConstraintsBlockOutputStream::write(const Block & block)
+{
+    for (auto & constraint_expr: expressions)
+        if (!checkConstraintOnBlock(block, constraint_expr))
+            throw Exception("Some constraints are not satisfied", ErrorCodes::QUERY_WAS_CANCELLED);
+    output->write(block);
+}
+
+void CheckConstraintsBlockOutputStream::flush()
+{
+    output->flush();
+}
+
+void CheckConstraintsBlockOutputStream::writePrefix()
+{
+    output->writePrefix();
+}
+
+void CheckConstraintsBlockOutputStream::writeSuffix()
+{
+    output->writeSuffix();
+}
+
+bool CheckConstraintsBlockOutputStream::checkConstraintOnBlock(const Block & block, const ExpressionActionsPtr & constraint)
+{
+    Block res = block;
+    constraint->execute(res);
+    assert(block.columns() == res.columns() - 1);
+    ColumnWithTypeAndName res_column = res.safeGetByPosition(res.columns() - 1);
+    size_t column_size = res_column.column->size();
+    for (size_t i = 0; i < column_size; ++i)
+        if (!res_column.column->getBool(i))
+            return false;
+    return true;
+}
+
+}
diff --git a/dbms/src/DataStreams/CheckConstraintsBlockOutputStream.h b/dbms/src/DataStreams/CheckConstraintsBlockOutputStream.h
new file mode 100644
index 00000000000..623eccc8172
--- /dev/null
+++ b/dbms/src/DataStreams/CheckConstraintsBlockOutputStream.h
@@ -0,0 +1,48 @@
+#pragma once
+
+#include <DataStreams/IBlockOutputStream.h>
+#include <Storages/ConstraintsDescription.h>
+#include <DataStreams/OneBlockInputStream.h>
+#include <DataStreams/ExpressionBlockInputStream.h>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int CONSTRAINTS_ARE_NOT_SATISFIED;
+}
+
+class CheckConstraintsBlockOutputStream : public IBlockOutputStream
+{
+public:
+    CheckConstraintsBlockOutputStream(
+            const BlockOutputStreamPtr & output_,
+            const Block & header_,
+            const ConstraintsDescription & constraints_,
+            const Context & context_)
+            : output(output_),
+              header(header_),
+              constraints(constraints_),
+              expressions(constraints_.getExpressions(context_, header.getNamesAndTypesList())),
+              context(context_)
+    { }
+
+    Block getHeader() const override { return header; }
+    void write(const Block & block) override;
+
+    void flush() override;
+
+    void writePrefix() override;
+    void writeSuffix() override;
+
+    bool checkConstraintOnBlock(const Block & block, const ExpressionActionsPtr & constraint);
+
+private:
+    BlockOutputStreamPtr output;
+    Block header;
+    const ConstraintsDescription constraints;
+    const ConstraintsExpressions expressions;
+    const Context & context;
+};
+}
diff --git a/dbms/src/Interpreters/InterpreterInsertQuery.cpp b/dbms/src/Interpreters/InterpreterInsertQuery.cpp
index e4391f52247..fa6df1599ea 100644
--- a/dbms/src/Interpreters/InterpreterInsertQuery.cpp
+++ b/dbms/src/Interpreters/InterpreterInsertQuery.cpp
@@ -5,6 +5,7 @@
 
 #include <DataStreams/AddingDefaultBlockOutputStream.h>
 #include <DataStreams/AddingDefaultsBlockInputStream.h>
+#include <DataStreams/CheckConstraintsBlockOutputStream.h>
 #include <DataStreams/OwningBlockInputStream.h>
 #include <DataStreams/ConvertingBlockInputStream.h>
 #include <DataStreams/CountingBlockOutputStream.h>
@@ -117,6 +118,9 @@ BlockIO InterpreterInsertQuery::execute()
     out = std::make_shared<AddingDefaultBlockOutputStream>(
         out, query_sample_block, table->getSampleBlock(), table->getColumns().getDefaults(), context);
 
+    out = std::make_shared<CheckConstraintsBlockOutputStream>(
+            out, query_sample_block, table->getConstraints(), context);
+
     auto out_wrapper = std::make_shared<CountingBlockOutputStream>(out);
     out_wrapper->setProcessListElement(context.getProcessListElement());
     out = std::move(out_wrapper);
diff --git a/dbms/src/Storages/ConstraintsDescription.h b/dbms/src/Storages/ConstraintsDescription.h
index c2954d94428..fbb0f5167fc 100644
--- a/dbms/src/Storages/ConstraintsDescription.h
+++ b/dbms/src/Storages/ConstraintsDescription.h
@@ -1,12 +1,13 @@
 #pragma once
 
 #include <Parsers/ASTConstraintDeclaration.h>
-
+#include <Interpreters/ExpressionAnalyzer.h>
 
 namespace DB
 {
 
 using ConstraintsASTs = std::vector<std::shared_ptr<ASTConstraintDeclaration>>;
+using ConstraintsExpressions = std::vector<ExpressionActionsPtr>;
 
 struct ConstraintsDescription
 {
@@ -18,6 +19,17 @@ struct ConstraintsDescription
     String toString() const;
 
     static ConstraintsDescription parse(const String & str);
+
+    ConstraintsExpressions getExpressions(const Context & context, const NamesAndTypesList & source_columns_) const {
+        ConstraintsExpressions res;
+        res.reserve(constraints.size());
+        for (const auto & constraint : constraints) {
+            ASTPtr expr = constraint->expr->clone();
+            auto syntax_result = SyntaxAnalyzer(context).analyze(expr, source_columns_);
+            res.push_back(ExpressionAnalyzer(constraint->expr->clone(), syntax_result, context).getActions(false));
+        }
+        return res;
+    }
 };
 
 }