Feature/clickhouse connector

chaos_genius/connectors/__init__.py

@@ -7,6 +7,7 @@
 from chaos_genius.connectors.druid import Druid
 from chaos_genius.connectors.aws_athena import AwsAthenaDb
 from chaos_genius.connectors.databricks import Databricks
+from chaos_genius.connectors.clickhouse import ClickHouseDb
 
 
 DB_CLASS_MAPPER = {
@@ -18,6 +19,7 @@
     "Druid": Druid,
     "AWS Athena": AwsAthenaDb,
     "Databricks": Databricks,
+    "ClickHouse": ClickHouseDb
 }
 
 

chaos_genius/connectors/clickhouse.py

@@ -0,0 +1,76 @@
+import pandas as pd
+from sqlalchemy import create_engine
+from sqlalchemy import text
+from .base_db import BaseDb
+from .connector_utils import merge_dataframe_chunks
+
+
+class ClickHouseDb(BaseDb):
+
+    __SQL_IDENTIFIER = "`"
+
+    @property
+    def sql_identifier(self):
+        """Used to quote any SQL identifier in case of it using special characters or keywords."""
+        return self.__SQL_IDENTIFIER
+
+    db_name = "ClickHouse"
+    test_db_query = "SELECT 1"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_db_uri(self):
+        db_info = self.ds_info
+        host = db_info.get("host")
+        port = int(db_info.get("port"))
+        username = db_info.get("username")
+        database = db_info.get("database")
+        password = db_info.get("password")
+        if not (host and port and username and password and database):
+            raise NotImplementedError(
+                "Database Credential not found for ClickHouse."
+            )
+        self.sqlalchemy_db_uri = (
+            f"clickhouse+native://{username}:{password}@{host}:{port}/{database}"
+        )
+        return self.sqlalchemy_db_uri
+
+    def get_db_engine(self):
+        db_uri = self.get_db_uri()
+        self.engine = create_engine(db_uri, echo=self.debug)
+        return self.engine
+
+    def test_connection(self):
+        if not hasattr(self, "engine") or not self.engine:
+            self.engine = self.get_db_engine()
+        query_text = text(self.test_db_query)
+        status, message = None, ""
+        try:
+            with self.engine.connect() as connection:
+                cursor = connection.execute(query_text)
+                results = cursor.all()
+                if results[0][0] == 1:
+                    status = True
+                else:
+                    status = False
+        except Exception as err_msg:
+            status = False
+            message = str(err_msg)
+        return status, message
+
+    def run_query(self, query, as_df=True):
+        engine = self.get_db_engine()
+        if as_df == True:
+            return merge_dataframe_chunks(
+                pd.read_sql_query(query, engine, chunksize=self.CHUNKSIZE)
+            )
+        else:
+            return []
+
+    def get_schema(self):
+        self.schema = self.ds_info.get("database")
+        return self.schema
+
+    def get_schema_names_list(self):
+        return None

chaos_genius/third_party/integration_server_config.py

@@ -66,6 +66,7 @@
     "5B45DB62-303C-4E70-92DA-419D3CDBD506": False,  # Druid
     "F2A65364-07BA-4542-BA1B-0805E5A0572D": False,  # AWS Athena
     "ba49e8e6-eed6-11ec-8ea0-0242ac120002": False,  # Databricks
+    "15323c3e-222d-11ed-861d-0242ac120002": False  # Clickhouse
     # "29b409d9-30a5-4cc8-ad50-886eb846fea3", # Quickbooks
 }
 
@@ -125,6 +126,14 @@
         "access_token": "access_token",
         "catalog": "catalog",
         "db_type": "databricks"
+    },
+    "15323c3e-222d-11ed-861d-0242ac120002": {
+        "host": "host",
+        "port": "port",
+        "database": "database",
+        "username": "username",
+        "password": "password",
+        "db_type": "ClickHouse"
     }
 }
 

chaos_genius/utils/metadata_api_config.py

@@ -6,7 +6,8 @@
     "Snowflake": True,
     "Druid": False,
     "AWS Athena": True,
-    "Databricks": True
+    "Databricks": True,
+    "ClickHouse": False
 }
 
 TABLE_VIEW_MATERIALIZED_VIEW_AVAILABILITY = {
@@ -66,6 +67,13 @@
         "supported_aggregations": ["mean", "sum", "count"],
         "supports_multidim_dd": True
     },
+    "ClickHouse": {
+        "tables": True,
+        "views": True,
+        "materialized_views": True,
+        "supported_aggregations": ["mean", "sum", "count"],
+        "supports_multidim_dd": True
+    },
 }
 
 TABLE_VIEW_MATERIALIZED_VIEW_AVAILABILITY_THIRD_PARTY = {

requirements/prod.txt

@@ -23,6 +23,7 @@ sqlalchemy-redshift==0.8.6
 pydruid[sqlalchemy]~=0.6.2
 PyAthena[SQLAlchemy,Pandas]==2.5.2
 sqlalchemy-databricks==0.2.0
+clickhouse-sqlalchemy==0.2.2
 
 # Migrations
 Flask-Migrate==2.7.0