Firebolt Query Runner: now uses firebold-sdk python package (getredas…

…h#5689) * Added firebolt-sdk in place of firebolt-sqlalchemy * fixed connection issue * fixed connection issue * final commit * Moved firebolt-sdk's imports to try block Co-authored-by: rajeshSigmoid <rajeshk@sigmoidanalytics.com>
michifrey · Feb 1, 2022 · 26ac8ab · 26ac8ab
1 parent 12c4750
commit 26ac8ab
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/redash/query_runner/firebolt.py b/redash/query_runner/firebolt.py
@@ -1,5 +1,6 @@
 try:
-    from firebolt_db.firebolt_connector import Connection
+    from firebolt.db import connect
+    from firebolt.client import DEFAULT_API_URL
     enabled = True
 except ImportError:
     enabled = False
@@ -19,14 +20,14 @@ def configuration_schema(cls):
         return {
             "type": "object",
             "properties": {
-                "host": {"type": "string", "default": "localhost"},
-                "port": {"type": "string", "default": 8123},
+                "api_endpoint": {"type": "string", "default": DEFAULT_API_URL},
+                "engine_name": {"type": "string"},
                 "DB": {"type": "string"},
                 "user": {"type": "string"},
                 "password": {"type": "string"},
             },
-            "order": ["host", "port", "user", "password", "DB"],
-            "required": ["user","password"],
+            "order": ["user", "password", "api_endpoint", "engine_name", "DB"],
+            "required": ["user", "password", "engine_name", "DB"],
             "secret": ["password"],
         }
 
@@ -35,12 +36,12 @@ def enabled(cls):
         return enabled
 
     def run_query(self, query, user):
-        connection = Connection(
-            host=self.configuration["host"],
-            port=self.configuration["port"],
+        connection = connect(
+            api_endpoint=(self.configuration.get("api_endpoint") or DEFAULT_API_URL),
+            engine_name=(self.configuration.get("engine_name") or None),
             username=(self.configuration.get("user") or None),
             password=(self.configuration.get("password") or None),
-            db_name=(self.configuration.get("DB") or None),
+            database=(self.configuration.get("DB") or None),
         )
 
         cursor = connection.cursor()

diff --git a/requirements_all_ds.txt b/requirements_all_ds.txt
@@ -40,5 +40,5 @@ trino~=0.305
 cmem-cmempy==21.2.3
 xlrd==2.0.1
 openpyxl==3.0.7
-firebolt-sqlalchemy
+firebolt-sdk
 pandas==1.3.4