[dy] Update Bigquery query conditions (mage-ai#2214)

ravediamond · Mar 17, 2023 · 01e581d · 01e581d
1 parent caf5de1
commit 01e581d
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/mage_integrations/mage_integrations/destinations/bigquery/__init__.py b/mage_integrations/mage_integrations/destinations/bigquery/__init__.py
@@ -5,6 +5,7 @@
     MAX_QUERY_BUFFER,
     MAX_QUERY_PARAMETERS,
     MAX_QUERY_PARAMETERS_SIZE,
+    MAX_QUERY_STRING_SIZE,
     QUERY_JOB_MAX_TIMEOUT_SECONDS,
 )
 from mage_integrations.destinations.bigquery.utils import (
@@ -512,15 +513,19 @@ def __insert_with_limits(
         jobs = []
         job_results = []
 
+        insert_statement = f"INSERT INTO {full_table_name} ({insert_columns}) VALUES"
+
         while len(insert_values) >= 1:
+            query_size = len(insert_statement)
             query_payload_size = 0
             query_parameters = []
 
             row_values = []
             row_idx = -1
 
             while query_payload_size < (MAX_QUERY_PARAMETERS_SIZE * MAX_QUERY_BUFFER) and \
-                len(query_parameters) < (MAX_QUERY_PARAMETERS * MAX_QUERY_BUFFER) and \
+                    len(query_parameters) < (MAX_QUERY_PARAMETERS * MAX_QUERY_BUFFER) and \
+                    query_size < (MAX_QUERY_STRING_SIZE * MAX_QUERY_BUFFER) and \
                     row_idx + 1 < len(insert_values):
 
                 row_idx += 1
@@ -551,11 +556,13 @@ def __insert_with_limits(
                         arr.append(f'@{variable_name}')
 
                         query_payload_size += sys.getsizeof(value)
-
-                row_values.append(f'({",".join(arr)})')
+
+                row_value = f'({",".join(arr)})'
+                query_size += len(row_value)
+                row_values.append(row_value)
 
             query_arr = [
-                f"INSERT INTO {full_table_name} ({insert_columns}) VALUES {','.join(row_values)};",
+                f"{insert_statement} {','.join(row_values)};",
             ]
 
             if count_rows: