google-pso-data-validator 7.6.1__tar.gz → 7.7.0__tar.gz

{google_pso_data_validator-7.6.1/google_pso_data_validator.egg-info → google_pso_data_validator-7.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-pso-data-validator
-Version: 7.6.1
+Version: 7.7.0
 Summary: A package to enable easy data validation
 Author: PSO DVT Engineering team
 Author-email: data-validator-eng@google.com
@@ -125,8 +125,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
@@ -584,6 +586,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
 

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/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
 

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/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 = {

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/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

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/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"

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/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

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0/google_pso_data_validator.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-pso-data-validator
-Version: 7.6.1
+Version: 7.7.0
 Summary: A package to enable easy data validation
 Author: PSO DVT Engineering team
 Author-email: data-validator-eng@google.com
@@ -125,8 +125,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
@@ -584,6 +586,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
 

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/google_pso_data_validator.egg-info/SOURCES.txt

@@ -43,8 +43,8 @@ google_pso_data_validator.egg-info/top_level.txt
 third_party/ibis/ibis_addon/__init__.py
 third_party/ibis/ibis_addon/api.py
 third_party/ibis/ibis_addon/operations.py
-third_party/ibis/ibis_biquery/__init__.py
-third_party/ibis/ibis_biquery/api.py
+third_party/ibis/ibis_bigquery/__init__.py
+third_party/ibis/ibis_bigquery/api.py
 third_party/ibis/ibis_cloud_spanner/__init__.py
 third_party/ibis/ibis_cloud_spanner/api.py
 third_party/ibis/ibis_cloud_spanner/client.py

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/setup.py

@@ -18,7 +18,7 @@ import setuptools
 
 name = "google-pso-data-validator"
 description = "A package to enable easy data validation"
-version = "7.6.1"
+version = "7.7.0"
 release_status = "Development Status :: 4 - Beta"
 
 with open("README.md", "r") as fh:

{google_pso_data_validator-7.6.1 → google_pso_data_validator-7.7.0}/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
@@ -329,40 +329,6 @@ def sa_format_binary_length_oracle(translator, op):
     return sa.func.dbms_lob.getlength(arg)
 
 
-def sa_cast_decimal_when_scale_padded_fmt_fm(t, op):
-    """Caters for engines that fully pad scale with 0s when casting decimal to string and support FM format."""
-    # Add cast from numeric to string
-    arg = op.arg
-    typ = op.to
-    arg_dtype = arg.output_dtype
-
-    # Specialize going from numeric(p,s>0) to string
-    if (
-        arg_dtype.is_decimal()
-        and arg_dtype.scale
-        and arg_dtype.scale > 0
-        and typ.is_string()
-    ):
-        sa_arg = t.translate(arg)
-        # When casting a number to string PostgreSQL and Snowflake include the full scale, e.g.:
-        #   SELECT CAST(CAST(100 AS DECIMAL(5,2)) AS VARCHAR(10));
-        #     100.00
-        # This doesn't match most engines which would return "100".
-        # Using to_char() function instead of cast to return a more typical value.
-        # We've wrapped to_char in rtrim(".") due to whole numbers having a trailing ".".
-        # Would have liked to use trim_scale but this is only available in PostgreSQL 13+
-        #     return (sa.cast(sa.func.trim_scale(arg), typ))
-        precision = arg_dtype.precision or 38
-        fmt = (
-            "FM"
-            + ("9" * (precision - arg_dtype.scale - 1))
-            + "0."
-            + ("9" * arg_dtype.scale)
-        )
-        return sa.func.rtrim(sa.func.to_char(sa_arg, fmt), ".")
-    return None
-
-
 def sa_cast_hive(t, op):
     arg = op.arg
     typ = op.to
@@ -387,27 +353,6 @@ def sa_cast_hive(t, op):
         return cast_expr
 
 
-def sa_cast_postgres(t, op):
-    custom_cast = sa_cast_decimal_when_scale_padded_fmt_fm(t, op)
-    if custom_cast is not None:
-        return custom_cast
-
-    arg = op.arg
-    typ = op.to
-    arg_dtype = arg.output_dtype
-
-    sa_arg = t.translate(arg)
-    if arg_dtype.is_binary() and typ.is_string():
-        # Binary to string cast is a "to hex" conversion for DVT.
-        return sa.func.encode(sa_arg, sa.literal("hex"))
-    elif arg_dtype.is_string() and typ.is_binary():
-        # Binary from string cast is a "from hex" conversion for DVT.
-        return sa.func.decode(sa_arg, sa.literal("hex"))
-
-    # Follow the original Ibis code path.
-    return sa_fixed_cast(t, op)
-
-
 def sa_cast_mssql(t, op):
     arg = op.arg
     typ = op.to
@@ -477,15 +422,33 @@ def sa_cast_mysql(t, op):
 
 
 def sa_cast_snowflake(t, op):
-    custom_cast = sa_cast_decimal_when_scale_padded_fmt_fm(t, op)
-    if custom_cast is not None:
-        return custom_cast
-
     arg = op.arg
     typ = op.to
     arg_dtype = arg.output_dtype
-
     sa_arg = t.translate(arg)
+
+    # Specialize going from numeric(p,s>0) to string
+    if (
+        arg_dtype.is_decimal()
+        and arg_dtype.scale
+        and arg_dtype.scale > 0
+        and typ.is_string()
+    ):
+        # When casting a number to string Snowflake includes the full scale, e.g.:
+        #   SELECT CAST(CAST(100 AS DECIMAL(5,2)) AS VARCHAR(10));
+        #     100.00
+        # This doesn't match most engines which would return "100".
+        # Using to_char() function instead of cast to return a more typical value.
+        # We've wrapped to_char in rtrim(".") due to whole numbers having a trailing ".".
+        precision = arg_dtype.precision or 38
+        fmt = (
+            "FM"
+            + ("9" * (precision - arg_dtype.scale - 1))
+            + "0."
+            + ("9" * arg_dtype.scale)
+        )
+        return sa.func.rtrim(sa.func.to_char(sa_arg, fmt), ".")
+
     if arg_dtype.is_binary() and typ.is_string():
         # Binary to string cast is a "to hex" conversion for DVT.
         return sa.func.hex_encode(sa_arg, sa.literal(0))
@@ -648,7 +611,7 @@ PostgreSQLExprTranslator._registry[
 ] = postgres_registry.sa_format_hashbytes
 PostgreSQLExprTranslator._registry[RawSQL] = sa_format_raw_sql
 PostgreSQLExprTranslator._registry[ToChar] = sa_format_to_char
-PostgreSQLExprTranslator._registry[ops.Cast] = sa_cast_postgres
+PostgreSQLExprTranslator._registry[ops.Cast] = postgres_registry.sa_cast_postgres
 PostgreSQLExprTranslator._registry[BinaryLength] = sa_format_binary_length
 PostgreSQLExprTranslator._registry[
     ops.ExtractEpochSeconds

google_pso_data_validator-7.7.0/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)

google_pso_data_validator-7.7.0/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

google_pso_data_validator-7.7.0/third_party/ibis/ibis_postgres/registry.py

@@ -0,0 +1,86 @@
+# 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.
+
+import sqlalchemy as sa
+from ibis.backends.base.sql.alchemy.registry import _cast as sa_fixed_cast
+
+
+def sa_format_hashbytes(translator, op):
+    arg = translator.translate(op.arg)
+    convert = sa.func.convert_to(arg, sa.sql.literal_column("'UTF8'"))
+    hash_func = sa.func.sha256(convert)
+    return sa.func.encode(hash_func, sa.sql.literal_column("'hex'"))
+
+
+def sa_epoch_seconds(translator, op):
+    """Override for standard ExtractEpochSeconds but catering for larger second values.
+
+    This expression also truncates fractional seconds from the incoming datetime.
+    This matches behaviour on other SQL engines' epoch seconds expressions."""
+    arg = translator.translate(op.arg)
+    return sa.cast(
+        sa.extract("epoch", sa.func.date_trunc(sa.sql.literal_column("'second'"), arg)),
+        sa.BIGINT,
+    )
+
+
+def sa_cast_postgres(t, op):
+    arg = op.arg
+    typ = op.to
+    arg_dtype = arg.output_dtype
+    sa_arg = t.translate(arg)
+
+    if arg_dtype.is_decimal() and typ.is_string():
+        if arg_dtype.scale is None:
+            # trim_scale() is only available in PostgreSQL 13+ but solves a lot of problems
+            # when trying to get consistently formatted numerics.
+            # We've documented a workaround for PostgreSQL 12 and prior versions.
+            return sa.cast(sa.func.trim_scale(sa_arg), t.get_sqla_type(typ))
+        elif arg_dtype.scale > 0:
+            # When casting a number to string PostgreSQL includes the full scale, e.g.:
+            #   SELECT CAST(CAST(100 AS DECIMAL(5,2)) AS VARCHAR(10));
+            #     100.00
+            # This doesn't match most engines which would return "100".
+            # Using to_char() function instead of cast to return a more typical value.
+            # We've wrapped to_char in rtrim(".") due to whole numbers having a trailing ".".
+            #
+            # Potentially we could use trim_scale() here like when scale is None but this is
+            # a much more common scenario than decimal(scale=None) so I'm trying to minimize
+            # risk of needing a UDF.
+            precision = arg_dtype.precision or 38
+            fmt = (
+                "FM"
+                + ("9" * (precision - arg_dtype.scale - 1))
+                + "0."
+                + ("9" * arg_dtype.scale)
+            )
+            return sa.func.rtrim(sa.func.to_char(sa_arg, fmt), ".")
+    elif arg_dtype.is_binary() and typ.is_string():
+        # Binary to string cast is a "to hex" conversion for DVT.
+        return sa.func.encode(sa_arg, sa.literal("hex"))
+    elif arg_dtype.is_string() and typ.is_binary():
+        # Binary from string cast is a "from hex" conversion for DVT.
+        return sa.func.decode(sa_arg, sa.literal("hex"))
+
+    # Follow the original Ibis code path.
+    return sa_fixed_cast(t, op)
+
+
+def sa_format_postgres_padded_char_length(translator, op):
+    """Inject concat of empty string into length(bpchar) to protect trailing spaces.
+
+    Without this workaround the bpchar value is implicitly cast to varchar and loses trailing spaces.
+    """
+    arg = translator.translate(op.arg)
+    return sa.func.char_length(sa.func.concat(arg, sa.text("''")))

google_pso_data_validator-7.6.1/third_party/ibis/ibis_biquery/api.py

@@ -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

google_pso_data_validator-7.6.1/third_party/ibis/ibis_postgres/registry.py

@@ -1,43 +0,0 @@
-# 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.
-
-import sqlalchemy as sa
-
-
-def sa_format_hashbytes(translator, op):
-    arg = translator.translate(op.arg)
-    convert = sa.func.convert_to(arg, sa.sql.literal_column("'UTF8'"))
-    hash_func = sa.func.sha256(convert)
-    return sa.func.encode(hash_func, sa.sql.literal_column("'hex'"))
-
-
-def sa_epoch_seconds(translator, op):
-    """Override for standard ExtractEpochSeconds but catering for larger second values.
-
-    This expression also truncates fractional seconds from the incoming datetime.
-    This matches behaviour on other SQL engines' epoch seconds expressions."""
-    arg = translator.translate(op.arg)
-    return sa.cast(
-        sa.extract("epoch", sa.func.date_trunc(sa.sql.literal_column("'second'"), arg)),
-        sa.BIGINT,
-    )
-
-
-def sa_format_postgres_padded_char_length(translator, op):
-    """Inject concat of empty string into length(bpchar) to protect trailing spaces.
-
-    Without this workaround the bpchar value is implicitly cast to varchar and loses trailing spaces.
-    """
-    arg = translator.translate(op.arg)
-    return sa.func.char_length(sa.func.concat(arg, sa.text("''")))