google-pso-data-validator 7.4.0__tar.gz → 7.5.0__tar.gz

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: google-pso-data-validator
-Version: 7.4.0
+Version: 7.5.0
 Summary: A package to enable easy data validation
 Author: PSO DVT Engineering team
 Author-email: data-validator-eng@google.com
@@ -51,6 +51,7 @@ Dynamic: author-email
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/data_validation/combiner.py

@@ -22,10 +22,12 @@ import datetime
 import functools
 import json
 import logging
+from typing import TYPE_CHECKING
+
 import ibis
 import ibis.expr.datatypes as dt
+import pandas
 
-from typing import TYPE_CHECKING
 from data_validation import consts
 
 if TYPE_CHECKING:
@@ -34,22 +36,109 @@ if TYPE_CHECKING:
     from data_validation.metadata import RunMetadata, ValidationMetadata
 
 
+# At around 140 columns we hit RecursionError when unioning Ibis subqueries.
+# This constant is a threshold at which we slice up the input Dataframes
+# and stitch them back together again after Ibis processing.
+COMBINER_COLUMN_SLICE_WIDTH = 120
+
+COMBINER_GET_SUMMARY_EXC_TEXT = (
+    "Error while generating summary report of row validation results"
+)
+
+
 def generate_report(
-    client,
     run_metadata: "RunMetadata",
-    source: "IbisTable",
-    target: "IbisTable",
+    source_df: "DataFrame",
+    target_df: "DataFrame",
     join_on_fields=(),
     is_value_comparison=False,
     verbose=False,
-):
+) -> "DataFrame":
+    """Combine results into a report.
+
+    This function is a wrapper around _generate_report_slice(). _generate_report_slice() does the main work, this
+    wrapper simply manages the input columns and stitches the results back together.
+    This is because validations of > 140(ish) columns trigger a RecursionError when unioning Ibis subqueries.
+    In this method we pass in column slices of the incoming Dataframes and combine the results.
+
+    It is a bit of a hack but I cannot find a way to optimize the Ibis processing. It appears to be
+    inefficient in that we create a subquery for each validation (column) in _calculate_differences() and
+    then union them all. We then do the same on the source/target table expressions to join it all back
+    together again. I (nj1973) spent a singificant amount of time trying to understand/optimize the Ibis
+    processing but fell back on this simpler (less risky) workaround.
+
+    Returns:
+        pandas.DataFrame:
+            A pandas DataFrame with the results of the validation in the same
+            schema as the report table.
+    """
+    _check_schema_names(source_df, target_df)
+
+    join_on_fields = tuple(join_on_fields)
+
+    validation_columns = run_metadata.validations.keys()
+    # slice_thresholds is a list of points at which we should break up the Dataframe by column.
+    # e.g. [10, 20, 30] would mean process columns 0-9, 10-19 and 20-the max column.
+    # 1. len(...) / COMBINER_COLUMN_SLICE_WIDTH: Divides total columns by the slice width to get the number of slices.
+    # 2. int(...) + 1: int()+1 is effectively ceil() which is what we want to get the actual whole number of slices
+    # 3. _ * COMBINER_COLUMN_SLICE_WIDTH: Multiplies each number by the slice width to get actual column counts for each slice.
+    slice_thresholds = [
+        (_ * COMBINER_COLUMN_SLICE_WIDTH)
+        for _ in range(int(len(validation_columns) / COMBINER_COLUMN_SLICE_WIDTH) + 1)
+    ]
+
+    result_df = None
+    # Process the input Dataframes in slices of columns to avoid "RecursionError"s.
+    for slice_start in slice_thresholds:
+        columns_in_vertical_slice = list(validation_columns)[
+            slice_start : slice_start + COMBINER_COLUMN_SLICE_WIDTH
+        ]
+        # Ensure any join columns are included in the column slice.
+        columns_in_vertical_slice.extend(
+            set(join_on_fields) - set(columns_in_vertical_slice)
+        )
+        interim_result_df = _generate_report_slice(
+            run_metadata,
+            source_df[columns_in_vertical_slice],
+            target_df[columns_in_vertical_slice],
+            join_on_fields=join_on_fields,
+            is_value_comparison=is_value_comparison,
+            verbose=verbose,
+        )
+        if result_df is None:
+            result_df = interim_result_df
+        else:
+            result_df = pandas.concat([result_df, interim_result_df])
+
+    # Get the first validation metadata object to fill source and/or target empty table names.
+    first = run_metadata.validations[next(iter(run_metadata.validations))]
+    if first.validation_type != consts.CUSTOM_QUERY:
+        result_df.source_table_name.fillna(
+            first.get_table_name(consts.RESULT_TYPE_SOURCE), inplace=True
+        )
+        result_df.target_table_name.fillna(
+            first.get_table_name(consts.RESULT_TYPE_TARGET), inplace=True
+        )
+
+    _get_summary(run_metadata, result_df, source_df, target_df)
+
+    return result_df
+
+
+def _generate_report_slice(
+    run_metadata: "RunMetadata",
+    source_df: "DataFrame",
+    target_df: "DataFrame",
+    join_on_fields=(),
+    is_value_comparison=False,
+    verbose=False,
+) -> "DataFrame":
     """Combine results into a report.
 
     Args:
-        client (ibis.client.Client): Ibis client used to combine results.
         run_metadata: Metadata about the run and validations.
-        source (ibis.QUERY): Ibis query / table object.
-        target (ibis.QUERY): Ibis query / table object.
+        source_df: Dataframe contains results of source query.
+        target_df: Dataframe contains results of target query.
         join_on_fields (Sequence[str]):
             A collection of column names to use to join source and target.
             These are the columns that both the source and target queries
@@ -63,16 +152,15 @@ def generate_report(
             A pandas DataFrame with the results of the validation in the same
             schema as the report table.
     """
