Repository: GoogleCloudPlatform/professional-services-data-validator

Issue:
Issues with custom bigquery api endpoint
In issue https://github.com/GoogleCloudPlatform/professional-services-data-validator/issues/1337 we added support for customers to define a custom endpoint. This was tested by the requesting and confirmed as working, however, as they have continued their testing they have run into further issues.

Example (sanitised) command:
```
$ data-validation -v --log-level DEBUG validate column -sc bq_conn -tc bq_conn \
-tbls OWN.TAB=OWN.TAB --filters 'TS_COL > TIMESTAMP(CURRENT_DATE())'
...
11/26/2024 01:23:36 PM-DEBUG: Starting new HTTPS connection (1): bigquery-priv.p.googleapis.com:443
11/26/2024 01:23:36 PM-DEBUG: https://bigquery-priv.p.googleapis.com:443 "GET /bigquery/v2/projects/proj/datasets/OWN/tables/TAB?prettyPrint=false HTTP/11" 200 None
11/26/2024 01:23:36 PM-DEBUG: https://bigquery-priv.p.googleapis.com:443 "GET /bigquery/v2/projects/proj/datasets/OWN/tables/TAB?prettyPrint=false HTTP/11" 200 None
11/26/2024 01:23:36 PM-INFO: {'data_client': <ibis.backends.bigquery.Backend object at 0x7f4a5d49da90>, 'schema_name': 'OWN', 'table_name': 'TAB', 'source_query': None}
11/26/2024 01:23:36 PM-INFO: -- ** Source Query ** --
11/26/2024 01:23:36 PM-INFO: SELECT count(1) AS `count`
FROM `proj.OWN.TAB` t0
WHERE TS_COL > TIMESTAMP(CURRENT_DATE())
11/26/2024 01:23:36 PM-DEBUG: Converted retries value: 3 -> Retry(total=3, connect=None, read=None, redirect=None, status=None)

... same as above but for target connection ...

11/26/2024 01:24:03 PM-DEBUG: Retrying due to 503 failed to connect to all addresses; last error: UNKNOWN: ipv4:172.x.y.z:443: Failed to connect to remote host: Timeout occurred: FD Shutdown, sleeping 0.0s ...
11/26/2024 01:24:03 PM-DEBUG: Retrying due to 503 failed to connect to all addresses; last error: UNKNOWN: ipv4:142.x.y.z:443: Failed to connect to remote host: Timeout occurred: FD Shutdown, sleeping 0.1s ...
11/26/2024 01:24:03 PM-DEBUG: Retrying due to 503 failed to connect to all addresses; last error: UNKNOWN: ipv4:172.x.y.z:443: Failed to connect to remote host: Timeout occurred: FD Shutdown, sleeping 0.0s …
```

We can see some successful interactions with the correct private endpoint:
```
11/26/2024 01:23:36 PM-DEBUG: https://bigquery-priv.p.googleapis.com:443 "GET /bigquery/v2/projects/proj/datasets/OWN/tables/TAB?prettyPrint=false HTTP/11" 200 None
```

But start trying to access blocked IPs when executing the query.

Hints from issue comments:
@nj1973: Interestingly this query works:
```
$ data-validation -v --log-level DEBUG query --conn bq_conn --query 'SELECT count(1) AS `count` FROM `proj.O_TEST.T1` t0'
...
[(1,)]
```

As does this:
```
$ data-validation -v --log-level DEBUG query --conn bq_conn --query 'SELECT count(1) AS `count` FROM `proj.OWN.TAB` where TS_COL > TIMESTAMP(CURRENT_DATE()) '
...
[(349893,)]
```

The second query above it the same query as DVT generated in the failing command. The only difference appears to be that one is from `data-validation query` and the other from `data-validation validate column`.
@nj1973: We are missing an override for the BigQuery Storage API endpoint.

`data-validation query` is not using the storage API, when I use `data-validation validate column` we pass through some Ibis `to_arrow` code which sends us down a different path and interacts with the storage API.

In Ibis v6 and upwards we have the option of passing in a BigQuery client and a BigQuery storage API client which has two benefits.

