From b8c9276be011d971b941b583fd3d4417d438067f Mon Sep 17 00:00:00 2001
From: Jared Chapman <thejaredchapman@gmail.com>
Date: Thu, 5 Oct 2023 14:03:46 -0500
Subject: [PATCH] docs: revised relax column mode sample (#1467)

* docs: Revised relax_column sample

* add todo for snippets.py cleanup

---------

Co-authored-by: Tim Swast <swast@google.com>
---
 docs/snippets.py                      |  2 ++
 samples/snippets/relax_column.py      | 52 +++++++++++++++++++++++++++
 samples/snippets/relax_column_test.py | 46 ++++++++++++++++++++++++
 3 files changed, 100 insertions(+)
 create mode 100644 samples/snippets/relax_column.py
 create mode 100644 samples/snippets/relax_column_test.py

diff --git a/docs/snippets.py b/docs/snippets.py
index 7f9b4f59e..62b0b6fd6 100644
--- a/docs/snippets.py
+++ b/docs/snippets.py
@@ -265,6 +265,8 @@ def test_relax_column(client, to_delete):
     dataset = client.create_dataset(dataset)
     to_delete.append(dataset)
 
+    # TODO(tswast): remove code sample once references to it on
+    # cloud.google.com are updated to samples/snippets/relax_column.py
     # [START bigquery_relax_column]
     # from google.cloud import bigquery
     # client = bigquery.Client()
diff --git a/samples/snippets/relax_column.py b/samples/snippets/relax_column.py
new file mode 100644
index 000000000..bcd79cee8
--- /dev/null
+++ b/samples/snippets/relax_column.py
@@ -0,0 +1,52 @@
+# Copyright 2022 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
+#
+#     https://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 google.cloud import bigquery
+
+
+def relax_column(table_id: str) -> bigquery.Table:
+    orig_table_id = table_id
+
+    # [START bigquery_relax_column]
+    from google.cloud import bigquery
+
+    client = bigquery.Client()
+
+    # TODO(dev): Change table_id to full name of the table you want to create.
+    table_id = "your-project.your_dataset.your_table"
+
+    # [END bigquery_relax_column]
+    table_id = orig_table_id
+
+    # [START bigquery_relax_column]
+    table = client.get_table(table_id)
+    new_schema = []
+    for field in table.schema:
+        if field.mode != "REQUIRED":
+            new_schema.append(field)
+        else:
+            # SchemaField properties cannot be edited after initialization.
+            # To make changes, construct new SchemaField objects.
+            new_field = field.to_api_repr()
+            new_field["mode"] = "NULLABLE"
+            relaxed_field = bigquery.SchemaField.from_api_repr(new_field)
+            new_schema.append(relaxed_field)
+
+    table.schema = new_schema
+    table = client.update_table(table, ["schema"])
+
+    print(f"Updated {table_id} schema: {table.schema}.")
+
+    # [END bigquery_relax_column]
+    return table
diff --git a/samples/snippets/relax_column_test.py b/samples/snippets/relax_column_test.py
new file mode 100644
index 000000000..b40b13fa1
--- /dev/null
+++ b/samples/snippets/relax_column_test.py
@@ -0,0 +1,46 @@
+# Copyright 2022 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
+#
+#     https://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 typing
+
+from google.cloud import bigquery
+
+import relax_column
+
+if typing.TYPE_CHECKING:
+    import pytest
+
+
+def test_relax_column(
+    capsys: "pytest.CaptureFixture[str]",
+    bigquery_client: bigquery.Client,
+    random_table_id: str,
+) -> None:
+    table = bigquery.Table(
+        random_table_id,
+        schema=[
+            bigquery.SchemaField("string_col", "STRING", mode="NULLABLE"),
+            bigquery.SchemaField("string_col2", "STRING", mode="REQUIRED"),
+        ],
+    )
+
+    bigquery_client.create_table(table)
+    table = relax_column.relax_column(random_table_id)
+
+    out, _ = capsys.readouterr()
+
+    assert all(field.mode == "NULLABLE" for field in table.schema)
+    assert "REQUIRED" not in out
+    assert "NULLABLE" in out
+    assert random_table_id in out