From 54db1127fe6bf3163beb3d9122d9a7bd622165a4 Mon Sep 17 00:00:00 2001
From: George Leung <phiSgr@users.noreply.github.com>
Date: Tue, 20 Feb 2024 22:25:15 -0800
Subject: [PATCH] Use merge tree engine for data table (#451)

We go from not supporting primary key to requiring primary key.
---
 .../src/infrastructure/olap/clickhouse.rs     |  3 +-
 .../infrastructure/olap/clickhouse/queries.rs | 41 +++++++++++++------
 apps/framework-cli/tests/psl/event.prisma     |  8 ++--
 apps/framework-cli/tests/psl/simple.prisma    |  8 ++--
 4 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/apps/framework-cli/src/infrastructure/olap/clickhouse.rs b/apps/framework-cli/src/infrastructure/olap/clickhouse.rs
index 920a9e353..d8d43d794 100644
--- a/apps/framework-cli/src/infrastructure/olap/clickhouse.rs
+++ b/apps/framework-cli/src/infrastructure/olap/clickhouse.rs
@@ -9,6 +9,7 @@ use log::debug;
 
 use serde::{Deserialize, Serialize};
 
+use crate::infrastructure::olap::clickhouse::queries::ClickhouseEngine;
 use crate::{
     framework::schema::{FieldArity, UnsupportedDataTypeError},
     utilities::constants::REDPANDA_CONTAINER_NAME,
@@ -182,7 +183,7 @@ impl ClickhouseTable {
         )
     }
     pub fn create_data_table_query(&self) -> Result<String, UnsupportedDataTypeError> {
-        CreateTableQuery::build(self.clone(), "Memory".to_string())
+        CreateTableQuery::build(self.clone(), ClickhouseEngine::MergeTree)
     }
 
     pub fn drop_kafka_table_query(&self) -> Result<String, UnsupportedDataTypeError> {
diff --git a/apps/framework-cli/src/infrastructure/olap/clickhouse/queries.rs b/apps/framework-cli/src/infrastructure/olap/clickhouse/queries.rs
index c896c6851..0adb85eaa 100644
--- a/apps/framework-cli/src/infrastructure/olap/clickhouse/queries.rs
+++ b/apps/framework-cli/src/infrastructure/olap/clickhouse/queries.rs
@@ -36,6 +36,11 @@ pub struct CreateTableQuery;
 static KAFKA_SETTINGS: &str =
     "kafka_skip_broken_messages = 1, date_time_input_format = 'best_effort'";
 
+pub enum ClickhouseEngine {
+    MergeTree,
+    Kafka(String, u16, String),
+}
+
 impl CreateTableQuery {
     pub fn kafka(
         table: ClickhouseTable,
@@ -45,16 +50,13 @@ impl CreateTableQuery {
     ) -> Result<String, UnsupportedDataTypeError> {
         CreateTableQuery::build(
             table,
-            format!(
-                "Kafka('{}:{}', '{}', 'clickhouse-group', 'JSONEachRow') SETTINGS {}",
-                kafka_host, kafka_port, topic, KAFKA_SETTINGS,
-            ),
+            ClickhouseEngine::Kafka(kafka_host, kafka_port, topic),
         )
     }
 
     pub fn build(
         table: ClickhouseTable,
-        engine: String,
+        engine: ClickhouseEngine,
     ) -> Result<String, UnsupportedDataTypeError> {
         let mut tt = TinyTemplate::new();
         tt.set_default_formatter(&format_unescaped); // by default it formats HTML-escaped and messes up single quotes
@@ -78,14 +80,29 @@ struct CreateTableContext {
 impl CreateTableContext {
     fn new(
         table: ClickhouseTable,
-        engine: String,
+        engine: ClickhouseEngine,
     ) -> Result<CreateTableContext, UnsupportedDataTypeError> {
-        let primary_key = table
-            .columns
-            .iter()
-            .filter(|column| column.primary_key)
-            .map(|column| column.name.clone())
-            .collect::<Vec<String>>();
+        let (engine, ignore_primary_key) = match engine {
+            ClickhouseEngine::MergeTree => ("MergeTree".to_string(), false),
+            ClickhouseEngine::Kafka(kafka_host, kafka_port, topic) => (
+                format!(
+                    "Kafka('{}:{}', '{}', 'clickhouse-group', 'JSONEachRow') SETTINGS {}",
+                    kafka_host, kafka_port, topic, KAFKA_SETTINGS,
+                ),
+                true,
+            ),
+        };
+
+        let primary_key = if ignore_primary_key {
+            Vec::new()
+        } else {
+            table
+                .columns
+                .iter()
+                .filter(|column| column.primary_key)
+                .map(|column| column.name.clone())
+                .collect::<Vec<String>>()
+        };
 
         Ok(CreateTableContext {
             db_name: table.db_name,
diff --git a/apps/framework-cli/tests/psl/event.prisma b/apps/framework-cli/tests/psl/event.prisma
index 18ec4d257..b25a6a2b4 100644
--- a/apps/framework-cli/tests/psl/event.prisma
+++ b/apps/framework-cli/tests/psl/event.prisma
@@ -1,5 +1,5 @@
 model Awesome {
-    id Int
-    name String
-    description String
-}
\ No newline at end of file
+  id          Int    @id
+  name        String
+  description String
+}
diff --git a/apps/framework-cli/tests/psl/simple.prisma b/apps/framework-cli/tests/psl/simple.prisma
index 6001b7c42..00c5298c1 100644
--- a/apps/framework-cli/tests/psl/simple.prisma
+++ b/apps/framework-cli/tests/psl/simple.prisma
@@ -1,5 +1,5 @@
 model User {
-  id      Int     
-  email   String   
-  name    String?
-}
\ No newline at end of file
+  id    Int     @id
+  email String
+  name  String?
+}