Trino connector

chaos_genius/connectors/__init__.py

@@ -5,6 +5,7 @@
 from chaos_genius.connectors.snowflake import SnowflakeDb
 from chaos_genius.connectors.redshift import Redshift
 from chaos_genius.connectors.druid import Druid
+from chaos_genius.connectors.trino import Trino
 
 
 DB_CLASS_MAPPER = {
@@ -14,6 +15,7 @@
     "Snowflake": SnowflakeDb,
     "Redshift": Redshift,
     "Druid": Druid,
+    "Trino": Trino,
 }
 
 

chaos_genius/connectors/trino.py

@@ -0,0 +1,104 @@
+"""Apache Trino connector."""
+
+import pandas as pd
+from sqlalchemy import create_engine, text
+
+from chaos_genius.connectors.base_db import BaseDb
+
+from .connector_utils import merge_dataframe_chunks
+
+
+class Trino(BaseDb):
+    """Apache Trino connector."""
+
+    test_db_query = "SELECT 1"
+    __SQL_DATE_FORMAT = "timestamp '%Y-%m-%d 00:00:00{}'"
+    __SQL_STRPTIME_FORMAT = "timestamp '%Y-%m-%d %H:%M:%S%z'"
+    __SQL_STRFTIME_FORMAT = "timestamp '%Y-%m-%d %H:%M:%S'"
+
+    @property
+    def sql_date_format(self):
+        """String format to convert date to datetime along with an offset."""
+        return self.__SQL_DATE_FORMAT
+
+    @property
+    def sql_strptime_format(self):
+        """Format to convert strings into dates."""
+        return self.__SQL_STRPTIME_FORMAT
+
+    @property
+    def sql_strftime_format(self):
+        """Format to convert dates into strings."""
+        return self.__SQL_STRFTIME_FORMAT
+
+    def get_db_uri(self):
+        """Create SQLAlchemy URI from data source info."""
+        db_info = self.ds_info
+        if db_info is None:
+            raise Exception("Datasource info not found for Trino.")
+
+        host = db_info.get("host")
+        port = int(db_info.get("port"))
+        username = db_info.get("username")
+        catalog = db_info.get("catalog")
+        password = db_info.get("password")
+        if not (host and port and catalog):
+            raise Exception("Database Credential not found for Trino.")
+
+        if not (username and password):
+            self.sqlalchemy_db_uri = f"trino://{host}:{port}/{catalog}"
+        else:
+            self.sqlalchemy_db_uri = (
+                f"trino://{username}:{password}@{host}:{port}/{catalog}"
+            )
+        return self.sqlalchemy_db_uri
+
+    def get_db_engine(self):
+        """Create an SQLAlchemy engine from data source info."""
+        db_uri = self.get_db_uri()
+        self.engine = create_engine(
+            db_uri, echo=self.debug, connect_args={"http_scheme": "https", "verify": False}
+        )
+        return self.engine
+
+    def test_connection(self):
+        """Test data source connection."""
+        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:  # noqa: B902
+            status = False
+            message = str(err_msg)
+        return status, message
+
+    def run_query(self, query, as_df=True):
+        """Run a SQL query."""
+        engine = self.get_db_engine()
+        if as_df:
+            return merge_dataframe_chunks(
+                pd.read_sql_query(query, engine, chunksize=self.CHUNKSIZE)
+            )
+        else:
+            return []
+
+    def get_schema(self):
+        """Get schema name."""
+        schema_name = self.ds_info.get("schema") if self.ds_info is not None else None
+        if schema_name:
+            self.schema = schema_name
+        else:
+            self.schema = "public"
+        return self.schema
+
+    def get_schema_names_list(self):
+        data = self.inspector.get_schema_names()
+        return data

chaos_genius/core/utils/kpi_validation.py

@@ -40,7 +40,7 @@ def validate_kpi(kpi_info: Dict[str, Any], check_tz_aware: bool = False) -> Tupl
     connection_info = DataSource.get_by_id(
         kpi_info["data_source"]
     ).as_dict
-    supports_date_string_parsing = connection_info["name"] == "Druid"
+    supports_date_string_parsing = connection_info["connection_type"] in {"Druid", "Trino"}
 
     status, message = _validate_kpi_from_df(
         df,

chaos_genius/third_party/integration_server_config.py

@@ -64,6 +64,7 @@
     "e2d65910-8c8b-40a1-ae7d-ee2416b2bfa2": False,  # Snowflake
     "e87ffa8e-a3b5-f69c-9076-6011339de1f6": False,  # Redshift
     "5B45DB62-303C-4E70-92DA-419D3CDBD506": False,  # Druid
+    "ea6b0872-b25b-4591-8829-7e46d3276a5b": False,  # Trino
     # "29b409d9-30a5-4cc8-ad50-886eb846fea3", # Quickbooks
 }
 
@@ -116,6 +117,14 @@
         "username": "username",
         "password": "password",
         "db_type": "druid"
+    },
+    "ea6b0872-b25b-4591-8829-7e46d3276a5b": {
+        "host": "host",
+        "port": "port",
+        "username": "username",
+        "password": "password",
+        "catalog": "catalog",
+        "db_type": "trino"
     }
 }
 

chaos_genius/utils/metadata_api_config.py

@@ -4,7 +4,8 @@
     "Redshift": True,
     "BigQuery": False,
     "Snowflake": True,
-    "Druid": False
+    "Druid": False,
+    "Trino": True,
 }
 
 TABLE_VIEW_MATERIALIZED_VIEW_AVAILABILITY = {
@@ -49,6 +50,13 @@
         "materialized_views": True,
         "supported_aggregations": ["sum", "count"],
         "supports_multidim_dd": False
+    },
+    "Trino": {
+        "tables": True,
+        "views": True,
+        "materialized_views": True,
+        "supported_aggregations": ["mean", "sum", "count"],
+        "supports_multidim_dd": True,
     }
 }
 
@@ -66,4 +74,4 @@
     if conf["supports_multidim_dd"]
 ]
 
-NON_THIRD_PARTY_DATASOURCES = TABLE_VIEW_MATERIALIZED_VIEW_AVAILABILITY.keys()
+NON_THIRD_PARTY_DATASOURCES = TABLE_VIEW_MATERIALIZED_VIEW_AVAILABILITY.keys()

requirements/prod.txt

@@ -21,6 +21,8 @@ snowflake-sqlalchemy==1.2.4
 sqlalchemy-redshift==0.8.6
 # For apache druid
 pydruid[sqlalchemy]~=0.6.2
+# For Trino
+trino[sqlalchemy]==0.313.0
 
 # Migrations
 Flask-Migrate==2.7.0