-    join_on_fields = tuple(join_on_fields)
-
-    source_names = source.schema().names
-    target_names = target.schema().names
+    client = ibis.pandas.connect(
+        {
+            consts.RESULT_TYPE_SOURCE: source_df,
+            consts.RESULT_TYPE_TARGET: target_df,
+        }
+    )
+    source = client.table(consts.RESULT_TYPE_SOURCE)
+    target = client.table(consts.RESULT_TYPE_TARGET)
 
-    if source_names != target_names:
-        raise ValueError(
-            "Expected source and target to have same schema, got "
-            f"{consts.RESULT_TYPE_SOURCE}: {source_names}; {consts.RESULT_TYPE_TARGET}: {target_names}"
-        )
     differences_pivot = _calculate_differences(
         source, target, join_on_fields, run_metadata.validations, is_value_comparison
     )
@@ -81,18 +169,18 @@ def generate_report(
     source_pivot = _pivot_result(
         source, join_on_fields, run_metadata.validations, consts.RESULT_TYPE_SOURCE
     )
-    source_df = client.execute(source_pivot)
+    source_pivot_df = client.execute(source_pivot)
 
     target_pivot = _pivot_result(
         target, join_on_fields, run_metadata.validations, consts.RESULT_TYPE_TARGET
     )
-    target_df = client.execute(target_pivot)
+    target_pivot_df = client.execute(target_pivot)
 
     con = ibis.pandas.connect(
         {
-            consts.RESULT_TYPE_SOURCE: source_df,
-            "differences": differences_df,
-            consts.RESULT_TYPE_TARGET: target_df,
+            consts.RESULT_TYPE_SOURCE: source_pivot_df,
+            consts.RESULT_TYPE_DIFFERENCES: differences_df,
+            consts.RESULT_TYPE_TARGET: target_pivot_df,
         }
     )
     joined = _join_pivots(
@@ -107,19 +195,6 @@ def generate_report(
 
     result_df = client.execute(documented)
     result_df.validation_status.fillna(consts.VALIDATION_STATUS_FAIL, inplace=True)
-
-    # get the first validation metadata object to fill source and/or target empty table names
-    first = run_metadata.validations[next(iter(run_metadata.validations))]
-    if first.validation_type != consts.CUSTOM_QUERY:
-        result_df.source_table_name.fillna(
-            first.get_table_name(consts.RESULT_TYPE_SOURCE), inplace=True
-        )
-        result_df.target_table_name.fillna(
-            first.get_table_name(consts.RESULT_TYPE_TARGET), inplace=True
-        )
-
-    _get_summary(run_metadata, result_df, source_df, target_df)
-
     return result_df
 
 
@@ -279,6 +354,18 @@ def _calculate_differences(
     return differences_pivot
 
 
+def _check_schema_names(source_df, target_df):
+    """Check that the two input Dataframes have matching column names."""
+    source_names = tuple(source_df.columns)
+    target_names = tuple(target_df.columns)
+
+    if source_names != target_names:
+        raise ValueError(
+            "Expected source and target to have same schema, got "
+            f"{consts.RESULT_TYPE_SOURCE}: {source_names}; {consts.RESULT_TYPE_TARGET}: {target_names}"
+        )
+
+
 def _pivot_result(
     result: "IbisTable",
     join_on_fields: tuple,
@@ -441,12 +528,19 @@ def _get_summary(
 ):
     """Logs a summary report/stats of row validation results."""
     try:
-        if result_df.loc[0, consts.VALIDATION_TYPE] == consts.ROW_VALIDATION:
-            # Vectorized calculations for all counts
+        if result_df.empty:
+            return
+
+        if (result_df.loc[0, consts.VALIDATION_TYPE] == consts.ROW_VALIDATION) or (
+            # Check for custom-query row validation which always should have primary keys (not null)
+            result_df.loc[0, consts.VALIDATION_TYPE] == consts.CUSTOM_QUERY
+            and result_df.loc[0, consts.CONFIG_PRIMARY_KEYS]
+        ):
+            # Vectorized calculations for all counts.
             success_condition = (
                 result_df[consts.VALIDATION_STATUS] == consts.VALIDATION_STATUS_SUCCESS
             )
-            fail_condition = ~success_condition  # Invert success for fail condition
+            fail_condition = ~success_condition  # Invert success for fail condition.
 
             source_not_in_target = (
                 result_df[consts.SOURCE_AGG_VALUE].notnull()
@@ -462,33 +556,32 @@ def _get_summary(
             )
 
             logging.info(
-                {
-                    consts.CONFIG_RUN_ID: run_metadata.run_id,
-                    consts.CONFIG_START_TIME: run_metadata.start_time.strftime(
-                        "%Y-%m-%d %H:%M:%S %Z"
-                    ),
-                    consts.CONFIG_END_TIME: run_metadata.end_time.strftime(
-                        "%Y-%m-%d %H:%M:%S %Z"
-                    ),
-                    consts.TOTAL_SOURCE_ROWS: source_df.shape[0],
-                    consts.TOTAL_TARGET_ROWS: target_df.shape[0],
-                    consts.TOTAL_ROWS_VALIDATED: result_df.shape[0],
-                    # Using .sum() on boolean Series for much faster counting
-                    consts.TOTAL_ROWS_SUCCESS: success_condition.sum(),
-                    consts.TOTAL_ROWS_FAIL: fail_condition.sum(),
-                    consts.FAILED_SOURCE_NOT_IN_TARGET: (
-                        fail_condition & source_not_in_target
-                    ).sum(),
-                    consts.FAILED_TARGET_NOT_IN_SOURCE: (
-                        fail_condition & target_not_in_source
-                    ).sum(),
-                    consts.FAILED_PRESENT_IN_BOTH_TABLES: (
-                        fail_condition & present_in_both_tables
-                    ).sum(),
-                }
+                json.dumps(
+                    {
+                        consts.CONFIG_RUN_ID: run_metadata.run_id,
+                        consts.CONFIG_START_TIME: run_metadata.start_time.isoformat(),
+                        consts.CONFIG_END_TIME: run_metadata.end_time.isoformat(),
+                        # Explicit conversion of numpy's int64 values to int for JSON serializability
+                        consts.TOTAL_SOURCE_ROWS: int(source_df.shape[0]),
+                        consts.TOTAL_TARGET_ROWS: int(target_df.shape[0]),
+                        consts.TOTAL_ROWS_VALIDATED: int(result_df.shape[0]),
+                        # Using .sum() on boolean Series for much faster counting
+                        consts.TOTAL_ROWS_SUCCESS: int(success_condition.sum()),
+                        consts.TOTAL_ROWS_FAIL: int(fail_condition.sum()),
+                        consts.FAILED_SOURCE_NOT_IN_TARGET: int(
+                            (fail_condition & source_not_in_target).sum()
+                        ),
+                        consts.FAILED_TARGET_NOT_IN_SOURCE: int(
+                            (fail_condition & target_not_in_source).sum()
+                        ),
+                        consts.FAILED_PRESENT_IN_BOTH_TABLES: int(
+                            (fail_condition & present_in_both_tables).sum()
+                        ),
+                    }
+                )
             )
     except Exception as e:
         logging.warning(
-            f"Error while generating summary report of row validation results: {e}",
+            f"{COMBINER_GET_SUMMARY_EXC_TEXT}: {e}",
             exc_info=True,
         )

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/data_validation/config_manager.py

@@ -16,7 +16,7 @@ import copy
 import logging
 import string
 import random
-from typing import TYPE_CHECKING, Optional, Union, List, Dict
+from typing import TYPE_CHECKING, Dict, List, Optional, Union, Tuple
 
 import google.oauth2.service_account
 import ibis.expr.datatypes as dt
@@ -66,6 +66,9 @@ class ConfigManager(object):
         if self.validation_type not in consts.CONFIG_TYPES:
             raise ValueError(f"Unknown Configuration Type: {self.validation_type}")
         self._comparison_max_col_length = None
+        # For some engines we need to know the actual raw data type rather than the Ibis canonical type.
+        self._source_raw_data_types = None
+        self._target_raw_data_types = None
 
     @property
     def config(self):
@@ -94,6 +97,50 @@ class ConfigManager(object):
 
         return self._target_conn
 
+    def get_source_raw_data_types(self) -> Dict[str, Tuple]:
+        """Return raw data type information from source system.
+
+        The raw data type is the source/target engine type, for example it might
+        be "NCLOB" or "char" when the Ibis type simply states "string".
+        The data is cached in state when fetched for the first time.
+        The retuen value is keyed on the casefolded column name and the tuple is
+        the remaining 6 elements of the DB API cursor description specification."""
+        if self._source_raw_data_types is None:
+            if hasattr(self.source_client, "raw_column_metadata"):
+                raw_data_types = self.source_client.raw_column_metadata(
+                    database=self.source_schema,
+                    table=self.source_table,
+                    query=self.source_query,
+                )
+                self._source_raw_data_types = {
+                    _[0].casefold(): _[1:] for _ in raw_data_types
+                }
+            else:
+                self._source_raw_data_types = {}
+        return self._source_raw_data_types
+
+    def get_target_raw_data_types(self) -> Dict[str, Tuple]:
+        """Return raw data type information from target system.
+
+        The raw data type is the source/target engine type, for example it might
+        be "NCLOB" or "char" when the Ibis type simply states "string".
+        The data is cached in state when fetched for the first time.
+        The retuen value is keyed on the casefolded column name and the tuple is
+        the remaining 6 elements of the DB API cursor description specification."""
+        if self._target_raw_data_types is None:
+            if hasattr(self.target_client, "raw_column_metadata"):
+                raw_data_types = self.target_client.raw_column_metadata(
+                    database=self.target_schema,
+                    table=self.target_table,
+                    query=self.target_query,
+                )
+                self._target_raw_data_types = {
+                    _[0].casefold(): _[1:] for _ in raw_data_types
+                }
+            else:
+                self._target_raw_data_types = {}
+        return self._target_raw_data_types
+
     def close_client_connections(self):
         """Attempt to clean up any source/target connections, based on the client types.
 
@@ -138,10 +185,6 @@ class ConfigManager(object):
         """Return if the validation should perform a case insensitive match."""
         return self._config.get(consts.CONFIG_CASE_INSENSITIVE_MATCH) or False
 
-    def process_in_memory(self):
-        """Return whether to process in memory or on a remote platform."""
-        return True
-
     @property
     def max_recursive_query_size(self):
         """Return Aggregates from Config"""

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/data_validation/consts.py

@@ -166,6 +166,7 @@ COUNT_STAR = "{count_star}"
 # Validation metadata
 RESULT_TYPE_SOURCE = "source"
 RESULT_TYPE_TARGET = "target"
+RESULT_TYPE_DIFFERENCES = "differences"
 
 FORMAT_TYPE_CSV = "csv"
 FORMAT_TYPE_JSON = "json"

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/data_validation/data_validation.py

@@ -115,9 +115,7 @@ class DataValidation(object):
                 "Schema validation", self.schema_validator.execute
             )
         else:
-            result_df = self._execute_validation(
-                self.validation_builder, process_in_memory=True
-            )
+            result_df = self._execute_validation(self.validation_builder)
 
         # Call Result Handler to Manage Results
         return self.result_handler.execute(result_df)
@@ -227,13 +225,10 @@ class DataValidation(object):
         clause recursively until the individual row differences can be
         identified.
         """
-        process_in_memory = self.config_manager.process_in_memory()
         past_results = []
         if len(grouped_fields) > 0:
             validation_builder.add_query_group(grouped_fields[0])
-            result_df = self._execute_validation(
-                validation_builder, process_in_memory=process_in_memory
-            )
+            result_df = self._execute_validation(validation_builder)
 
             for grouped_key in result_df[consts.GROUP_BY_COLUMNS].unique():
                 # Validations are viewed separtely, but queried together.
@@ -267,11 +262,7 @@ class DataValidation(object):
                         )
                     )
         elif self.config_manager.primary_keys and len(grouped_fields) == 0:
-            past_results.append(
-                self._execute_validation(
-                    validation_builder, process_in_memory=process_in_memory
-                )
-            )
+            past_results.append(self._execute_validation(validation_builder))
 
         # elif self.config_manager.primary_keys:
         #     validation_builder.add_config_query_groups(self.config_manager.primary_keys)
@@ -302,7 +293,7 @@ class DataValidation(object):
             }
             validation_builder.add_filter(filter_field)
 
-    def _execute_validation(self, validation_builder, process_in_memory=True):
+    def _execute_validation(self, validation_builder):
         """Execute Against a Supplied Validation Builder"""
 
         self.run_metadata.validations = validation_builder.get_metadata()
@@ -329,67 +320,48 @@ class DataValidation(object):
             )
         )
 
-        if process_in_memory:
-            futures = []
-            with ThreadPoolExecutor() as executor:
-                # Submit the two query network calls concurrently
-                futures.append(
-                    executor.submit(
-                        util.timed_call,
-                        "Source query",
-                        self.config_manager.source_client.execute,
-                        source_query,
-                    )
+        futures = []
+        with ThreadPoolExecutor() as executor:
+            # Submit the two query network calls concurrently
+            futures.append(
+                executor.submit(
+                    util.timed_call,
+                    "Source query",
+                    self.config_manager.source_client.execute,
+                    source_query,
                 )
-                futures.append(
-                    executor.submit(
-                        util.timed_call,
-                        "Target query",
-                        self.config_manager.target_client.execute,
-                        target_query,
-                    )
+            )
+            futures.append(
+                executor.submit(
+                    util.timed_call,
+                    "Target query",
+                    self.config_manager.target_client.execute,
+                    target_query,
                 )
-                source_df = futures[0].result()
-                target_df = futures[1].result()
-
-            pandas_client = ibis.pandas.connect(
-                {
-                    consts.RESULT_TYPE_SOURCE: source_df,
-                    consts.RESULT_TYPE_TARGET: target_df,
-                }
             )
+            source_df = futures[0].result()
+            target_df = futures[1].result()
 
-            try:
-                result_df = util.timed_call(
-                    "Generate report",
-                    combiner.generate_report,
-                    pandas_client,
-                    self.run_metadata,
-                    pandas_client.table(consts.RESULT_TYPE_SOURCE),
-                    pandas_client.table(consts.RESULT_TYPE_TARGET),
-                    join_on_fields=join_on_fields,
-                    is_value_comparison=is_value_comparison,
-                    verbose=self.verbose,
-                )
-            except Exception as e:
-                if self.verbose:
-                    logging.error("-- ** Logging Source DF ** --")
-                    logging.error(source_df.dtypes)
-                    logging.error(source_df)
-                    logging.error("-- ** Logging Target DF ** --")
-                    logging.error(target_df.dtypes)
-                    logging.error(target_df)
-                raise e
-        else:
-            result_df = combiner.generate_report(
-                self.config_manager.source_client,
+        try:
+            result_df = util.timed_call(
+                "Generate report",
+                combiner.generate_report,
                 self.run_metadata,
-                source_query,
-                target_query,
+                source_df,
+                target_df,
                 join_on_fields=join_on_fields,
                 is_value_comparison=is_value_comparison,
                 verbose=self.verbose,
             )
+        except Exception as e:
+            if self.verbose:
+                logging.error("-- ** Logging Source DF ** --")
+                logging.error(source_df.dtypes)
+                logging.error(source_df)
+                logging.error("-- ** Logging Target DF ** --")
+                logging.error(target_df.dtypes)
+                logging.error(target_df)
+            raise e
 
         return result_df
 

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/data_validation/partition_builder.py

@@ -17,6 +17,7 @@ import ibis
 import pandas
 import logging
 import re
+import datetime
 from typing import List, Dict, TYPE_CHECKING
 
 from data_validation import cli_tools, consts
@@ -253,12 +254,13 @@ class PartitionBuilder:
 
             def less_than_value(table, keys, values):
                 key_column = table.__getattr__(keys[0])
-                if key_column.type().is_date():
-                    # Ensure date PKs are treated as date literals as per #1191
-                    value = values[0].date()
-                else:
-                    value = values[0]
-
+                # Due to issue 1474, the type can be datetime.datetime or datetime.date
+                value = (
+                    values[0].date()
+                    if key_column.type().is_date()
+                    and isinstance(values[0], datetime.datetime)
+                    else values[0]
+                )
                 if len(keys) == 1:
                     return key_column < value
                 else:
@@ -269,10 +271,13 @@ class PartitionBuilder:
 
             def geq_value(table, keys, values):
                 key_column = table.__getattr__(keys[0])
-                if key_column.type().is_date():
-                    value = values[0].date()
-                else:
-                    value = values[0]
+                # Due to issue 1474, the type can be datetime.datetime or datetime.date
+                value = (
+                    values[0].date()
+                    if key_column.type().is_date()
+                    and isinstance(values[0], datetime.datetime)
+                    else values[0]
+                )
 
                 if len(keys) == 1:
                     return key_column >= value

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/google_pso_data_validator.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: google-pso-data-validator
-Version: 7.4.0
+Version: 7.5.0
 Summary: A package to enable easy data validation
 Author: PSO DVT Engineering team
 Author-email: data-validator-eng@google.com
@@ -51,6 +51,7 @@ Dynamic: author-email
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/setup.py

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

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/third_party/ibis/ibis_oracle/__init__.py

@@ -11,6 +11,9 @@
 # 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 Iterable, Literal, Tuple
+
 import sqlalchemy as sa
 from sqlalchemy.dialects.oracle.base import (
     OracleIdentifierPreparer,
@@ -20,7 +23,6 @@ from sqlalchemy.dialects.oracle.cx_oracle import OracleDialect_cx_oracle
 import re
 
 import ibis.expr.datatypes as dt
-from typing import Iterable, Literal, Tuple
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
 from third_party.ibis.ibis_oracle.compiler import OracleCompiler
 from third_party.ibis.ibis_oracle.datatypes import _get_type
@@ -178,3 +180,35 @@ class Backend(BaseAlchemyBackend):
                 list_pk_col_sql, parameters=(database.upper(), table.upper())
             )
             return [_[0] for _ in result.cursor.fetchall()]
+
+    def raw_column_metadata(
+        self, database: str = None, table: str = None, query: str = None
+    ) -> Iterable[Tuple]:
+        """Partner method to _metadata that retains raw data type information instead of converting to Ibis types.
+        This works in the same way as _metadata by running a query over the DVT source, either schema.table or a
+        custom query, and fetching the first row. From the cursor we can detect data types of the row's columns.
+
+        Returns:
+            list: A list of tuples containing the standard 7 DB API fields:
+                  https://peps.python.org/pep-0249/#description
+        """
+
+        def strip_prefix(s: str):
+            if s.startswith("DB_TYPE_"):
+                return s[8:]
+            else:
+                return s
+
+        assert (database and table) or query, "We should never receive all args=None"
+        if database and table:
+            source = f'"{database}"."{table}"'.upper()
+        elif query:
+            source = f"({query})"
+
+        with self.begin() as con:
+            result = con.exec_driver_sql(f"SELECT * FROM {source} t0 WHERE ROWNUM <= 1")
+            cursor = result.cursor
+            yield from (
+                (column[0], strip_prefix(column[1].name), *column[2:])
+                for column in cursor.description
+            )

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/third_party/ibis/ibis_postgres/client.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Literal
+from typing import Iterable, Literal, Tuple
 
 import ibis.expr.datatypes as dt
 import ibis.expr.schema as sch
@@ -156,7 +156,57 @@ def _list_primary_key_columns(self, database: str, table: str) -> list:
         return [_[0] for _ in result.cursor.fetchall()]
 
 
+def _raw_column_metadata(
+    self, database: str = None, table: str = None, query: str = None
+) -> Iterable[Tuple]:
+    """Partner method to _metadata that retains raw data type information instead of converting to Ibis types.
+    This works in the same way as _metadata by running a query over the DVT source, either schema.table or a
+    custom query, and fetching the first row. From the cursor we can detect data types of the row's columns.
+
+    Returns:
+        list: A list of tuples containing the standard 7 DB API fields:
+                https://peps.python.org/pep-0249/#description
+    """
+
+    assert (database and table) or query, "We should never receive all args=None"
+    if database and table:
+        source = f'"{database}"."{table}"'
+    elif query:
+        source = f"({query})"
+
+    with self.begin() as con:
+        cur = con.exec_driver_sql(f"SELECT * FROM {source} t0 LIMIT 0")
+        qry_cols = [
+            f"('{column.name}'::text,"
+            + f"{column.type_code},"
+            + f"{column.display_size if column.display_size else 'NULL'}::int,"
+            + f"{column.internal_size if column.internal_size else 'NULL'}::int,"
+            + f"{column.precision if column.precision else 'NULL'}::int,"
+            + f"{column.scale if column.scale is not None else 'NULL'}::int,"
+            + f"{column.null_ok if column.null_ok else 'NULL'}::int, {idx})"
+            for idx, column in enumerate(cur.cursor.description)
+        ]
+        sql = f"""SELECT name, format_type(t0.type_code, NULL) AS data_type
+                ,    display_size, internal_size, precision, scale, null_ok
+                FROM UNNEST(array[{','.join(qry_cols)}])
+                     AS t0(name text, type_code int, display_size int, internal_size int, precision int, scale int, null_ok int, col_order int)
+                ORDER BY col_order"""
+        yield from (
+            (
+                column.name,
+                column.data_type,
+                column.display_size,
+                column.internal_size,
+                column.precision,
+                column.scale,
+                column.null_ok,
+            )
+            for column in con.exec_driver_sql(sql)
+        )
+
+
 PostgresBackend._metadata = _metadata
 PostgresBackend.list_databases = list_schemas
 PostgresBackend.do_connect = do_connect
 PostgresBackend.list_primary_key_columns = _list_primary_key_columns
+PostgresBackend.raw_column_metadata = _raw_column_metadata

{google_pso_data_validator-7.4.0 → google_pso_data_validator-7.5.0}/third_party/ibis/ibis_teradata/__init__.py

@@ -11,6 +11,9 @@
 # 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 Iterable, Tuple
+
 import pandas
 import warnings
 
@@ -272,6 +275,33 @@ class Backend(BaseSQLBackend):
         )
         return list(tables_df.ColumnName.str.rstrip())
 
+    def raw_column_metadata(
+        self, database: str = None, table: str = None, query: str = None
+    ) -> Iterable[Tuple]:
+        """Partner method to _get_schema_using_query that retains raw data type information instead of converting
+        to Ibis types.  This works in the same way as _get_schema_using_query by running a query over the DVT
+        source, either schema.table or a custom query, and fetching the first row. From the cursor we can detect
+        data types of the row's columns.
+
+        Returns:
+            list: A list of tuples containing the standard 7 DB API fields:
+                  https://peps.python.org/pep-0249/#description
+        """
+
+        assert (database and table) or query, "We should never receive all args=None"
+        if database and table:
+            source = f'"{database}"."{table}"'.upper()
+        elif query:
+            source = f"({query})"
+
+        cur = self.raw_sql(f"SELECT TOP 1 * FROM {source} t0")
+        # resets the state of the cursor and closes operation
+        cur.fetchall()
+        yield from (
+            (column[0], raw_type, *column[2:])
+            for column, raw_type in zip(cur.description, cur.columntypename)
+        )
+
     # Methods we need to implement for BaseSQLBackend
     def create_table(self):
         pass