diff --git a/tests/system/data_sources/test_hive.py b/tests/system/data_sources/test_hive.py index eb9ea737..dd2f876b 100644 --- a/tests/system/data_sources/test_hive.py +++ b/tests/system/data_sources/test_hive.py @@ -339,3 +339,28 @@ def test_row_validation_binary_pk_to_bigquery(): ) df = run_test_from_cli_args(args) binary_key_assertions(df) + + +@mock.patch( + "data_validation.state_manager.StateManager.get_connection_config", + new=mock_get_connection_config, +) +def test_custom_query_validation_core_types(): + """Hive to Hive dvt_core_types custom-query validation""" + parser = cli_tools.configure_arg_parser() + args = parser.parse_args( + [ + "validate", + "custom-query", + "column", + "-sc=mock-conn", + "-tc=mock-conn", + "--source-query=select * from pso_data_validator.dvt_core_types", + "--target-query=select * from pso_data_validator.dvt_core_types", + "--filter-status=fail", + "--count=*", + ] + ) + df = run_test_from_cli_args(args) + # With filter on failures the data frame should be empty + assert len(df) == 0 diff --git a/third_party/ibis/ibis_impala/api.py b/third_party/ibis/ibis_impala/api.py index 00d18596..3f5783e2 100644 --- a/third_party/ibis/ibis_impala/api.py +++ b/third_party/ibis/ibis_impala/api.py @@ -197,6 +197,10 @@ def _get_schema_using_query(self, query): # Removing LIMIT 0 around query since it returns no results in Hive cur = self.raw_sql(f"SELECT * FROM ({query}) t0 LIMIT 1") cur.fetchall() + cur.description = [ + (description[0].replace("t0.", "", 1), *description[1:]) + for description in cur.description + ] ibis_fields = self._adapt_types(cur.description) cur.release()