From 8e86fc3fa881246981ab63561d5d67e012035d89 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 10 Dec 2025 10:34:30 -0500 Subject: [PATCH 01/19] starting the bigquery connector --- .../sources/google_bigquery/__init__.py | 9 + .../sources/google_bigquery/client.py | 48 +++ .../sources/google_bigquery/datasource.py | 316 ++++++++++++++++++ .../tests/sources/test_google_bigquery.py | 113 +++++++ 4 files changed, 486 insertions(+) create mode 100644 app/connectors_service/connectors/sources/google_bigquery/__init__.py create mode 100644 app/connectors_service/connectors/sources/google_bigquery/client.py create mode 100644 app/connectors_service/connectors/sources/google_bigquery/datasource.py create mode 100644 app/connectors_service/tests/sources/test_google_bigquery.py diff --git a/app/connectors_service/connectors/sources/google_bigquery/__init__.py b/app/connectors_service/connectors/sources/google_bigquery/__init__.py new file mode 100644 index 000000000..73a3152b0 --- /dev/null +++ b/app/connectors_service/connectors/sources/google_bigquery/__init__.py @@ -0,0 +1,9 @@ +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License 2.0; +# you may not use this file except in compliance with the Elastic License 2.0. +# +from .client import GoogleBigqueryClient +from .datasource import GoogleBigqueryDataSource + +__all__ = ["GoogleBigqueryClient", "GoogleBigqueryDataSource"] diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py new file mode 100644 index 000000000..20ba630db --- /dev/null +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -0,0 +1,48 @@ +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License 2.0; +# you may not use this file except in compliance with the Elastic License 2.0. +# +"""Google Bigquery module which fetches rows from a Bigquery table.""" + +from aiogoogle.auth.creds import ServiceAccountCreds +from connectors_sdk.logger import logger +from google.cloud import bigquery +from google.oauth2 import service_account +import os +import tempfile + +RUNNING_FTEST = ( + "RUNNING_FTEST" in os.environ +) # Flag to check if a connector is run for ftest or not. + +class GoogleBigqueryClient: + """A client to make api calls to Google Bigquery.""" + + def __init__(self, json_credentials): + """Initialize the ServiceAccountCreds instance. + + Args: + json_credentials(dict): Service account credentials json.""" + self.json_credentials = json_credentials + self._logger = logger + + + def client(self, project_id=None): + """Returns an instance of a bigquery client, using the configured credentials, + optionally to a different project_id (configured creds must have permissions to + access and create query jobs in it). + + Args: + project_id (string, optional): If connecting to a project other than the + service account credentials default, pass this as a string. + + Returns: + bigquery.Client instance + + """ + credentials = service_account.Credentials.from_service_account_info(self.json_credentials) + if project_id is not None: + return bigquery.Client(credentials=credentials, project=project_id) + else: + return bigquery.Client(credentials=credentials) diff --git a/app/connectors_service/connectors/sources/google_bigquery/datasource.py b/app/connectors_service/connectors/sources/google_bigquery/datasource.py new file mode 100644 index 000000000..875aae7da --- /dev/null +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -0,0 +1,316 @@ +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License 2.0; +# you may not use this file except in compliance with the Elastic License 2.0. +# + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from connectors_sdk.source import ( + BaseDataSource, + DataSourceConfiguration +) +from connectors.sources.google_bigquery.client import GoogleBigqueryClient +from connectors.sources.shared.database.generic_database import ( + DEFAULT_FETCH_SIZE, + DEFAULT_RETRY_COUNT +) +from connectors.sources.shared.google import ( + load_service_account_json, + validate_service_account_json, +) +from connectors.utils import get_pem_format + +from functools import cached_property, partial +import os + +RUNNING_FTEST = ( + "RUNNING_FTEST" in os.environ +) # Flag to check if a connector is run for ftest or not. + +executor = ThreadPoolExecutor(1) + +class GoogleBigqueryDataSource(BaseDataSource): + """Google Bigquery""" + + name = "Google Bigquery" + service_type = "google_bigquery" + incremental_sync_enabled = False + advanced_rules_enabled = True + dls_enabled = False + + + def __init__(self, configuration: DataSourceConfiguration): + """Set up the connection to Google Bigquery. + + Args: + configuration (DataSourceConfiguration): Instance of DataSourceConfiguration. + """ + super().__init__(configuration=configuration) + + def _set_internal_logger(self): + self._google_bigquery_client.set_logger(self._logger) + + @cached_property + def _google_bigquery_client(self) -> GoogleBigqueryClient: + """Initialize and return an instance of GoogleBigqueryClient + + Returns: + GoogleBigqueryClient: An instance of GoogleBigqueryClient. + """ + REQUIRED_CREDENTIAL_KEYS = [ + "type", + "project_id", + "private_key_id", + "private_key", + "client_email", + "client_id", + "auth_uri", + "token_uri", + ] + + json_credentials = load_service_account_json( + self.configuration["service_account_credentials"], "Google Bigquery" + ) + + if ( + json_credentials.get("private_key") + and "\n" not in json_credentials["private_key"] + ): + json_credentials["private_key"] = get_pem_format( + key=json_credentials["private_key"].strip(), + postfix="-----END PRIVATE KEY-----", + ) + + required_credentials = { + key: value + for key, value in json_credentials.items() + if key in REQUIRED_CREDENTIAL_KEYS + } + + return GoogleBigqueryClient(json_credentials=required_credentials) + + + @classmethod + def get_default_configuration(cls) -> dict: + """Get the default configuration for Google Bigquery. + + Returns: + dictionary: Default configuration. + """ + return { + "service_account_credentials": { + "display": "textarea", + "label": "Google Cloud service account JSON", + "sensitive": True, + "order": 1, + "type": "str" + }, + "dataset": { + "display": "text", + "label": "Dataset the table is in.", + "order": 2, + "type": "str", + }, + "table": { + "display": "text", + "label": "Table to sync.", + "order": 3, + "type": "str", + }, + "project_id": { + "display": "text", + "label": "Google Cloud project.", + "order": 4, + "type": "str", + "required": False, + "default_value": "", + "tooltip": "Defaults to the service account project.", + }, + "index_settings": { + "display": "textarea", + "label": "Additional index settings, if any. For example, to set mode to lookup.", + "order": 5, + "required": False, + "default_value": "", + "type": "str", + "ui_restrictions": ["advanced"], + }, + "columns": { + "display": "textarea", + "label": "Columns to fetch. Defaults to * if none are set.", + "order": 6, + "required": False, + "default_value": "*", + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "Comma-separated, as in a SQL SELECT." + }, + "predicates": { + "display": "textarea", + "label": "Predicates for the query.", + "order": 7, + "required": False, + "default_value": "", + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "A SQL WHERE clause. May be required for some partitioned table configurations." + }, + "fetch_size": { + "default_value": DEFAULT_FETCH_SIZE, + "display": "numeric", + "label": "Rows fetched per request", + "order": 8, + "required": False, + "type": "int", + "ui_restrictions": ["advanced"], + }, + "retry_count": { + "default_value": DEFAULT_RETRY_COUNT, + "display": "numeric", + "label": "Retries per request", + "order": 9, + "required": False, + "type": "int", + "ui_restrictions": ["advanced"], + }, + } + + async def validate_config(self): + """Validates whether user inputs are valid for configuration. + + Raises: + Exception: The format of service account json is invalid. + """ + await super().validate_config() + validate_service_account_json( + self.configuration["service_account_credentials"], + "Google Bigquery" + ) + + # if they set a project_id and have a wrong service account, a log message + # will give them a fighting chance :) + if self.configuration["project_id"] and self.configuration["service_account_credentials"]["project_id"] != self.configuration["project_id"]: + self._logger.info("A project_id is configured and does not match the project_id for the service_account_credentials block. If authorization fails, this could be why!") + self.project_id = self.resolve_project() + + async def ping(self): + """Verify the connection with Google Bigquery""" + if RUNNING_FTEST: + return + try: + await anext( + self._google_bigquery_client.api_call( + resource="projects", + method="serviceAccount", + sub_method="get", + projectId=self._google_bigquery_client.user_project_id)) + except Exception: + self._logger.exception("Error while connecting to Google Bigquery.") + raise + + + def resolve_project(self): + """ + Resolves a project_id, as a string. This should be a Google Cloud project. + Chooses configured project_id first, or if that is unset, falls back to the + project_id on the service account. + + Returns: + string: The project name. + + Raises: + ConfigurableFieldValueError: service account JSON was invalid/unparseable. + """ + + if self.configuration["project_id"]: + return self.configuration["project_id"] + json_credentials = load_service_account_json( + self.configuration["service_account_credentials"], "Google Cloud Storage" + ) + return json_credentials["project_id"] + + + def resolve_table(self): + """Inspects the configuration and produces a fully qualified BigQuery table + identifier. + + Returns: + string: A BQ table in the form `project.dataset.table` + """ + + return "`%s.%s.%s`" % (self.resolve_project(), self.configuration["dataset"], self.configuration["table"]) + + + def build_query(self): + """ + Builds the query that will be run on BigQuery. + + Returns: + string: The query. + """ + # create a sub-config because full config contains secrets + conf = {k: self.configuration[k] for k in ("predicates", "columns")} + conf["resolved_table"] = self.resolve_table() + if not conf["columns"]: + conf["columns"] = "*" + + query = """SELECT %(columns)s FROM %(resolved_table)s""" % conf + + if conf["predicates"]: + query = query + " " + conf["predicates"] + return query.strip() + + + async def get_docs(self, filtering=None): + """Returns results as (rowdict,None) on the configured query results. Realizes + results in fetch_size chunks. + + Args: + filtering: Unused; part of the BaseDataSource contract. + + Yields: + dictionary: (rowdict, None) pairs, per the BaseDataSource contract. + + """ + self._logger.info("Connected to Google Bigquery.") + sql = self.build_query() + + # job is a QueryJob instance + # https://docs.cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob + job = self._google_bigquery_client.client().query(sql) + + # do not block on the bigquery job completing + loop = asyncio.get_running_loop() + await loop.run_in_executor(executor, job.done) + + fetch_size = self.configuration.get("fetch_size") or DEFAULT_FETCH_SIZE + + def get_next_chunk(job_iter): + try: + chunk = [] + for _ in range(fetch_size): + try: + row = next(job_iter) + chunk.append(dict(row)) # Row -> dict + except StopIteration: + break + return chunk + except Exception as e: + # TODO: retry_count support here? + self._logger.error(f"Error fetching chunk from BigQuery job: {e}") + raise e + + # we don't use .result() as it is synchronous + # as of this writing, awaitable jobs are not yet implemented + # which is why this asyncio machinery is necessary presently + # but in a future version this should all become awaitable + # https://github.com/googleapis/python-bigquery/issues/18 + job_iter = iter(job) + + while True: + chunk = await loop.run_in_executor(executor, get_next_chunk, job_iter) + if not chunk: + break + for row_dict in chunk: + yield row_dict, None diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py new file mode 100644 index 000000000..eca4b6d1b --- /dev/null +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -0,0 +1,113 @@ +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License 2.0; +# you may not use this file except in compliance with the Elastic License 2.0. +# +"""Tests the Google Bigquery source.""" + +import asyncio +from contextlib import asynccontextmanager +import json +import pytest +from unittest import mock +from unittest.mock import Mock, patch + +from connectors.sources.google_bigquery import GoogleBigqueryDataSource +from connectors_sdk.source import ConfigurableFieldValueError, DataSourceConfiguration +from tests.sources.support import create_source + +SERVICE_ACCOUNT_CREDENTIALS = '{"project_id": "dummy123"}' + +# a very small default table to use, ~496 rows +test_bq_small = { + "project_id": "bigquery-public-data", + "dataset": "new_york_subway", + "table": "stations" +} + +# a >10k table to use, 19640 to be exact +test_bq_10k = { + "project_id": "bigquery-public-data", + "dataset": "new_york_subway", + "table": "trips" +} + +@asynccontextmanager +async def create_bigquery_source( + service_account_credentials=SERVICE_ACCOUNT_CREDENTIALS, + dataset=test_bq_small["dataset"], + table=test_bq_small["table"] +): + async with create_source( + GoogleBigqueryDataSource, + service_account_credentials=SERVICE_ACCOUNT_CREDENTIALS, + dataset=dataset, + table=table + ) as source: + yield source + +@pytest.mark.asyncio +async def test_empty_configuration(): + """Tests GoogleBigqueryDataSource's validation of the provided configuration.""" + error_configs = [ + # service account credentials must be present and not blank + DataSourceConfiguration({"service_account_credentials": ""}) + + ] + + for error_config in error_configs: + bq = GoogleBigqueryDataSource(configuration=error_config) + with pytest.raises(ConfigurableFieldValueError): + await bq.validate_config() + +@pytest.mark.asyncio +async def test_build_query(): + """Tests query building for GoogleBigqueryDataSource.""" + testcreds = json.dumps({"project_id": "testprojectid"}) + + # minimum defaults + config = DataSourceConfiguration( + {"service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + }) + bq = GoogleBigqueryDataSource(configuration=config) + query = bq.build_query() + expected = """SELECT * FROM `testprojectid.testdataset.testtable`""" + assert query == expected + + # custom project + config = DataSourceConfiguration( + {"service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + "project_id": "different", + }) + bq = GoogleBigqueryDataSource(configuration=config) + query = bq.build_query() + expected = """SELECT * FROM `different.testdataset.testtable`""" + assert query == expected + + # custom columns + config = DataSourceConfiguration( + {"service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + "columns": "account_id, timestamp" + }) + bq = GoogleBigqueryDataSource(configuration=config) + query = bq.build_query() + expected = """SELECT account_id, timestamp FROM `testprojectid.testdataset.testtable`""" + assert query == expected + + # custom predicates + config = DataSourceConfiguration( + {"service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + "predicates": "WHERE foo=1" + }) + bq = GoogleBigqueryDataSource(configuration=config) + query = bq.build_query() + expected = """SELECT * FROM `testprojectid.testdataset.testtable` WHERE foo=1""" + assert query == expected From 706fca5839ffa5bc2f17a3fd5ae1b500f8e37e40 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Thu, 18 Dec 2025 15:34:04 -0500 Subject: [PATCH 02/19] google-cloud-bigquery pyproject dep --- app/connectors_service/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/app/connectors_service/pyproject.toml b/app/connectors_service/pyproject.toml index 657c7d06e..03a602008 100644 --- a/app/connectors_service/pyproject.toml +++ b/app/connectors_service/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "exchangelib==5.4.0", "fastjsonschema==2.16.2", "gidgethub==5.2.1", + "google-cloud-bigquery==3.39.0", "graphql-core==3.2.3", "httpx-ntlm==1.4.0", "httpx==0.27.0", From cab3a087cd9151a9ff58a215c14722cd028413a0 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Tue, 13 Jan 2026 15:37:35 -0500 Subject: [PATCH 03/19] register google_bigquery as an available service_type --- app/connectors_service/connectors/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/app/connectors_service/connectors/config.py b/app/connectors_service/connectors/config.py index 606548f92..5d7995d39 100644 --- a/app/connectors_service/connectors/config.py +++ b/app/connectors_service/connectors/config.py @@ -119,6 +119,7 @@ def _default_config(): "github": "connectors.sources.github:GitHubDataSource", "gitlab": "connectors.sources.gitlab:GitLabDataSource", "gmail": "connectors.sources.gmail:GMailDataSource", + "google_bigquery": "connectors.sources.google_bigquery:GoogleBigqueryDataSource", "google_cloud_storage": "connectors.sources.google_cloud_storage:GoogleCloudStorageDataSource", "google_drive": "connectors.sources.google_drive:GoogleDriveDataSource", "graphql": "connectors.sources.graphql:GraphQLDataSource", From 3e0f790bce3c4e34b7a8d05cb9154e93e24f4f78 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Tue, 13 Jan 2026 15:41:41 -0500 Subject: [PATCH 04/19] remove index_settings config option MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unfortunately, indexes are created as part of the pre-setup workflow in kibana. The user has no opportunity to provide custom settings. In fact, if you want custom settings (like we do, for "index.mode": "lookup"), you have to let it create the index, pause before config, DELETE the index it created, and create a new one with the same name, set up the way you want it. You can't even pre-create it. The workflow fails on "index already exists." ಠ_ಠ Anyway yeah so this can't actually happen here so I removed it. fix project_id validation in terms of resolve_project() use new style querying to implement ping() implement row2dict() - allow a doc_id_column config for the ES doc _id - use a 22 char urlsafe uuid4 if none is configured, same as ES would do - allow timestamp_column to provide ES doc _timestamp - use run start timestamp if none is configured implement set_logger() on client wrapper --- .../sources/google_bigquery/client.py | 2 + .../sources/google_bigquery/datasource.py | 103 ++++++++++++++---- 2 files changed, 83 insertions(+), 22 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py index 20ba630db..e14b43993 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/client.py +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -27,6 +27,8 @@ def __init__(self, json_credentials): self.json_credentials = json_credentials self._logger = logger + def set_logger(self, logger_): + self._logger = logger_ def client(self, project_id=None): """Returns an instance of a bigquery client, using the configured credentials, diff --git a/app/connectors_service/connectors/sources/google_bigquery/datasource.py b/app/connectors_service/connectors/sources/google_bigquery/datasource.py index 875aae7da..70b059ee3 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/datasource.py +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -20,9 +20,10 @@ validate_service_account_json, ) from connectors.utils import get_pem_format - +from datetime import datetime, timezone from functools import cached_property, partial import os +import uuid RUNNING_FTEST = ( "RUNNING_FTEST" in os.environ @@ -127,15 +128,6 @@ def get_default_configuration(cls) -> dict: "default_value": "", "tooltip": "Defaults to the service account project.", }, - "index_settings": { - "display": "textarea", - "label": "Additional index settings, if any. For example, to set mode to lookup.", - "order": 5, - "required": False, - "default_value": "", - "type": "str", - "ui_restrictions": ["advanced"], - }, "columns": { "display": "textarea", "label": "Columns to fetch. Defaults to * if none are set.", @@ -144,12 +136,32 @@ def get_default_configuration(cls) -> dict: "default_value": "*", "type": "str", "ui_restrictions": ["advanced"], - "tooltip": "Comma-separated, as in a SQL SELECT." + "tooltip": "Comma-separated, as in a SQL SELECT.", + }, + "doc_id_column": { + "display": "text", + "label": "Document _id column", + "order": 7, + "required": False, + "default_value": None, + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "Use the value of this column as the ES document _id instead of generating a UUID.", + }, + "timestamp_column": { + "display": "text", + "label": "Timestamp column", + "order": 8, + "required": False, + "default_value": None, + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "Use the value of this column as the ES document _timestamp instead of using the sync start time." }, "predicates": { "display": "textarea", "label": "Predicates for the query.", - "order": 7, + "order": 9, "required": False, "default_value": "", "type": "str", @@ -160,7 +172,7 @@ def get_default_configuration(cls) -> dict: "default_value": DEFAULT_FETCH_SIZE, "display": "numeric", "label": "Rows fetched per request", - "order": 8, + "order": 10, "required": False, "type": "int", "ui_restrictions": ["advanced"], @@ -169,7 +181,7 @@ def get_default_configuration(cls) -> dict: "default_value": DEFAULT_RETRY_COUNT, "display": "numeric", "label": "Retries per request", - "order": 9, + "order": 11, "required": False, "type": "int", "ui_restrictions": ["advanced"], @@ -190,7 +202,7 @@ async def validate_config(self): # if they set a project_id and have a wrong service account, a log message # will give them a fighting chance :) - if self.configuration["project_id"] and self.configuration["service_account_credentials"]["project_id"] != self.configuration["project_id"]: + if self.configuration["project_id"] and self.resolve_project() != self.configuration["project_id"]: self._logger.info("A project_id is configured and does not match the project_id for the service_account_credentials block. If authorization fails, this could be why!") self.project_id = self.resolve_project() @@ -199,12 +211,9 @@ async def ping(self): if RUNNING_FTEST: return try: - await anext( - self._google_bigquery_client.api_call( - resource="projects", - method="serviceAccount", - sub_method="get", - projectId=self._google_bigquery_client.user_project_id)) + sql = "SELECT 1=1" + job = self._google_bigquery_client.client().query(sql) + return job.done() # signal we can indeed run queries except Exception: self._logger.exception("Error while connecting to Google Bigquery.") raise @@ -261,6 +270,54 @@ def build_query(self): query = query + " " + conf["predicates"] return query.strip() + def url_safe_uuid(self): + return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip('b=').decode('ascii') + + def generate_doc_id(self, row): + """Creates and returns a string suitable for use as a doc _id. If the user + configured a doc_id_column to use for this, uses that from the row. If no id is + configured, documents will be assigned a random uuid, similar to what + Elasticsearch itself would do. + + Args: + row (dict): A row from BQ + + Returns: + str: A string for the _id + + """ + if self.configuration["doc_id_column"]: + return row.get(self.configuration["doc_id_column"]) + return self.url_safe_uuid() + + def generate_doc_timestamp(self, row): + """Creates and returns a timestamp string. If the user configured a + timestamp_column from their table, uses that. If not, returns None, and current + sync run start time in UTC will be assigned. + + Args: + row (dict): A row from BQ + + Returns: + str, or None: a timestamp string + """ + if self.configuration["timestamp_column"]: + return row.get(self.configuration["timestamp_column"]) + return None + + def row2doc(self, row): + doc = dict(row) + doc_id = self.generate_doc_id(row) + doc_timestamp = self.generate_doc_timestamp(row) + if doc_timestamp is None: + doc_timestamp = self._run_timestamp + row.update( + { + "_id_": doc_id, + "_timestamp": doc_timestamp, + } + ) + return row async def get_docs(self, filtering=None): """Returns results as (rowdict,None) on the configured query results. Realizes @@ -274,6 +331,8 @@ async def get_docs(self, filtering=None): """ self._logger.info("Connected to Google Bigquery.") + # collect the start time of the sync, used when no timestamp_column is configured + self._run_timestamp = datetime.now(timezone.utc).isoformat() sql = self.build_query() # job is a QueryJob instance @@ -292,7 +351,7 @@ def get_next_chunk(job_iter): for _ in range(fetch_size): try: row = next(job_iter) - chunk.append(dict(row)) # Row -> dict + chunk.append(row2doc(row)) # Row -> dict except StopIteration: break return chunk From 694c5e8a0f853cdc7cb26a8eada3ce38454d3977 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Thu, 15 Jan 2026 13:57:50 -0500 Subject: [PATCH 05/19] assorted fixes and cleanups --- .../sources/google_bigquery/datasource.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/datasource.py b/app/connectors_service/connectors/sources/google_bigquery/datasource.py index 70b059ee3..7f99679d4 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/datasource.py +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -5,6 +5,7 @@ # import asyncio +import base64 from concurrent.futures import ThreadPoolExecutor from connectors_sdk.source import ( BaseDataSource, @@ -271,7 +272,7 @@ def build_query(self): return query.strip() def url_safe_uuid(self): - return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip('b=').decode('ascii') + return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip(b'=').decode('ascii') def generate_doc_id(self, row): """Creates and returns a string suitable for use as a doc _id. If the user @@ -306,18 +307,18 @@ def generate_doc_timestamp(self, row): return None def row2doc(self, row): - doc = dict(row) + doc = dict(row) # cast to dict from the google class doc_id = self.generate_doc_id(row) doc_timestamp = self.generate_doc_timestamp(row) if doc_timestamp is None: doc_timestamp = self._run_timestamp - row.update( + doc.update( { - "_id_": doc_id, + "_id": doc_id, "_timestamp": doc_timestamp, } ) - return row + return doc async def get_docs(self, filtering=None): """Returns results as (rowdict,None) on the configured query results. Realizes @@ -351,7 +352,7 @@ def get_next_chunk(job_iter): for _ in range(fetch_size): try: row = next(job_iter) - chunk.append(row2doc(row)) # Row -> dict + chunk.append(self.row2doc(row)) # Row -> dict except StopIteration: break return chunk From b43b97d171422b0f01536d50252a29fbe14cd2c3 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Tue, 20 Jan 2026 10:49:09 -0500 Subject: [PATCH 06/19] client mocks and test_ping() test_get_docs() test_ping_negative() --- .../tests/sources/test_google_bigquery.py | 64 ++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py index eca4b6d1b..49756abe1 100644 --- a/app/connectors_service/tests/sources/test_google_bigquery.py +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -12,7 +12,10 @@ from unittest import mock from unittest.mock import Mock, patch -from connectors.sources.google_bigquery import GoogleBigqueryDataSource +from connectors.sources.google_bigquery import ( + GoogleBigqueryClient, + GoogleBigqueryDataSource, +) from connectors_sdk.source import ConfigurableFieldValueError, DataSourceConfiguration from tests.sources.support import create_source @@ -46,6 +49,42 @@ async def create_bigquery_source( ) as source: yield source +class Credentials: + """Creates a mocked google oauth2 service_account Credentials instance.""" + def __init__(self, creds): + self.json_credentials = creds + + +class Job: + """Creates a mocked google bigquery Job. Job is an iterator. Iteration returns a + fixture result a few times and then stops.""" + def __init__(self, *args, **kw): + self.doc_count = 5 + + def __iter__(self): + return self + + def __next__(self): + if self.doc_count < 1: + raise StopIteration + self.doc_count = self.doc_count - 1 + return {"foo": "bar"} + + def done(self): + return True + + +class Client: + """This class creates a mocked bigquery client.""" + def __init__(self, *args, **kw): + self.credentials = kw.get('credentials') + self.project = kw.get('project') + + def query(self, sql): + """Returns mocked Job.""" + return Job() + + @pytest.mark.asyncio async def test_empty_configuration(): """Tests GoogleBigqueryDataSource's validation of the provided configuration.""" @@ -111,3 +150,26 @@ async def test_build_query(): query = bq.build_query() expected = """SELECT * FROM `testprojectid.testdataset.testtable` WHERE foo=1""" assert query == expected +@pytest.mark.asyncio +async def test_ping(): + async with create_bigquery_source() as source: + with patch.object(GoogleBigqueryClient, "client", return_value=Client()): + await source.ping() + +@pytest.mark.asyncio +@patch("connectors.utils.time_to_sleep_between_retries", Mock(return_value=0)) +async def test_ping_negative(): + with pytest.raises(Exception): + async with create_bigquery_source() as source: + with patch.object(GoogleBigqueryClient, "client", return_value=Client()): + with patch.object(Client, "query", side_effect=Exception()): + await source.ping() + +@pytest.mark.asyncio +async def test_get_docs(): + async with create_bigquery_source() as source: + with patch.object(GoogleBigqueryClient, "client", return_value=Client()): + docs = [] + async for doc, _ in source.get_docs(): + docs.append(doc) + assert len(docs) == 5 From 49f8e93099784374e1e825cdf8c0da24fa7ebf02 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Tue, 20 Jan 2026 12:40:49 -0500 Subject: [PATCH 07/19] provides a commented-out unit test against live real bigquery --- .../tests/sources/test_google_bigquery.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py index 49756abe1..caa6e6075 100644 --- a/app/connectors_service/tests/sources/test_google_bigquery.py +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -8,6 +8,7 @@ import asyncio from contextlib import asynccontextmanager import json +import os import pytest from unittest import mock from unittest.mock import Mock, patch @@ -150,6 +151,8 @@ async def test_build_query(): query = bq.build_query() expected = """SELECT * FROM `testprojectid.testdataset.testtable` WHERE foo=1""" assert query == expected + + @pytest.mark.asyncio async def test_ping(): async with create_bigquery_source() as source: @@ -173,3 +176,29 @@ async def test_get_docs(): async for doc, _ in source.get_docs(): docs.append(doc) assert len(docs) == 5 + +# To test against live bigquery, you can uncomment this and: +# - have an active google cloud service account with BQ execute permission +# - set up those creds and provide the location in GOOGLE_APPLICATION_CREDENTIALS +# - run the tests with fail-slow increased to 10 (live BQ isn't 'under 1s' fast) + +# @pytest.mark.asyncio +# async def test_get_docs_from_live(): +# """Tests get_docs against the small bq fixture.""" +# credsfile = os.getenv("GOOGLE_APPLICATION_CREDENTIALS") +# if credsfile is None: +# print("GOOGLE_APPLICATION_CREDENTIALS was unset; skipping test_get_docs") +# return +# +# with open(credsfile, 'r') as gac: +# creds = gac.read() +# +# config = DataSourceConfiguration( +# test_bq_small | {"service_account_credentials": creds,} +# ) +# +# bq = GoogleBigqueryDataSource(configuration=config) +# docs = [] +# async for doc, _ in bq.get_docs(): +# docs.append(doc) +# assert len(docs) > 0 From 27c1f4a5fa5519996650813e0b0a5e2df26079d0 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Tue, 20 Jan 2026 12:41:44 -0500 Subject: [PATCH 08/19] cleanups fix typo unused as of now --- .../connectors/sources/google_bigquery/client.py | 3 --- app/connectors_service/tests/sources/test_google_bigquery.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py index e14b43993..1a45803d6 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/client.py +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -12,9 +12,6 @@ import os import tempfile -RUNNING_FTEST = ( - "RUNNING_FTEST" in os.environ -) # Flag to check if a connector is run for ftest or not. class GoogleBigqueryClient: """A client to make api calls to Google Bigquery.""" diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py index caa6e6075..3b1aaf639 100644 --- a/app/connectors_service/tests/sources/test_google_bigquery.py +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -187,7 +187,7 @@ async def test_get_docs(): # """Tests get_docs against the small bq fixture.""" # credsfile = os.getenv("GOOGLE_APPLICATION_CREDENTIALS") # if credsfile is None: -# print("GOOGLE_APPLICATION_CREDENTIALS was unset; skipping test_get_docs") +# print("GOOGLE_APPLICATION_CREDENTIALS was unset; skipping test_get_docs_from_live") # return # # with open(credsfile, 'r') as gac: From 586e33ce007768b1d9ffe9596b1ab700549f9667 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 21 Jan 2026 12:00:32 -0500 Subject: [PATCH 09/19] underscore prefixes for the internals --- .../sources/google_bigquery/datasource.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/datasource.py b/app/connectors_service/connectors/sources/google_bigquery/datasource.py index 7f99679d4..17f29652e 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/datasource.py +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -203,9 +203,9 @@ async def validate_config(self): # if they set a project_id and have a wrong service account, a log message # will give them a fighting chance :) - if self.configuration["project_id"] and self.resolve_project() != self.configuration["project_id"]: + if self.configuration["project_id"] and self._resolve_project() != self.configuration["project_id"]: self._logger.info("A project_id is configured and does not match the project_id for the service_account_credentials block. If authorization fails, this could be why!") - self.project_id = self.resolve_project() + self.project_id = self._resolve_project() async def ping(self): """Verify the connection with Google Bigquery""" @@ -220,7 +220,7 @@ async def ping(self): raise - def resolve_project(self): + def _resolve_project(self): """ Resolves a project_id, as a string. This should be a Google Cloud project. Chooses configured project_id first, or if that is unset, falls back to the @@ -241,7 +241,7 @@ def resolve_project(self): return json_credentials["project_id"] - def resolve_table(self): + def _resolve_table(self): """Inspects the configuration and produces a fully qualified BigQuery table identifier. @@ -249,7 +249,7 @@ def resolve_table(self): string: A BQ table in the form `project.dataset.table` """ - return "`%s.%s.%s`" % (self.resolve_project(), self.configuration["dataset"], self.configuration["table"]) + return "`%s.%s.%s`" % (self._resolve_project(), self.configuration["dataset"], self.configuration["table"]) def build_query(self): @@ -261,7 +261,7 @@ def build_query(self): """ # create a sub-config because full config contains secrets conf = {k: self.configuration[k] for k in ("predicates", "columns")} - conf["resolved_table"] = self.resolve_table() + conf["resolved_table"] = self._resolve_table() if not conf["columns"]: conf["columns"] = "*" @@ -271,10 +271,10 @@ def build_query(self): query = query + " " + conf["predicates"] return query.strip() - def url_safe_uuid(self): + def _url_safe_uuid(self): return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip(b'=').decode('ascii') - def generate_doc_id(self, row): + def _generate_doc_id(self, row): """Creates and returns a string suitable for use as a doc _id. If the user configured a doc_id_column to use for this, uses that from the row. If no id is configured, documents will be assigned a random uuid, similar to what @@ -289,9 +289,9 @@ def generate_doc_id(self, row): """ if self.configuration["doc_id_column"]: return row.get(self.configuration["doc_id_column"]) - return self.url_safe_uuid() + return self._url_safe_uuid() - def generate_doc_timestamp(self, row): + def _generate_doc_timestamp(self, row): """Creates and returns a timestamp string. If the user configured a timestamp_column from their table, uses that. If not, returns None, and current sync run start time in UTC will be assigned. @@ -308,8 +308,8 @@ def generate_doc_timestamp(self, row): def row2doc(self, row): doc = dict(row) # cast to dict from the google class - doc_id = self.generate_doc_id(row) - doc_timestamp = self.generate_doc_timestamp(row) + doc_id = self._generate_doc_id(row) + doc_timestamp = self._generate_doc_timestamp(row) if doc_timestamp is None: doc_timestamp = self._run_timestamp doc.update( From 7d81fdd1900f94e80a3302591c28302e6e6e1ac3 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 21 Jan 2026 12:02:55 -0500 Subject: [PATCH 10/19] make autoformat --- .../sources/google_bigquery/client.py | 7 +- .../sources/google_bigquery/datasource.py | 59 ++++++------ .../tests/sources/test_google_bigquery.py | 92 +++++++++++-------- 3 files changed, 88 insertions(+), 70 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py index 1a45803d6..10896fc11 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/client.py +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -5,12 +5,9 @@ # """Google Bigquery module which fetches rows from a Bigquery table.""" -from aiogoogle.auth.creds import ServiceAccountCreds from connectors_sdk.logger import logger from google.cloud import bigquery from google.oauth2 import service_account -import os -import tempfile class GoogleBigqueryClient: @@ -40,7 +37,9 @@ def client(self, project_id=None): bigquery.Client instance """ - credentials = service_account.Credentials.from_service_account_info(self.json_credentials) + credentials = service_account.Credentials.from_service_account_info( + self.json_credentials + ) if project_id is not None: return bigquery.Client(credentials=credentials, project=project_id) else: diff --git a/app/connectors_service/connectors/sources/google_bigquery/datasource.py b/app/connectors_service/connectors/sources/google_bigquery/datasource.py index 17f29652e..00b88e0df 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/datasource.py +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -6,25 +6,24 @@ import asyncio import base64 +import os +import uuid from concurrent.futures import ThreadPoolExecutor -from connectors_sdk.source import ( - BaseDataSource, - DataSourceConfiguration -) +from datetime import datetime, timezone +from functools import cached_property + +from connectors_sdk.source import BaseDataSource, DataSourceConfiguration + from connectors.sources.google_bigquery.client import GoogleBigqueryClient from connectors.sources.shared.database.generic_database import ( DEFAULT_FETCH_SIZE, - DEFAULT_RETRY_COUNT + DEFAULT_RETRY_COUNT, ) from connectors.sources.shared.google import ( load_service_account_json, validate_service_account_json, ) from connectors.utils import get_pem_format -from datetime import datetime, timezone -from functools import cached_property, partial -import os -import uuid RUNNING_FTEST = ( "RUNNING_FTEST" in os.environ @@ -32,6 +31,7 @@ executor = ThreadPoolExecutor(1) + class GoogleBigqueryDataSource(BaseDataSource): """Google Bigquery""" @@ -41,7 +41,6 @@ class GoogleBigqueryDataSource(BaseDataSource): advanced_rules_enabled = True dls_enabled = False - def __init__(self, configuration: DataSourceConfiguration): """Set up the connection to Google Bigquery. @@ -76,8 +75,8 @@ def _google_bigquery_client(self) -> GoogleBigqueryClient: ) if ( - json_credentials.get("private_key") - and "\n" not in json_credentials["private_key"] + json_credentials.get("private_key") + and "\n" not in json_credentials["private_key"] ): json_credentials["private_key"] = get_pem_format( key=json_credentials["private_key"].strip(), @@ -92,7 +91,6 @@ def _google_bigquery_client(self) -> GoogleBigqueryClient: return GoogleBigqueryClient(json_credentials=required_credentials) - @classmethod def get_default_configuration(cls) -> dict: """Get the default configuration for Google Bigquery. @@ -106,7 +104,7 @@ def get_default_configuration(cls) -> dict: "label": "Google Cloud service account JSON", "sensitive": True, "order": 1, - "type": "str" + "type": "str", }, "dataset": { "display": "text", @@ -157,7 +155,7 @@ def get_default_configuration(cls) -> dict: "default_value": None, "type": "str", "ui_restrictions": ["advanced"], - "tooltip": "Use the value of this column as the ES document _timestamp instead of using the sync start time." + "tooltip": "Use the value of this column as the ES document _timestamp instead of using the sync start time.", }, "predicates": { "display": "textarea", @@ -167,7 +165,7 @@ def get_default_configuration(cls) -> dict: "default_value": "", "type": "str", "ui_restrictions": ["advanced"], - "tooltip": "A SQL WHERE clause. May be required for some partitioned table configurations." + "tooltip": "A SQL WHERE clause. May be required for some partitioned table configurations.", }, "fetch_size": { "default_value": DEFAULT_FETCH_SIZE, @@ -197,14 +195,18 @@ async def validate_config(self): """ await super().validate_config() validate_service_account_json( - self.configuration["service_account_credentials"], - "Google Bigquery" + self.configuration["service_account_credentials"], "Google Bigquery" ) # if they set a project_id and have a wrong service account, a log message # will give them a fighting chance :) - if self.configuration["project_id"] and self._resolve_project() != self.configuration["project_id"]: - self._logger.info("A project_id is configured and does not match the project_id for the service_account_credentials block. If authorization fails, this could be why!") + if ( + self.configuration["project_id"] + and self._resolve_project() != self.configuration["project_id"] + ): + self._logger.info( + "A project_id is configured and does not match the project_id for the service_account_credentials block. If authorization fails, this could be why!" + ) self.project_id = self._resolve_project() async def ping(self): @@ -214,12 +216,11 @@ async def ping(self): try: sql = "SELECT 1=1" job = self._google_bigquery_client.client().query(sql) - return job.done() # signal we can indeed run queries + return job.done() # signal we can indeed run queries except Exception: self._logger.exception("Error while connecting to Google Bigquery.") raise - def _resolve_project(self): """ Resolves a project_id, as a string. This should be a Google Cloud project. @@ -240,7 +241,6 @@ def _resolve_project(self): ) return json_credentials["project_id"] - def _resolve_table(self): """Inspects the configuration and produces a fully qualified BigQuery table identifier. @@ -249,8 +249,11 @@ def _resolve_table(self): string: A BQ table in the form `project.dataset.table` """ - return "`%s.%s.%s`" % (self._resolve_project(), self.configuration["dataset"], self.configuration["table"]) - + return "`%s.%s.%s`" % ( + self._resolve_project(), + self.configuration["dataset"], + self.configuration["table"], + ) def build_query(self): """ @@ -272,7 +275,7 @@ def build_query(self): return query.strip() def _url_safe_uuid(self): - return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip(b'=').decode('ascii') + return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip(b"=").decode("ascii") def _generate_doc_id(self, row): """Creates and returns a string suitable for use as a doc _id. If the user @@ -307,7 +310,7 @@ def _generate_doc_timestamp(self, row): return None def row2doc(self, row): - doc = dict(row) # cast to dict from the google class + doc = dict(row) # cast to dict from the google class doc_id = self._generate_doc_id(row) doc_timestamp = self._generate_doc_timestamp(row) if doc_timestamp is None: @@ -352,7 +355,7 @@ def get_next_chunk(job_iter): for _ in range(fetch_size): try: row = next(job_iter) - chunk.append(self.row2doc(row)) # Row -> dict + chunk.append(self.row2doc(row)) # Row -> dict except StopIteration: break return chunk diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py index 3b1aaf639..ef4f4d5f2 100644 --- a/app/connectors_service/tests/sources/test_google_bigquery.py +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -5,19 +5,17 @@ # """Tests the Google Bigquery source.""" -import asyncio -from contextlib import asynccontextmanager import json -import os -import pytest -from unittest import mock +from contextlib import asynccontextmanager from unittest.mock import Mock, patch +import pytest +from connectors_sdk.source import ConfigurableFieldValueError, DataSourceConfiguration + from connectors.sources.google_bigquery import ( GoogleBigqueryClient, GoogleBigqueryDataSource, ) -from connectors_sdk.source import ConfigurableFieldValueError, DataSourceConfiguration from tests.sources.support import create_source SERVICE_ACCOUNT_CREDENTIALS = '{"project_id": "dummy123"}' @@ -26,32 +24,35 @@ test_bq_small = { "project_id": "bigquery-public-data", "dataset": "new_york_subway", - "table": "stations" + "table": "stations", } # a >10k table to use, 19640 to be exact test_bq_10k = { "project_id": "bigquery-public-data", "dataset": "new_york_subway", - "table": "trips" + "table": "trips", } + @asynccontextmanager async def create_bigquery_source( - service_account_credentials=SERVICE_ACCOUNT_CREDENTIALS, - dataset=test_bq_small["dataset"], - table=test_bq_small["table"] + service_account_credentials=SERVICE_ACCOUNT_CREDENTIALS, + dataset=test_bq_small["dataset"], + table=test_bq_small["table"], ): async with create_source( - GoogleBigqueryDataSource, - service_account_credentials=SERVICE_ACCOUNT_CREDENTIALS, - dataset=dataset, - table=table + GoogleBigqueryDataSource, + service_account_credentials=SERVICE_ACCOUNT_CREDENTIALS, + dataset=dataset, + table=table, ) as source: yield source + class Credentials: """Creates a mocked google oauth2 service_account Credentials instance.""" + def __init__(self, creds): self.json_credentials = creds @@ -59,6 +60,7 @@ def __init__(self, creds): class Job: """Creates a mocked google bigquery Job. Job is an iterator. Iteration returns a fixture result a few times and then stops.""" + def __init__(self, *args, **kw): self.doc_count = 5 @@ -77,9 +79,10 @@ def done(self): class Client: """This class creates a mocked bigquery client.""" + def __init__(self, *args, **kw): - self.credentials = kw.get('credentials') - self.project = kw.get('project') + self.credentials = kw.get("credentials") + self.project = kw.get("project") def query(self, sql): """Returns mocked Job.""" @@ -92,7 +95,6 @@ async def test_empty_configuration(): error_configs = [ # service account credentials must be present and not blank DataSourceConfiguration({"service_account_credentials": ""}) - ] for error_config in error_configs: @@ -100,6 +102,7 @@ async def test_empty_configuration(): with pytest.raises(ConfigurableFieldValueError): await bq.validate_config() + @pytest.mark.asyncio async def test_build_query(): """Tests query building for GoogleBigqueryDataSource.""" @@ -107,10 +110,12 @@ async def test_build_query(): # minimum defaults config = DataSourceConfiguration( - {"service_account_credentials": testcreds, - "dataset": "testdataset", - "table": "testtable", - }) + { + "service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + } + ) bq = GoogleBigqueryDataSource(configuration=config) query = bq.build_query() expected = """SELECT * FROM `testprojectid.testdataset.testtable`""" @@ -118,11 +123,13 @@ async def test_build_query(): # custom project config = DataSourceConfiguration( - {"service_account_credentials": testcreds, - "dataset": "testdataset", - "table": "testtable", - "project_id": "different", - }) + { + "service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + "project_id": "different", + } + ) bq = GoogleBigqueryDataSource(configuration=config) query = bq.build_query() expected = """SELECT * FROM `different.testdataset.testtable`""" @@ -130,23 +137,29 @@ async def test_build_query(): # custom columns config = DataSourceConfiguration( - {"service_account_credentials": testcreds, - "dataset": "testdataset", - "table": "testtable", - "columns": "account_id, timestamp" - }) + { + "service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + "columns": "account_id, timestamp", + } + ) bq = GoogleBigqueryDataSource(configuration=config) query = bq.build_query() - expected = """SELECT account_id, timestamp FROM `testprojectid.testdataset.testtable`""" + expected = ( + """SELECT account_id, timestamp FROM `testprojectid.testdataset.testtable`""" + ) assert query == expected # custom predicates config = DataSourceConfiguration( - {"service_account_credentials": testcreds, - "dataset": "testdataset", - "table": "testtable", - "predicates": "WHERE foo=1" - }) + { + "service_account_credentials": testcreds, + "dataset": "testdataset", + "table": "testtable", + "predicates": "WHERE foo=1", + } + ) bq = GoogleBigqueryDataSource(configuration=config) query = bq.build_query() expected = """SELECT * FROM `testprojectid.testdataset.testtable` WHERE foo=1""" @@ -159,6 +172,7 @@ async def test_ping(): with patch.object(GoogleBigqueryClient, "client", return_value=Client()): await source.ping() + @pytest.mark.asyncio @patch("connectors.utils.time_to_sleep_between_retries", Mock(return_value=0)) async def test_ping_negative(): @@ -168,6 +182,7 @@ async def test_ping_negative(): with patch.object(Client, "query", side_effect=Exception()): await source.ping() + @pytest.mark.asyncio async def test_get_docs(): async with create_bigquery_source() as source: @@ -177,6 +192,7 @@ async def test_get_docs(): docs.append(doc) assert len(docs) == 5 + # To test against live bigquery, you can uncomment this and: # - have an active google cloud service account with BQ execute permission # - set up those creds and provide the location in GOOGLE_APPLICATION_CREDENTIALS From a66392a6bcfe12dcb87b4172d97ce6bd187b2a6f Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 28 Jan 2026 10:23:39 -0500 Subject: [PATCH 11/19] add ftest support to client --- .../sources/google_bigquery/client.py | 43 ++++++++++++++++--- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py index 10896fc11..e4cf3425f 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/client.py +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -6,8 +6,15 @@ """Google Bigquery module which fetches rows from a Bigquery table.""" from connectors_sdk.logger import logger +from google.api_core.client_options import ClientOptions +from google.auth.credentials import AnonymousCredentials from google.cloud import bigquery from google.oauth2 import service_account +import os + +RUNNING_FTEST = ( + "RUNNING_FTEST" in os.environ +) # Flag to check if a connector is run for ftest or not. class GoogleBigqueryClient: @@ -37,10 +44,36 @@ def client(self, project_id=None): bigquery.Client instance """ - credentials = service_account.Credentials.from_service_account_info( - self.json_credentials - ) + + if RUNNING_FTEST: + self._logger.debug("GoogleBigqueryClient: ftest detected, using AnonymousCredentials.") + credentials = AnonymousCredentials() + # even AnonymousCredentials will pick up a project_id it finds in the user's + # ADC config if they have one, so we MUST override that here too. + project_id = "test" + else: + credentials = service_account.Credentials.from_service_account_info( + self.json_credentials + ) + + # Extend here if configurable BQ client options are required in future, + # potentially useful to implement private universes for example if needed. + client_options = None + + # When running local ftest, connect to local container for api. + if RUNNING_FTEST: + client_options = ClientOptions(api_endpoint="http://0.0.0.0:9050") + + if project_id is not None: - return bigquery.Client(credentials=credentials, project=project_id) + self._logger.debug(f"GoogleBigqueryClient setting project_id: {project_id}.") + if client_options is not None: + return bigquery.Client(credentials=credentials, project=project_id, client_options=client_options) + else: + return bigquery.Client(credentials=credentials, project=project_id) else: - return bigquery.Client(credentials=credentials) + self._logger.debug(f"GoogleBigqueryClient setting default project_id.") + if client_options is not None: + return bigquery.Client(credentials=credentials, client_options=client_options) + else: + return bigquery.Client(credentials=credentials) From 2aebee0834e19aabc864c2d7882455adc9b1cb05 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 28 Jan 2026 10:34:06 -0500 Subject: [PATCH 12/19] ftest fixture for google_bigquery --- .../fixtures/google_bigquery/config.yml | 6 ++ .../fixtures/google_bigquery/connector.json | 101 ++++++++++++++++++ .../google_bigquery/docker-compose.yml | 40 +++++++ .../fixtures/google_bigquery/fixture.py | 74 +++++++++++++ 4 files changed, 221 insertions(+) create mode 100644 app/connectors_service/tests/sources/fixtures/google_bigquery/config.yml create mode 100644 app/connectors_service/tests/sources/fixtures/google_bigquery/connector.json create mode 100644 app/connectors_service/tests/sources/fixtures/google_bigquery/docker-compose.yml create mode 100644 app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py diff --git a/app/connectors_service/tests/sources/fixtures/google_bigquery/config.yml b/app/connectors_service/tests/sources/fixtures/google_bigquery/config.yml new file mode 100644 index 000000000..168c0c72f --- /dev/null +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/config.yml @@ -0,0 +1,6 @@ +service.idling: 1 + +connectors: + - + connector_id: 'google_bigquery' + service_type: 'google_bigquery' diff --git a/app/connectors_service/tests/sources/fixtures/google_bigquery/connector.json b/app/connectors_service/tests/sources/fixtures/google_bigquery/connector.json new file mode 100644 index 000000000..805043109 --- /dev/null +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/connector.json @@ -0,0 +1,101 @@ +{ + "configuration": { + "service_account_credentials": { + "depends_on": [], + "display": "textarea", + "tooltip": null, + "default_value": null, + "label": "Google Cloud service account JSON", + "sensitive": true, + "type": "str", + "required": true, + "options": [], + "validations": [], + "value": "{\"type\": \"service_account\", \"project_id\": \"dummy_project_id\", \"private_key_id\": \"abc\", \"private_key\": \"\\n-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDY3E8o1NEFcjMM\\nHW/5ZfFJw29/8NEqpViNjQIx95Xx5KDtJ+nWn9+OW0uqsSqKlKGhAdAo+Q6bjx2c\\nuXVsXTu7XrZUY5Kltvj94DvUa1wjNXs606r/RxWTJ58bfdC+gLLxBfGnB6CwK0YQ\\nxnfpjNbkUfVVzO0MQD7UP0Hl5ZcY0Puvxd/yHuONQn/rIAieTHH1pqgW+zrH/y3c\\n59IGThC9PPtugI9ea8RSnVj3PWz1bX2UkCDpy9IRh9LzJLaYYX9RUd7++dULUlat\\nAaXBh1U6emUDzhrIsgApjDVtimOPbmQWmX1S60mqQikRpVYZ8u+NDD+LNw+/Eovn\\nxCj2Y3z1AgMBAAECggEAWDBzoqO1IvVXjBA2lqId10T6hXmN3j1ifyH+aAqK+FVl\\nGjyWjDj0xWQcJ9ync7bQ6fSeTeNGzP0M6kzDU1+w6FgyZqwdmXWI2VmEizRjwk+/\\n/uLQUcL7I55Dxn7KUoZs/rZPmQDxmGLoue60Gg6z3yLzVcKiDc7cnhzhdBgDc8vd\\nQorNAlqGPRnm3EqKQ6VQp6fyQmCAxrr45kspRXNLddat3AMsuqImDkqGKBmF3Q1y\\nxWGe81LphUiRqvqbyUlh6cdSZ8pLBpc9m0c3qWPKs9paqBIvgUPlvOZMqec6x4S6\\nChbdkkTRLnbsRr0Yg/nDeEPlkhRBhasXpxpMUBgPywKBgQDs2axNkFjbU94uXvd5\\nznUhDVxPFBuxyUHtsJNqW4p/ujLNimGet5E/YthCnQeC2P3Ym7c3fiz68amM6hiA\\nOnW7HYPZ+jKFnefpAtjyOOs46AkftEg07T9XjwWNPt8+8l0DYawPoJgbM5iE0L2O\\nx8TU1Vs4mXc+ql9F90GzI0x3VwKBgQDqZOOqWw3hTnNT07Ixqnmd3dugV9S7eW6o\\nU9OoUgJB4rYTpG+yFqNqbRT8bkx37iKBMEReppqonOqGm4wtuRR6LSLlgcIU9Iwx\\nyfH12UWqVmFSHsgZFqM/cK3wGev38h1WBIOx3/djKn7BdlKVh8kWyx6uC8bmV+E6\\nOoK0vJD6kwKBgHAySOnROBZlqzkiKW8c+uU2VATtzJSydrWm0J4wUPJifNBa/hVW\\ndcqmAzXC9xznt5AVa3wxHBOfyKaE+ig8CSsjNyNZ3vbmr0X04FoV1m91k2TeXNod\\njMTobkPThaNm4eLJMN2SQJuaHGTGERWC0l3T18t+/zrDMDCPiSLX1NAvAoGBAN1T\\nVLJYdjvIMxf1bm59VYcepbK7HLHFkRq6xMJMZbtG0ryraZjUzYvB4q4VjHk2UDiC\\nlhx13tXWDZH7MJtABzjyg+AI7XWSEQs2cBXACos0M4Myc6lU+eL+iA+OuoUOhmrh\\nqmT8YYGu76/IBWUSqWuvcpHPpwl7871i4Ga/I3qnAoGBANNkKAcMoeAbJQK7a/Rn\\nwPEJB+dPgNDIaboAsh1nZhVhN5cvdvCWuEYgOGCPQLYQF0zmTLcM+sVxOYgfy8mV\\nfbNgPgsP5xmu6dw2COBKdtozw0HrWSRjACd1N4yGu75+wPCcX/gQarcjRcXXZeEa\\nNtBLSfcqPULqD+h7br9lEJio\\n-----END PRIVATE KEY-----\\n\", \"client_email\": \"123-abc@developer.gserviceaccount.com\", \"client_id\": \"123-abc.apps.googleusercontent.com\", \"auth_uri\": \"https://accounts.google.com/o/oauth2/auth\", \"token_uri\": \"http://localhost:4444/token\"}", + "order": 1, + "ui_restrictions": [] + }, + "dataset": { + "display": "text", + "label": "Dataset the table is in.", + "order": 2, + "type": "str", + "value": "testdata" + }, + "table": { + "display": "text", + "label": "Table to sync.", + "order": 3, + "type": "str", + "value": "testtable" + }, + "project_id": { + "display": "text", + "label": "Google Cloud project.", + "order": 4, + "type": "str", + "required": false, + "default_value": "", + "tooltip": "Defaults to the service account project.", + "value": "test" + }, + "columns": { + "display": "textarea", + "label": "Columns to fetch. Defaults to * if none are set.", + "order": 6, + "required": false, + "default_value": "*", + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "Comma-separated, as in a SQL SELECT." + }, + "doc_id_column": { + "display": "text", + "label": "Document _id column", + "order": 7, + "required": false, + "default_value": null, + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "Use the value of this column as the ES document _id instead of generating a UUID." + }, + "timestamp_column": { + "display": "text", + "label": "Timestamp column", + "order": 8, + "required": false, + "default_value": null, + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "Use the value of this column as the ES document _timestamp instead of using the sync start time." + }, + "predicates": { + "display": "textarea", + "label": "Predicates for the query.", + "order": 9, + "required": false, + "default_value": "", + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "A SQL WHERE clause. May be required for some partitioned table configurations." + }, + "fetch_size": { + "default_value": 50, + "display": "numeric", + "label": "Rows fetched per request", + "order": 10, + "required": false, + "type": "int", + "ui_restrictions": ["advanced"] + }, + "retry_count": { + "default_value": 3, + "display": "numeric", + "label": "Retries per request", + "order": 11, + "required": false, + "type": "int", + "ui_restrictions": ["advanced"] + } + } +} diff --git a/app/connectors_service/tests/sources/fixtures/google_bigquery/docker-compose.yml b/app/connectors_service/tests/sources/fixtures/google_bigquery/docker-compose.yml new file mode 100644 index 000000000..45e1555ed --- /dev/null +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/docker-compose.yml @@ -0,0 +1,40 @@ +version: '3.9' + +services: + elasticsearch: + image: ${ELASTICSEARCH_DRA_DOCKER_IMAGE} + container_name: elasticsearch + environment: + - cluster.name=docker-cluster + - bootstrap.memory_lock=true + - ES_JAVA_OPTS=-Xms2g -Xmx2g + - ELASTIC_PASSWORD=changeme + - xpack.security.enabled=true + - xpack.security.authc.api_key.enabled=true + - discovery.type=single-node + - action.destructive_requires_name=false + ulimits: + memlock: + soft: -1 + hard: -1 + volumes: + - esdata:/usr/share/elasticsearch/data + ports: + - 9200:9200 + networks: + - esnet + + bigquery-emulator: + image: ghcr.io/goccy/bigquery-emulator:latest + container_name: bigquery-emulator + platform: linux/amd64 + ports: + - "9050:9050" + command: --project=test --dataset=testdata --port=9050 + +networks: + esnet: + +volumes: + esdata: + driver: local diff --git a/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py new file mode 100644 index 000000000..5c36d8600 --- /dev/null +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py @@ -0,0 +1,74 @@ +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License 2.0; +# you may not use this file except in compliance with the Elastic License 2.0. +# +# ruff: noqa: T201 + +"""Google Bigquery module that generates fixture data on the local mocked Bigquery container.""" + +from google.api_core.client_options import ClientOptions +from google.auth.credentials import AnonymousCredentials +from google.cloud import bigquery +from google.cloud.bigquery import QueryJobConfig +from itertools import islice +import os + +client_options = ClientOptions(api_endpoint="http://0.0.0.0:9050") + +PROJECT = "test" +DATASET = "testdata" +TABLE = "testtable" + +DATA_SIZE = os.environ.get("DATA_SIZE", "medium").lower() + +match DATA_SIZE: + case "small": + RECORD_COUNT = 500 + case "medium": + RECORD_COUNT = 10000 + case "large": + RECORD_COUNT = 25000 + + +def partition_all(iterable, chunk_size): + iterator = iter(iterable) + while chunk := tuple(islice(iterator, chunk_size)): + yield chunk + +def insert_rows(client, table): + print(f"Loading BQ fixture data into {table}", end="") + rows = [] + for n in range(RECORD_COUNT): + rows.append({"testcolumn": f"row {n}"}) + for chunk in partition_all(rows, 50): + print(".", end="") + client.insert_rows_json(table, chunk) + print("done.") + +async def load(): + """Loads DATA_SIZE records into the Bigquery fixture container.""" + client = bigquery.Client( + PROJECT, + client_options=client_options, + credentials=AnonymousCredentials(), + ) + + schema = [ + bigquery.SchemaField("testcolumn", "STRING", mode="NULLABLE") + ] + + table_ref = client.dataset(DATASET).table(TABLE) + table = bigquery.Table(table_ref, schema=schema) + try: + table = client.create_table(table) + except Exception as e: + print(f"Error creating fixture table {table.full_table_id}: {e}") + + # fetch full created table instance back from BQ using the table_ref + table = client.get_table(table_ref) + + errors = insert_rows(client, table) + if errors is not None and errors != []: + print(f"Errors inserting fixture data: {errors}") + From fc545cd2a1e747595d86a67cf3cdee3d91bfffaa Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 28 Jan 2026 11:25:55 -0500 Subject: [PATCH 13/19] make autoformat for ftest fixture work --- .../sources/google_bigquery/client.py | 24 +++++++++++++------ .../fixtures/google_bigquery/fixture.py | 13 +++++----- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py index e4cf3425f..cfc8b173c 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/client.py +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -5,12 +5,13 @@ # """Google Bigquery module which fetches rows from a Bigquery table.""" +import os + from connectors_sdk.logger import logger from google.api_core.client_options import ClientOptions from google.auth.credentials import AnonymousCredentials from google.cloud import bigquery from google.oauth2 import service_account -import os RUNNING_FTEST = ( "RUNNING_FTEST" in os.environ @@ -46,7 +47,9 @@ def client(self, project_id=None): """ if RUNNING_FTEST: - self._logger.debug("GoogleBigqueryClient: ftest detected, using AnonymousCredentials.") + self._logger.debug( + "GoogleBigqueryClient: ftest detected, using AnonymousCredentials." + ) credentials = AnonymousCredentials() # even AnonymousCredentials will pick up a project_id it finds in the user's # ADC config if they have one, so we MUST override that here too. @@ -64,16 +67,23 @@ def client(self, project_id=None): if RUNNING_FTEST: client_options = ClientOptions(api_endpoint="http://0.0.0.0:9050") - if project_id is not None: - self._logger.debug(f"GoogleBigqueryClient setting project_id: {project_id}.") + self._logger.debug( + f"GoogleBigqueryClient setting project_id: {project_id}." + ) if client_options is not None: - return bigquery.Client(credentials=credentials, project=project_id, client_options=client_options) + return bigquery.Client( + credentials=credentials, + project=project_id, + client_options=client_options, + ) else: return bigquery.Client(credentials=credentials, project=project_id) else: - self._logger.debug(f"GoogleBigqueryClient setting default project_id.") + self._logger.debug("GoogleBigqueryClient setting default project_id.") if client_options is not None: - return bigquery.Client(credentials=credentials, client_options=client_options) + return bigquery.Client( + credentials=credentials, client_options=client_options + ) else: return bigquery.Client(credentials=credentials) diff --git a/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py index 5c36d8600..062118f3f 100644 --- a/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py @@ -7,12 +7,12 @@ """Google Bigquery module that generates fixture data on the local mocked Bigquery container.""" +import os +from itertools import islice + from google.api_core.client_options import ClientOptions from google.auth.credentials import AnonymousCredentials from google.cloud import bigquery -from google.cloud.bigquery import QueryJobConfig -from itertools import islice -import os client_options = ClientOptions(api_endpoint="http://0.0.0.0:9050") @@ -36,6 +36,7 @@ def partition_all(iterable, chunk_size): while chunk := tuple(islice(iterator, chunk_size)): yield chunk + def insert_rows(client, table): print(f"Loading BQ fixture data into {table}", end="") rows = [] @@ -46,6 +47,7 @@ def insert_rows(client, table): client.insert_rows_json(table, chunk) print("done.") + async def load(): """Loads DATA_SIZE records into the Bigquery fixture container.""" client = bigquery.Client( @@ -54,9 +56,7 @@ async def load(): credentials=AnonymousCredentials(), ) - schema = [ - bigquery.SchemaField("testcolumn", "STRING", mode="NULLABLE") - ] + schema = [bigquery.SchemaField("testcolumn", "STRING", mode="NULLABLE")] table_ref = client.dataset(DATASET).table(TABLE) table = bigquery.Table(table_ref, schema=schema) @@ -71,4 +71,3 @@ async def load(): errors = insert_rows(client, table) if errors is not None and errors != []: print(f"Errors inserting fixture data: {errors}") - From 4ecd51e1da8fd52e672e2835e529bac0fec2ee50 Mon Sep 17 00:00:00 2001 From: Dan Dillinger Date: Wed, 28 Jan 2026 13:34:58 -0500 Subject: [PATCH 14/19] add config option for a full custom sql query --- .../sources/google_bigquery/datasource.py | 94 ++++++++++++++----- 1 file changed, 68 insertions(+), 26 deletions(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/datasource.py b/app/connectors_service/connectors/sources/google_bigquery/datasource.py index 00b88e0df..09d9fcaec 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/datasource.py +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -12,7 +12,11 @@ from datetime import datetime, timezone from functools import cached_property -from connectors_sdk.source import BaseDataSource, DataSourceConfiguration +from connectors_sdk.source import ( + BaseDataSource, + ConfigurableFieldValueError, + DataSourceConfiguration, +) from connectors.sources.google_bigquery.client import GoogleBigqueryClient from connectors.sources.shared.database.generic_database import ( @@ -106,22 +110,42 @@ def get_default_configuration(cls) -> dict: "order": 1, "type": "str", }, + "query_type": { + "label": "Query Method", + "order": 2, + "type": "str", + "display": "dropdown", + "options": [ + {"label": "Table", "value": "table"}, + {"label": "Custom SQL Query", "value": "custom_query"}, + ], + "value": "table", + }, + "custom_query": { + "label": "Custom SQL Query", + "order": 3, + "type": "str", + "display": "textarea", + "depends_on": [{"field": "query_type", "value": "custom_query"}], + }, "dataset": { "display": "text", "label": "Dataset the table is in.", - "order": 2, + "order": 4, "type": "str", + "depends_on": [{"field": "query_type", "value": "table"}], }, "table": { "display": "text", "label": "Table to sync.", - "order": 3, + "order": 5, "type": "str", + "depends_on": [{"field": "query_type", "value": "table"}], }, "project_id": { "display": "text", "label": "Google Cloud project.", - "order": 4, + "order": 6, "type": "str", "required": False, "default_value": "", @@ -130,17 +154,29 @@ def get_default_configuration(cls) -> dict: "columns": { "display": "textarea", "label": "Columns to fetch. Defaults to * if none are set.", - "order": 6, + "order": 7, "required": False, "default_value": "*", "type": "str", "ui_restrictions": ["advanced"], "tooltip": "Comma-separated, as in a SQL SELECT.", + "depends_on": [{"field": "query_type", "value": "table"}], + }, + "predicates": { + "display": "textarea", + "label": "Predicates for the query.", + "order": 8, + "required": False, + "default_value": "", + "type": "str", + "ui_restrictions": ["advanced"], + "tooltip": "A SQL WHERE clause. May be required for some partitioned table configurations.", + "depends_on": [{"field": "query_type", "value": "table"}], }, "doc_id_column": { "display": "text", "label": "Document _id column", - "order": 7, + "order": 9, "required": False, "default_value": None, "type": "str", @@ -150,28 +186,18 @@ def get_default_configuration(cls) -> dict: "timestamp_column": { "display": "text", "label": "Timestamp column", - "order": 8, + "order": 10, "required": False, "default_value": None, "type": "str", "ui_restrictions": ["advanced"], "tooltip": "Use the value of this column as the ES document _timestamp instead of using the sync start time.", }, - "predicates": { - "display": "textarea", - "label": "Predicates for the query.", - "order": 9, - "required": False, - "default_value": "", - "type": "str", - "ui_restrictions": ["advanced"], - "tooltip": "A SQL WHERE clause. May be required for some partitioned table configurations.", - }, "fetch_size": { "default_value": DEFAULT_FETCH_SIZE, "display": "numeric", "label": "Rows fetched per request", - "order": 10, + "order": 11, "required": False, "type": "int", "ui_restrictions": ["advanced"], @@ -180,7 +206,7 @@ def get_default_configuration(cls) -> dict: "default_value": DEFAULT_RETRY_COUNT, "display": "numeric", "label": "Retries per request", - "order": 11, + "order": 12, "required": False, "type": "int", "ui_restrictions": ["advanced"], @@ -207,7 +233,16 @@ async def validate_config(self): self._logger.info( "A project_id is configured and does not match the project_id for the service_account_credentials block. If authorization fails, this could be why!" ) - self.project_id = self._resolve_project() + self.project_id = self._resolve_project() + + # validate that the UI wasn't bypassed, and an invalid config of both query + # types has been configured on the connector. + if ( + self.configuration["custom_query"] is not None + and self.configuration["table"] is not None + ): + msg = "Both table and custom_query cannot be set simultaneously." + raise ConfigurableFieldValueError(msg) async def ping(self): """Verify the connection with Google Bigquery""" @@ -262,6 +297,13 @@ def build_query(self): Returns: string: The query. """ + # if the user configured their own SQL use that + if ( + self.configuration["custom_query"] is not None + and self.configuration["custom_query"] != "" + ): + return self.configuration["custom_query"] + # create a sub-config because full config contains secrets conf = {k: self.configuration[k] for k in ("predicates", "columns")} conf["resolved_table"] = self._resolve_table() @@ -315,12 +357,12 @@ def row2doc(self, row): doc_timestamp = self._generate_doc_timestamp(row) if doc_timestamp is None: doc_timestamp = self._run_timestamp - doc.update( - { - "_id": doc_id, - "_timestamp": doc_timestamp, - } - ) + doc.update( + { + "_id": doc_id, + "_timestamp": doc_timestamp, + } + ) return doc async def get_docs(self, filtering=None): From f483891d5c744af777bcde514ec31f2733fc8258 Mon Sep 17 00:00:00 2001 From: David Leatherman Date: Fri, 20 Feb 2026 10:41:14 -0500 Subject: [PATCH 15/19] suppress 3.10 warning we'll need to address it, but for now let's see if this fixes the test run --- .../tests/sources/test_google_bigquery.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py index ef4f4d5f2..01dd6a1f6 100644 --- a/app/connectors_service/tests/sources/test_google_bigquery.py +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -5,6 +5,13 @@ # """Tests the Google Bigquery source.""" +# Python 3.10 id deprecated in google.api_core and will be unsupported +# on 2026-10-04, causing test failures on that version. Suppress +# these for now. +import warnings + +warnings.simplefilter(action="ignore", category=FutureWarning) + import json from contextlib import asynccontextmanager from unittest.mock import Mock, patch From bef6a1c4b9c720a98a8fe42b9404b408065bf07f Mon Sep 17 00:00:00 2001 From: David Leatherman Date: Tue, 24 Feb 2026 17:18:30 -0500 Subject: [PATCH 16/19] better way to suppress the 3.10 warning This way the imports can be sorted --- app/connectors_service/pyproject.toml | 2 ++ .../tests/sources/test_google_bigquery.py | 7 ------- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/app/connectors_service/pyproject.toml b/app/connectors_service/pyproject.toml index 03a602008..d8b18b9a5 100644 --- a/app/connectors_service/pyproject.toml +++ b/app/connectors_service/pyproject.toml @@ -131,6 +131,8 @@ addopts = [ ] filterwarnings = [ "error", + # google.api_core warns about Python 3.10 EOL; we still support 3.10 + "ignore::FutureWarning:google.api_core._python_version_support", # botocore has this warning that is reported by them to be irrelevant "ignore:.*urllib3.contrib.pyopenssl.*:DeprecationWarning:botocore.*", # latest main of aioresponses does not have this problem, but current package uses deprecated pkg_resources API diff --git a/app/connectors_service/tests/sources/test_google_bigquery.py b/app/connectors_service/tests/sources/test_google_bigquery.py index 01dd6a1f6..ef4f4d5f2 100644 --- a/app/connectors_service/tests/sources/test_google_bigquery.py +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -5,13 +5,6 @@ # """Tests the Google Bigquery source.""" -# Python 3.10 id deprecated in google.api_core and will be unsupported -# on 2026-10-04, causing test failures on that version. Suppress -# these for now. -import warnings - -warnings.simplefilter(action="ignore", category=FutureWarning) - import json from contextlib import asynccontextmanager from unittest.mock import Mock, patch From a3d0036abb468c42be666426b9d9efaa2d6db713 Mon Sep 17 00:00:00 2001 From: David Leatherman Date: Tue, 24 Feb 2026 17:19:20 -0500 Subject: [PATCH 17/19] change bigquery import to work around pyright bug google.cloud ns is shared across multiple packages and appears to confuse pyright (this has been reported several times in the past and appears to be a continual issue). --- .../connectors/sources/google_bigquery/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/connectors_service/connectors/sources/google_bigquery/client.py b/app/connectors_service/connectors/sources/google_bigquery/client.py index cfc8b173c..27111b856 100644 --- a/app/connectors_service/connectors/sources/google_bigquery/client.py +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -7,10 +7,10 @@ import os +import google.cloud.bigquery as bigquery from connectors_sdk.logger import logger from google.api_core.client_options import ClientOptions from google.auth.credentials import AnonymousCredentials -from google.cloud import bigquery from google.oauth2 import service_account RUNNING_FTEST = ( From 166b61b00b5bac3cc4a93e3767ebb44664f7cae0 Mon Sep 17 00:00:00 2001 From: David Leatherman Date: Tue, 24 Feb 2026 17:43:17 -0500 Subject: [PATCH 18/19] add test to Buildkite pipeline --- .buildkite/pipeline.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 36a2dc82e..6967f8893 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -360,6 +360,23 @@ steps: - DATA_SIZE=small - CONNECTOR=github + - path: + - "app/connectors_service/connectors/sources/google_bigquery/**" + - "app/connectors_service/tests/sources/fixtures/google_bigquery/**" + - "app/connectors_service/tests/sources/fixtures/fixture.py" + - "${DOCKERFILE_PATH}" + - "app/connectors_service/tests/Dockerfile.ftest" + - "app/connectors_service/pyproject.toml" + config: + label: "🔨 Google BigQuery" + <<: *test-agents + <<: *retries + <<: *ftest_defaults + env: + - PYTHON_VERSION=3.11 + - DATA_SIZE=small + - CONNECTOR=google_bigquery + - path: - "app/connectors_service/connectors/sources/google_drive/**" - "app/connectors_service/tests/sources/fixtures/google_drive/**" From 79267ced6aa34a83181cdf5cbe083ef44333dc10 Mon Sep 17 00:00:00 2001 From: David Leatherman Date: Wed, 25 Feb 2026 11:33:53 -0500 Subject: [PATCH 19/19] add get_num_docs to bigquery/fixture.py --- .../tests/sources/fixtures/google_bigquery/fixture.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py index 062118f3f..d0b92e74f 100644 --- a/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py @@ -31,6 +31,10 @@ RECORD_COUNT = 25000 +def get_num_docs(): + print(RECORD_COUNT) + + def partition_all(iterable, chunk_size): iterator = iter(iterable) while chunk := tuple(islice(iterator, chunk_size)):