Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,8 @@ class GCSToBigQueryOperator(BaseOperator):
If true, the extra values are ignored. If false, records with extra columns
are treated as bad records, and if there are too many bad records, an
invalid error is returned in the job result.
The values are compared against the schema the load job runs with, so this has no effect
when that schema is inferred from the source data itself - see ``autodetect``.
:param allow_quoted_newlines: Whether to allow quoted newlines (true) or not (false).
:param allow_jagged_rows: Accept rows that are missing trailing optional columns.
The missing values are treated as nulls. If false, records with missing trailing
Expand Down Expand Up @@ -162,12 +164,10 @@ class GCSToBigQueryOperator(BaseOperator):
by one or more columns. BigQuery supports clustering for both partitioned and
non-partitioned tables. The order of columns given determines the sort order.
Not applicable for external tables.
:param autodetect: [Optional] Indicates if we should automatically infer the
options and schema for CSV and JSON sources. (Default: ``True``).
Parameter must be set to True if 'schema_fields' and 'schema_object' are undefined.
It is suggested to set to True if table are create outside of Airflow.
If autodetect is None and no schema is provided (neither via schema_fields
nor a schema_object), assume the table already exists.
:param autodetect: [Optional] Whether to infer the schema from the source data for CSV and
JSON sources. If ``True``, the schema is inferred from the source data. If ``False``,
either ``schema_fields`` or ``schema_object`` must be provided. If ``None``, no schema is
supplied and the existing destination table's schema is used. (Default: ``True``).
:param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).

.. code-block:: python
Expand Down Expand Up @@ -239,7 +239,7 @@ def __init__(
time_partitioning=None,
range_partitioning=None,
cluster_fields=None,
autodetect=True,
autodetect: bool | None = True,
encryption_configuration=None,
location=None,
impersonation_chain: str | Sequence[str] | None = None,
Expand Down Expand Up @@ -765,6 +765,20 @@ def _use_existing_table(self):
if self.extra_config:
self.configuration["load"].update(self.extra_config)

# Checked against the assembled load config rather than the operator attributes, because
# src_fmt_configs and extra_config can still override any of these three keys.
load_config = self.configuration["load"]
if (
load_config.get("autodetect")
and load_config.get("ignoreUnknownValues")
and "schema" not in load_config
):
self.log.warning(
"`ignore_unknown_values` has no effect when `autodetect=True` and no schema is "
"provided. Set `autodetect=None` to use the existing destination table's schema "
"instead."
)

return self.configuration

def _validate_src_fmt_configs(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2124,6 +2124,141 @@ def test_schema_fields_is_templated(self):

assert operator.schema_fields == SCHEMA_FIELDS

@pytest.mark.parametrize(
("operator_kwargs", "expects_warning"),
[
pytest.param({"ignore_unknown_values": True}, True, id="autodetect_defaults_to_true"),
pytest.param({"ignore_unknown_values": True, "autodetect": True}, True, id="autodetect_true"),
pytest.param({"ignore_unknown_values": True, "autodetect": None}, False, id="autodetect_none"),
pytest.param(
{"ignore_unknown_values": True, "autodetect": False, "schema_fields": SCHEMA_FIELDS},
False,
id="autodetect_false",
),
pytest.param(
{"ignore_unknown_values": True, "schema_fields": SCHEMA_FIELDS},
False,
id="schema_fields_supplied",
),
pytest.param({"ignore_unknown_values": False}, False, id="ignore_unknown_values_off"),
pytest.param(
{"ignore_unknown_values": True, "extra_config": {"autodetect": None}},
False,
id="autodetect_cleared_by_extra_config",
),
pytest.param(
{"ignore_unknown_values": True, "extra_config": {"schema": {"fields": SCHEMA_FIELDS}}},
False,
id="schema_supplied_by_extra_config",
),
pytest.param(
{"ignore_unknown_values": False, "extra_config": {"ignoreUnknownValues": True}},
True,
id="ignore_unknown_values_set_by_extra_config",
),
],
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since autodetect is explicitly a three-state parameter, could we also cover autodetect=False here?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've added a test case for autodetect=False, supplying the schema_fields as required.

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_ignore_unknown_values_no_op_warning(self, bq_hook, operator_kwargs, expects_warning):
bq_hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
bq_hook.return_value.generate_job_id.return_value = REAL_JOB_ID
bq_hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)

operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
write_disposition=WRITE_DISPOSITION,
project_id=JOB_PROJECT_ID,
**operator_kwargs,
)

with mock.patch.object(operator.log, "warning") as mock_warning:
operator.execute(context=MagicMock())

if expects_warning:
mock_warning.assert_called_once()
else:
mock_warning.assert_not_called()

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_ignore_unknown_values_no_op_warning_names_the_fix(self, bq_hook):
bq_hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
bq_hook.return_value.generate_job_id.return_value = REAL_JOB_ID
bq_hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)

operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
write_disposition=WRITE_DISPOSITION,
project_id=JOB_PROJECT_ID,
ignore_unknown_values=True,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do these parametrized cases need to assert the complete warning text? The behaviour under test seems to be whether the warning is emitted. Using assert_called_once() / assert_not_called() would make these cases less coupled to the wording, with the message itself asserted in a single test if we want to protect it.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, good suggestion. I've updated to use assert_called_once() / assert_not_called(). On protecting the message, rather than pinning the full text, one test (test_ignore_unknown_values_no_op_warning_names_the_fix) now asserts only that the warning names autodetect=None, since that's the actionable part.
One thing worth flagging - these assertions don't discriminate which warning fired, so they'll break if an unrelated warning is added. Filtering on/asserting against a stable part of the message (i.e. ignore_unknown_values or autodetect=None), rather than the full message, could avoid that. Let me know if you think it's worth it.


with mock.patch.object(operator.log, "warning") as mock_warning:
operator.execute(context=MagicMock())

mock_warning.assert_called_once()
assert "autodetect=None" in mock_warning.call_args.args[0]

@mock.patch(GCS_TO_BQ_PATH.format("GCSHook"))
@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_no_ignore_unknown_values_warning_with_schema_object(self, bq_hook, gcs_hook):
bq_hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
bq_hook.return_value.generate_job_id.return_value = REAL_JOB_ID
bq_hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)
gcs_hook.return_value.download.return_value = bytes(json.dumps(SCHEMA_FIELDS), "utf-8")

operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
schema_object_bucket=SCHEMA_BUCKET,
schema_object=SCHEMA_OBJECT,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
write_disposition=WRITE_DISPOSITION,
ignore_unknown_values=True,
project_id=JOB_PROJECT_ID,
)

with mock.patch.object(operator.log, "warning") as mock_warning:
operator.execute(context=MagicMock())

mock_warning.assert_not_called()

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_no_ignore_unknown_values_warning_for_external_table(self, bq_hook):
bq_hook.return_value.generate_job_id.return_value = REAL_JOB_ID
bq_hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)

operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
write_disposition=WRITE_DISPOSITION,
external_table=True,
ignore_unknown_values=True,
project_id=JOB_PROJECT_ID,
)

with mock.patch.object(operator.log, "warning") as mock_warning:
operator.execute(context=MagicMock())

mock_warning.assert_not_called()


@pytest.fixture
def create_task_instance(create_task_instance_of_operator, session):
Expand Down