diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 5de542c86..8bbbfe4d7 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -403,6 +403,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/**" diff --git a/app/connectors_service/connectors/config.py b/app/connectors_service/connectors/config.py index eb27f75c4..e151499a5 100644 --- a/app/connectors_service/connectors/config.py +++ b/app/connectors_service/connectors/config.py @@ -120,6 +120,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", 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..27111b856 --- /dev/null +++ b/app/connectors_service/connectors/sources/google_bigquery/client.py @@ -0,0 +1,89 @@ +# +# 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.""" + +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.oauth2 import service_account + +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 set_logger(self, logger_): + 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 + + """ + + 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: + 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: + self._logger.debug("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) 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..09d9fcaec --- /dev/null +++ b/app/connectors_service/connectors/sources/google_bigquery/datasource.py @@ -0,0 +1,421 @@ +# +# 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 +import base64 +import os +import uuid +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from functools import cached_property + +from connectors_sdk.source import ( + BaseDataSource, + ConfigurableFieldValueError, + 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 + +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", + }, + "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": 4, + "type": "str", + "depends_on": [{"field": "query_type", "value": "table"}], + }, + "table": { + "display": "text", + "label": "Table to sync.", + "order": 5, + "type": "str", + "depends_on": [{"field": "query_type", "value": "table"}], + }, + "project_id": { + "display": "text", + "label": "Google Cloud project.", + "order": 6, + "type": "str", + "required": False, + "default_value": "", + "tooltip": "Defaults to the service account project.", + }, + "columns": { + "display": "textarea", + "label": "Columns to fetch. Defaults to * if none are set.", + "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": 9, + "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": 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.", + }, + "fetch_size": { + "default_value": DEFAULT_FETCH_SIZE, + "display": "numeric", + "label": "Rows fetched per request", + "order": 11, + "required": False, + "type": "int", + "ui_restrictions": ["advanced"], + }, + "retry_count": { + "default_value": DEFAULT_RETRY_COUNT, + "display": "numeric", + "label": "Retries per request", + "order": 12, + "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._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() + + # 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""" + if RUNNING_FTEST: + return + try: + 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 + + 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. + """ + # 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() + 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() + + 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) # 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 + doc.update( + { + "_id": doc_id, + "_timestamp": doc_timestamp, + } + ) + return doc + + 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.") + # 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 + # 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(self.row2doc(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/pyproject.toml b/app/connectors_service/pyproject.toml index 657c7d06e..d8b18b9a5 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", @@ -130,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/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..d0b92e74f --- /dev/null +++ b/app/connectors_service/tests/sources/fixtures/google_bigquery/fixture.py @@ -0,0 +1,77 @@ +# +# 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.""" + +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 + +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 get_num_docs(): + print(RECORD_COUNT) + + +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}") 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..ef4f4d5f2 --- /dev/null +++ b/app/connectors_service/tests/sources/test_google_bigquery.py @@ -0,0 +1,220 @@ +# +# 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 json +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 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 + + +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.""" + 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 + + +@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 + + +# 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_from_live") +# 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