1. We would be able to easily fix this problem
2. We could avoid the monkey patching we do now for the BigQuery client which ends up making two connections, the standard Ibis one and then our custom one which overrides the original.

Obviously upgrading Ibis is non trivial so I need to consider other options.@nj1973: This recent new feature is another reason why we need to get to the bottom of this.

https://cloud.google.com/bigquery/docs/regional-endpoints

Solution patch:
diff --git a/README.md b/README.md
index b3778f8..92ecafc 100644
--- a/README.md
+++ b/README.md
@@ -66,8 +66,10 @@ please review the [Connections](https://github.com/GoogleCloudPlatform/professio
 ### Running Validations
 
 The CLI is the main interface to use this tool and it has several different
-commands which can be used to create and run validations. Below are the command
-syntax and options for running validations.
+commands which can be used to create and run validations. DVT is designed to run in
+an environment connected to GCP services, specifically, BigQuery, GCS and Secret manager.
+If DVT is being run on-premises or in an environment with restricted access to GCP services, see
+[running DVT  at on-prem](#running-dvt-at-on-prem). Below are the command syntax and options for running validations.
 
 Alternatives to running DVT in the CLI include deploying DVT to Cloud Run, Cloud Functions, or Airflow
 ([Examples Here](https://github.com/GoogleCloudPlatform/professional-services-data-validator/tree/develop/samples)). See the [Validation Logic](https://github.com/GoogleCloudPlatform/professional-services-data-validator#validation-logic) section
@@ -525,6 +527,11 @@ For example, this flag can be used as follows:
     "target_query": "SELECT `hash__all`, `station_id`\nFROM ..."
 }
 ```
+#### Running DVT at on-prem
+On-premises environments can have limited access to GCP services. DVT supports using BigQuery for storing validation results, GCS for storage and
+the Secret Manager for secrets. You may also use BigQuery and Spanner as a source or target for validation. Service
+APIs (i.e. bigquery.googleapis.com) may not always be accessible due to firewall restrictions. Work with your network
+adminstrator to identify the way to access these services. They may set up a [Private Service Connect Endpoint](https://cloud.google.com/vpc/docs/about-accessing-google-apis-endpoints). DVT supports accessing source and target tables in Spanner and BigQuery via endpoints set up in your network. Connection Parameters for [Spanner](https://github.com/GoogleCloudPlatform/professional-services-data-validator/blob/develop/docs/connections.md#google-spanner) and [BigQuery](https://github.com/GoogleCloudPlatform/professional-services-data-validator/blob/develop/docs/connections.md#google-bigquery) outline regarding how to specify endpoints.
 
 ### Running DVT with YAML Configuration Files
 
diff --git a/data_validation/cli_tools.py b/data_validation/cli_tools.py
index fc6cf77..fec79a7 100644
--- a/data_validation/cli_tools.py
+++ b/data_validation/cli_tools.py
@@ -76,7 +76,11 @@ CONNECTION_SOURCE_FIELDS = {
         ["google_service_account_key_path", "(Optional) GCP SA Key Path"],
         [
             "api_endpoint",
-            '(Optional) GCP BigQuery API endpoint (e.g. "https://mybq.p.googleapis.com")',
+            '(Optional) GCP BigQuery API endpoint (e.g. "https://bigquery-mypsc.p.googleapis.com")',
+        ],
+        [
+            "storage_api_endpoint",
+            '(Optional) GCP BigQuery Storage API endpoint (e.g. "https://bigquerystorage-mypsc.p.googleapis.com")',
         ],
     ],
     consts.SOURCE_TYPE_TERADATA: [
@@ -140,7 +144,7 @@ CONNECTION_SOURCE_FIELDS = {
         ["google_service_account_key_path", "(Optional) GCP SA Key Path"],
         [
             "api_endpoint",
-            '(Optional) GCP Spanner API endpoint (e.g. "https://mycs.p.googleapis.com")',
+            '(Optional) GCP Spanner API endpoint (e.g. "https://spanner-mypsc.p.googleapis.com")',
         ],
     ],
     consts.SOURCE_TYPE_FILESYSTEM: [
@@ -1284,6 +1288,9 @@ def _get_result_handler(rc_value: str, sa_file=None) -> dict:
                 consts.PROJECT_ID: conn_from_file["project_id"],
                 consts.TABLE_ID: config[1],
                 consts.API_ENDPOINT: conn_from_file.get("api_endpoint", None),
+                consts.STORAGE_API_ENDPOINT: conn_from_file.get(
+                    "storage_api_endpoint", None
+                ),
             }
         elif conn_from_file[consts.SOURCE_TYPE] == consts.SOURCE_TYPE_POSTGRES:
             result_handler = {
diff --git a/data_validation/clients.py b/data_validation/clients.py
index 202e96e..a4f74d3 100644
--- a/data_validation/clients.py
+++ b/data_validation/clients.py
@@ -27,6 +27,8 @@ import pandas
 
 from data_validation import client_info, consts, exceptions
 from data_validation.secret_manager import SecretManagerBuilder
+
+from third_party.ibis.ibis_bigquery.api import bigquery_connect
 from third_party.ibis.ibis_cloud_spanner.api import spanner_connect
 from third_party.ibis.ibis_impala.api import impala_connect
 from third_party.ibis.ibis_mssql.api import mssql_connect
@@ -111,24 +113,42 @@ def get_google_bigquery_client(
     )
 
 
+def _get_google_bqstorage_client(credentials=None, api_endpoint: str = None):
+    options = None
+    if api_endpoint:
+        options = client_options.ClientOptions(api_endpoint=api_endpoint)
+    from google.cloud import bigquery_storage_v1 as bigquery_storage
+
+    return bigquery_storage.BigQueryReadClient(
+        credentials=credentials,
+        client_options=options,
+    )
+
+
 def get_bigquery_client(
-    project_id: str, dataset_id: str = "", credentials=None, api_endpoint: str = None
+    project_id: str,
+    dataset_id: str = "",
+    credentials=None,
+    api_endpoint: str = None,
+    storage_api_endpoint: str = None,
 ):
     google_client = get_google_bigquery_client(
         project_id, credentials=credentials, api_endpoint=api_endpoint
     )
+    bqstorage_client = None
+    if storage_api_endpoint:
+        bqstorage_client = _get_google_bqstorage_client(
+            credentials=credentials, api_endpoint=storage_api_endpoint
+        )
 
-    ibis_client = ibis.bigquery.connect(
+    return bigquery_connect(
         project_id=project_id,
         dataset_id=dataset_id,
         credentials=credentials,
+        bigquery_client=google_client,
+        bqstorage_client=bqstorage_client,
     )
 
-    # Override the BigQuery client object to ensure the correct user agent is
-    # included and any api_endpoint is used.
-    ibis_client.client = google_client
-    return ibis_client
-
 
 def get_pandas_client(table_name, file_path, file_type):
     """Return pandas client and env with file loaded into DataFrame
diff --git a/data_validation/consts.py b/data_validation/consts.py
index dd41d9e..21c3e1e 100644
--- a/data_validation/consts.py
+++ b/data_validation/consts.py
@@ -143,6 +143,7 @@ PROJECT_ID = "project_id"
 TABLE_ID = "table_id"
 GOOGLE_SERVICE_ACCOUNT_KEY_PATH = "google_service_account_key_path"
 API_ENDPOINT = "api_endpoint"
+STORAGE_API_ENDPOINT = "storage_api_endpoint"
 
 # Result Handler Output Table Fields
 VALIDATION_TYPE = "validation_type"
diff --git a/data_validation/result_handlers/bigquery.py b/data_validation/result_handlers/bigquery.py
index 5af8e4b..b00a631 100644
--- a/data_validation/result_handlers/bigquery.py
+++ b/data_validation/result_handlers/bigquery.py
@@ -65,7 +65,7 @@ class BigQueryResultHandler(BaseBackendResultHandler):
                 Explicit credentials to use in case default credentials
                 aren't working properly.
             status_list (list): provided status to filter the results with
-            api_endpoint (str): BigQuery API endpoint (e.g. https://mybq.p.googleapis.com)
+            api_endpoint (str): BigQuery API endpoint (e.g. https://bigquery-mypsc.p.googleapis.com)
             text_format (str, optional):
                 This allows the user to influence the text results written via logger.debug.
                 See: https://github.com/GoogleCloudPlatform/professional-services-data-validator/issues/871
diff --git a/docs/connections.md b/docs/connections.md
index 32f554f..6289df8 100644
--- a/docs/connections.md
+++ b/docs/connections.md
@@ -123,8 +123,12 @@ data-validation connections add
     --connection-name CONN_NAME BigQuery                Connection name
     --project-id MY_PROJECT                             Project ID where BQ data resides
     [--google-service-account-key-path PATH_TO_SA_KEY]  Path to SA key
-    [--api-endpoint ENDPOINT_URI]                       BigQuery API endpoint (e.g.
-                                                        "https://mybq.p.googleapis.com)
+    [--api-endpoint API_ENDPOINT]                       BigQuery API endpoint (e.g.
+                                                        "https://bigquery-mypsc.p.googleapis.com)
+    [--storage-api-endpoint STORAGE_API_ENDPOINT]       BigQuery Storage API endpoint (e.g.
+                                                        "bigquerystorage-mypsc.p.googleapis.com)
+                                                        Note this is a GRPC endpoint and does not
+                                                        include a URI scheme.
 ```
 
 ### User/Service account needs following BigQuery permissions to run DVT
@@ -150,8 +154,8 @@ data-validation connections add
     --instance-id MY_INSTANCE                           Spanner instance to connect to
     --database-id MY-DB                                 Spanner database (schema) to connect to
     [--google-service-account-key-path PATH_TO_SA_KEY]  Path to SA key
-    [--api-endpoint ENDPOINT_URI]                       Spanner API endpoint (e.g.
-                                                        "https://mycs.p.googleapis.com)
+    [--api-endpoint API_ENDPOINT]                       Spanner API endpoint (e.g.
+                                                        "https://spanner-mypsc.p.googleapis.com")
 ```
 
 ### User/Service account needs following Spanner role to run DVT
diff --git a/third_party/ibis/ibis_addon/operations.py b/third_party/ibis/ibis_addon/operations.py
index 46946c8..ea46bad 100644
--- a/third_party/ibis/ibis_addon/operations.py
+++ b/third_party/ibis/ibis_addon/operations.py
@@ -61,7 +61,7 @@ from ibis.backends.postgres.compiler import PostgreSQLExprTranslator
 from ibis.expr.types import BinaryValue, NumericValue, StringValue, TemporalValue
 
 # Do not remove these lines, they trigger patching of Ibis code.
-import third_party.ibis.ibis_biquery.api  # noqa
+import third_party.ibis.ibis_bigquery.api  # noqa
 import third_party.ibis.ibis_mysql.compiler  # noqa
 from third_party.ibis.ibis_mssql import registry as mssql_registry
 from third_party.ibis.ibis_postgres import registry as postgres_registry
diff --git a/third_party/ibis/ibis_bigquery/__init__.py b/third_party/ibis/ibis_bigquery/__init__.py
new file mode 100644
index 0000000..7627398
--- /dev/null
+++ b/third_party/ibis/ibis_bigquery/__init__.py
@@ -0,0 +1,142 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+import google.auth.credentials
+import google.cloud.bigquery as bq
+import pydata_google_auth
+from pydata_google_auth import cache
+
+from ibis.backends.bigquery import (
+    Backend as BigQueryBackend,
+    _create_client_info,
+    parse_project_and_dataset,
+    CLIENT_ID,
+    CLIENT_SECRET,
+    EXTERNAL_DATA_SCOPES,
+    SCOPES,
+)
+
+if TYPE_CHECKING:
+    import google.cloud.bigquery_storage_v1
+
+
+class Backend(BigQueryBackend):
+    def __init__(self):
+        super().__init__()
+        self.storage_client = None
+
+    def do_connect(
+        self,
+        project_id: str = None,
+        dataset_id: str = "",
+        credentials: google.auth.credentials.Credentials = None,
+        application_name: str = None,
+        auth_local_webserver: bool = True,
+        auth_external_data: bool = False,
+        auth_cache: str = "default",
+        partition_column: str = "PARTITIONTIME",
+        # Custom DVT arguments:
+        bigquery_client: bq.Client = None,
+        bqstorage_client: "google.cloud.bigquery_storage_v1.BigQueryReadClient" = None,
+    ):
+        """Copy of Ibis v5 BigQuery do_connect() customized for DVT, see original method for docs."""
+        default_project_id = ""
+
+        if credentials is None:
+            scopes = SCOPES
+            if auth_external_data:
+                scopes = EXTERNAL_DATA_SCOPES
+
+            if auth_cache == "default":
+                credentials_cache = cache.ReadWriteCredentialsCache(
+                    filename="ibis.json"
+                )
+            elif auth_cache == "reauth":
+                credentials_cache = cache.WriteOnlyCredentialsCache(
+                    filename="ibis.json"
+                )
+            elif auth_cache == "none":
+                credentials_cache = cache.NOOP
+            else:
+                raise ValueError(
+                    f"Got unexpected value for auth_cache = '{auth_cache}'. "
+                    "Expected one of 'default', 'reauth', or 'none'."
+                )
+
+            credentials, default_project_id = pydata_google_auth.default(
+                scopes,
+                client_id=CLIENT_ID,
+                client_secret=CLIENT_SECRET,
+                credentials_cache=credentials_cache,
+                use_local_webserver=auth_local_webserver,
+            )
+
+        project_id = project_id or default_project_id
+
+        (
+            self.data_project,
+            self.billing_project,
+            self.dataset,
+        ) = parse_project_and_dataset(project_id, dataset_id)
+
+        if bigquery_client is None:
+            self.client = bq.Client(
+                project=self.billing_project,
+                credentials=credentials,
+                client_info=_create_client_info(application_name),
+            )
+        else:
+            self.client = bigquery_client
+        self.partition_column = partition_column
+        self.storage_client = bqstorage_client
+
+    def _cursor_to_arrow(
+        self,
+        cursor,
+        *,
+        method=None,
+        chunk_size: int = None,
+    ):
+        """Copy of Ibis v5 BigQuery _cursor_to_arrow() except can use custom DVT storage client"""
+        if method is None:
+
+            def method(result, storage_client=self.storage_client):
+                return result.to_arrow(
+                    progress_bar_type=None,
+                    # Include DVT specific storage client.
+                    bqstorage_client=storage_client,
+                    create_bqstorage_client=bool(not self.storage_client),
+                )
+
+        query = cursor.query
+        query_result = query.result(page_size=chunk_size)
+        # workaround potentially not having the ability to create read sessions
+        # in the dataset project
+        orig_project = query_result._project
+        query_result._project = self.billing_project
+        try:
+            arrow_obj = method(query_result)
+        finally:
+            query_result._project = orig_project
+        return arrow_obj
+
+    def list_primary_key_columns(self, database: str, table: str) -> list:
+        """Return a list of primary key column names."""
+        # TODO: Related to issue-1253, it's not clear if this is possible, we should revisit if it becomes a requirement.
+        return None
+
+    def dvt_list_tables(self, like=None, database=None):
+        return self.list_tables(like=like, database=database)
diff --git a/third_party/ibis/ibis_bigquery/api.py b/third_party/ibis/ibis_bigquery/api.py
new file mode 100644
index 0000000..8330115
--- /dev/null
+++ b/third_party/ibis/ibis_bigquery/api.py
@@ -0,0 +1,35 @@
+# Copyright 2024 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from third_party.ibis.ibis_bigquery import Backend as BigQueryBackend
+
+
+def bigquery_connect(
+    project_id: str = None,
+    dataset_id: str = "",
+    credentials=None,
+    bigquery_client=None,
+    bqstorage_client=None,
+):
+    """Create a BigQuery Backend for use with Ibis."""
+    backend = BigQueryBackend()
+    backend.do_connect(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        credentials=credentials,
+        bigquery_client=bigquery_client,
+        bqstorage_client=bqstorage_client,
+    )
+    return backend
diff --git a/third_party/ibis/ibis_biquery/__init__.py b/third_party/ibis/ibis_biquery/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/third_party/ibis/ibis_biquery/api.py b/third_party/ibis/ibis_biquery/api.py
deleted file mode 100644
index 94b5c72..0000000
--- a/third_party/ibis/ibis_biquery/api.py
+++ /dev/null
@@ -1,15 +0,0 @@
-from ibis.backends.bigquery import Backend as BigQueryBackend
-
-
-def _list_primary_key_columns(self, database: str, table: str) -> list:
-    """Return a list of primary key column names."""
-    # TODO: Related to issue-1253, it's not clear if this is possible, we should revisit if it becomes a requirement.
-    return None
-
-
-def _dvt_list_tables(self, like=None, database=None):
-    return self.list_tables(like=like, database=database)
-
-
-BigQueryBackend.list_primary_key_columns = _list_primary_key_columns
-BigQueryBackend.dvt_list_tables = _dvt_list_tables


Test patch:
diff --git a/tests/unit/test__main.py b/tests/unit/test__main.py
index 1888e63..da3aeae 100644
--- a/tests/unit/test__main.py
+++ b/tests/unit/test__main.py
@@ -152,7 +152,8 @@ CONNECTION_ADD_ARGS = {
     consts.PROJECT_ID: "dummy-gcp-project",
     consts.GOOGLE_SERVICE_ACCOUNT_KEY_PATH: None,
     "connection_name": "dummy-bq-connection",
-    "api_endpoint": None,
+    consts.API_ENDPOINT: None,
+    consts.STORAGE_API_ENDPOINT: None,
 }
 CONNECTION_DESCRIBE_ARGS = {
     "verbose": False,
@@ -180,7 +181,7 @@ BROKEN_CONNECTION_CONFIG_INCORRECT_COMMAND = {
     consts.PROJECT_ID: "dummy-gcp-project",
     consts.GOOGLE_SERVICE_ACCOUNT_KEY_PATH: None,
     "connection_name": "dummy-bq-connection",
-    "api_endpoint": None,
+    consts.API_ENDPOINT: None,
 }  # same as CONNECTION_ADD_ARGS but with the command item replaced
 FIND_TABLES_ARGS = {
     "verbose": False,
diff --git a/tests/unit/test_cli_tools.py b/tests/unit/test_cli_tools.py
index 892c56a..6ef6c58 100644
--- a/tests/unit/test_cli_tools.py
+++ b/tests/unit/test_cli_tools.py
@@ -83,7 +83,9 @@ CLI_ADD_BQ_CONNECTION_ARGS = [
     "--project-id",
     "example-project",
     "--api-endpoint",
-    "https://mybq.p.googleapis.com",
+    "https://bigquery-mypsc.p.googleapis.com",
+    "--storage-api-endpoint",
+    "https://bigquerystorage-mypsc.p.googleapis.com",
 ]
 
 
@@ -309,7 +311,7 @@ def test_create_bq_connection(caplog, fs):
     assert bq_conn[consts.SOURCE_TYPE] == consts.SOURCE_TYPE_BIGQUERY
 
     conn_from_file = cli_tools.get_connection(args.connection_name)
-    assert conn_from_file["api_endpoint"] == "https://mybq.p.googleapis.com"
+    assert conn_from_file["api_endpoint"] == "https://bigquery-mypsc.p.googleapis.com"
 
 
 @mock.patch(
@@ -625,6 +627,7 @@ def test_get_result_handler_by_conn_file(fs):
         consts.PROJECT_ID: args.project_id,
         consts.TABLE_ID: "dataset.table",
         consts.API_ENDPOINT: args.api_endpoint,
+        consts.STORAGE_API_ENDPOINT: args.storage_api_endpoint,
     }
 
     # Plus check standard format still works.