Repository: ToucanToco/toucan-connectors

Issue:
[bigquery] Delegate parameters handling to the library
See https://github.com/ToucanToco/toucan-connectors/pull/594#discussion_r870425994


Discussion:

Resolving patch:
diff --git a/toucan_connectors/google_big_query/google_big_query_connector.py b/toucan_connectors/google_big_query/google_big_query_connector.py
--- a/toucan_connectors/google_big_query/google_big_query_connector.py
+++ b/toucan_connectors/google_big_query/google_big_query_connector.py
@@ -6,6 +6,7 @@
 import pandas
 import pandas as pd
 from google.cloud import bigquery
+from google.cloud.bigquery.dbapi import _helpers as bigquery_helpers
 from google.oauth2.service_account import Credentials
 from pydantic import Field
 
@@ -30,34 +31,16 @@ class GoogleBigQueryDataSource(ToucanDataSource):
 BigQueryParam = Union[bigquery.ScalarQueryParameter, bigquery.ArrayQueryParameter]
 
 
-# NOTE: This does not play nicely with dates. They're a bit tricky
-# though, as we'd have to try and parse dates from strings to
-# determine if something is a date or not. Until then, we can just
-# use a cast. eg: SELECT * FROM table WHERE STRING(date_col) IN UNNEST({{my_dates}})
-def _define_scalar_type(value: Any) -> str:
-    if isinstance(value, bool):
-        return 'BOOL'
-    elif isinstance(value, int):
-        return 'NUMERIC'
-    elif isinstance(value, float):
-        return 'FLOAT64'
-    elif isinstance(value, str):
-        return 'STRING'
-    # TODO - check bad return type
-    return 'STRING'
-
-
-def _define_array_type(name: str, values: List[Any]) -> BigQueryParam:
-    return bigquery.ArrayQueryParameter(
-        name, _define_scalar_type(values[0] if len(values) > 0 else ''), values
-    )
-
-
 def _define_query_param(name: str, value: Any) -> BigQueryParam:
     if isinstance(value, list):
-        return _define_array_type(name, value)
+        return (
+            bigquery_helpers.array_to_query_parameter(value=value, name=name)
+            if len(value) > 0
+            # array_to_query_parameter raises an exception in case of an empty list
+            else bigquery.ArrayQueryParameter(name=name, array_type='STRING', values=value)
+        )
     else:
-        return bigquery.ScalarQueryParameter(name, _define_scalar_type(value), value)
+        return bigquery_helpers.scalar_to_query_parameter(value=value, name=name)
 
 
 class GoogleBigQueryConnector(ToucanConnector):


Test patch:
diff --git a/tests/google_big_query/test_google_big_query.py b/tests/google_big_query/test_google_big_query.py
--- a/tests/google_big_query/test_google_big_query.py
+++ b/tests/google_big_query/test_google_big_query.py
@@ -46,12 +46,12 @@ def _fixture_scope():
     'input_value,expected_output',
     [
         ('test', ScalarQueryParameter('test_param', 'STRING', 'test')),
-        (0, ScalarQueryParameter('test_param', 'NUMERIC', 0)),
+        (0, ScalarQueryParameter('test_param', 'INT64', 0)),
         (0.0, ScalarQueryParameter('test_param', 'FLOAT64', 0.0)),
         (True, ScalarQueryParameter('test_param', 'BOOL', True)),
         ([], ArrayQueryParameter('test_param', 'STRING', [])),
         (['hi'], ArrayQueryParameter('test_param', 'STRING', ['hi'])),
-        ([0], ArrayQueryParameter('test_param', 'NUMERIC', [0])),
+        ([0], ArrayQueryParameter('test_param', 'INT64', [0])),
         ([0.0, 2], ArrayQueryParameter('test_param', 'FLOAT64', [0.0, 2])),
         ([True, False], ArrayQueryParameter('test_param', 'BOOL', [True, False])),
     ],