Fix/insert-into-connector-fix (#405)

* fixing db connector for redshift bigquery * PR comments resolve --------- Co-authored-by: Hari John Kuriakose <[email protected]>
Zipstack · Jun 20, 2024 · 49b9873 · 49b9873
1 parent 660ee19
commit 49b9873
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 18 deletions.
diff --git a/backend/workflow_manager/endpoint/constants.py b/backend/workflow_manager/endpoint/constants.py
@@ -92,9 +92,3 @@ class BigQuery:
     """
 
     TABLE_NAME_SIZE = 3
-    COLUMN_TYPES = [
-        "DATE",
-        "DATETIME",
-        "TIME",
-        "TIMESTAMP",
-    ]
diff --git a/backend/workflow_manager/endpoint/database_utils.py b/backend/workflow_manager/endpoint/database_utils.py
@@ -65,13 +65,6 @@ def get_sql_values_for_query(
                     sql_values[column] = f"parse_json($${values[column]}$$)"
                 else:
                     sql_values[column] = f"{values[column]}"
-            elif cls_name == DBConnectionClass.BIGQUERY:
-                col = column.lower()
-                type_x = column_types[col]
-                if type_x in BigQuery.COLUMN_TYPES:
-                    sql_values[column] = f"{type_x}({values[column]})"
-                else:
-                    sql_values[column] = f"{values[column]}"
             else:
                 # Default to Other SQL DBs
                 # TODO: Handle numeric types with no quotes

diff --git a/backend/workflow_manager/endpoint/tests/__init__.py b/backend/workflow_manager/endpoint/tests/__init__.py
@@ -1,3 +1,3 @@
-from backend.celery import app as celery_app
+from backend.celery_service import app as celery_app  # type: ignore
 
 __all__ = ["celery_app"]
diff --git a/unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py b/unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py
@@ -102,9 +102,10 @@ def execute_query(
         table_name = str(kwargs.get("table_name"))
         try:
             if sql_values:
-                engine.query(sql_query, job_config=sql_values)
+                query_job = engine.query(sql_query, job_config=sql_values)
             else:
-                engine.query(sql_query)
+                query_job = engine.query(sql_query)
+            query_job.result()
         except google.api_core.exceptions.Forbidden as e:
             logger.error(f"Forbidden exception in creating/inserting data: {str(e)}")
             raise BigQueryForbiddenException(

diff --git a/unstract/connectors/src/unstract/connectors/databases/redshift/redshift.py b/unstract/connectors/src/unstract/connectors/databases/redshift/redshift.py
@@ -76,12 +76,12 @@ def sql_to_db_mapping(value: str) -> str:
         """
         python_type = type(value)
         mapping = {
-            str: "SUPER",
+            str: "VARCHAR(65535)",
             int: "BIGINT",
             float: "DOUBLE PRECISION",
             datetime.datetime: "TIMESTAMP",
         }
-        return mapping.get(python_type, "SUPER")
+        return mapping.get(python_type, "VARCHAR(65535)")
 
     @staticmethod
     def get_create_table_query(table: str) -> str: