fix post commit yaml test (#39142)

* try to fix iceberg post commit test

* trigger file

* isort

* oops more lint

* lint

* fix tmp dir

* add timeouts per gemini comments

Author: derrickaw

.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "revision": 6
+  "revision": 7
 }

sdks/python/apache_beam/yaml/extended_tests/databases/iceberg.yaml

@@ -16,8 +16,18 @@
 #
 
 fixtures:
-  - name: TEMP_DIR
-    type: "tempfile.TemporaryDirectory"
+  - name: ICEBERG_FIXTURE
+    type: "apache_beam.yaml.integration_tests.temp_iceberg_table_with_pk"
+    config:
+      table_data:
+        name: "labels"
+        schema:
+          type: "struct"
+          schema-id: 0
+          fields:
+            - { id: 1, name: "label", required: true, type: "string" }
+            - { id: 2, name: "rank", required: true, type: "long" }
+          identifier-field-ids: [1]
 
 pipelines:
   - name: write
@@ -32,26 +42,26 @@ pipelines:
               - {label: "389a", rank: 2}
         - type: WriteToIceberg
           config:
-            table: db.labels
-            catalog_name: hadoop_catalog
+            table: "{ICEBERG_FIXTURE[table]}"
+            catalog_name: rest_catalog
             catalog_properties:
-              type: hadoop
-              warehouse: "{TEMP_DIR}"
+              type: rest
+              uri: "{ICEBERG_FIXTURE[api_url]}"
     options:
       project: "apache-beam-testing"
-      temp_location: "{TEMP_DIR}"
+      temp_location: "{ICEBERG_FIXTURE[temp_dir]}"
 
   - name: read
     pipeline:
       type: chain
       transforms:
         - type: ReadFromIceberg
           config:
-            table: db.labels
-            catalog_name: hadoop_catalog
+            table: "{ICEBERG_FIXTURE[table]}"
+            catalog_name: rest_catalog
             catalog_properties:
-              type: hadoop
-              warehouse: "{TEMP_DIR}"
+              type: rest
+              uri: "{ICEBERG_FIXTURE[api_url]}"
         - type: AssertEqual
           config:
             elements:
@@ -60,19 +70,19 @@ pipelines:
               - {label: "389a", rank: 2}
     options:
       project: "apache-beam-testing"
-      temp_location: "{TEMP_DIR}"
+      temp_location: "{ICEBERG_FIXTURE[temp_dir]}"
 
   - name: read_cdc_batch
     pipeline:
       type: chain
       transforms:
         - type: ReadFromIcebergCDC
           config:
-            table: db.labels
-            catalog_name: hadoop_catalog
+            table: "{ICEBERG_FIXTURE[table]}"
+            catalog_name: rest_catalog
             catalog_properties:
-              type: hadoop
-              warehouse: "{TEMP_DIR}"
+              type: rest
+              uri: "{ICEBERG_FIXTURE[api_url]}"
             from_timestamp: 1762819200000
             to_timestamp:  2078352000000
             filter: '"label" = ''11a'' or "rank" = 1'
@@ -86,19 +96,19 @@ pipelines:
               - {label: "37a", rank: 1}
     options:
       project: "apache-beam-testing"
-      temp_location: "{TEMP_DIR}"
+      temp_location: "{ICEBERG_FIXTURE[temp_dir]}"
 
   - name: read_cdc_streaming
     pipeline:
       type: chain
       transforms:
         - type: ReadFromIcebergCDC
           config:
-            table: db.labels
-            catalog_name: hadoop_catalog
+            table: "{ICEBERG_FIXTURE[table]}"
+            catalog_name: rest_catalog
             catalog_properties:
-              type: hadoop
-              warehouse: "{TEMP_DIR}"
+              type: rest
+              uri: "{ICEBERG_FIXTURE[api_url]}"
             streaming: True
             to_timestamp: 2078352000000
             filter: '"label" = ''11a'' or "rank" = 1'
@@ -112,4 +122,4 @@ pipelines:
               - {label: "37a", rank: 1}
     options:
       project: "apache-beam-testing"
-      temp_location: "{TEMP_DIR}"
+      temp_location: "{ICEBERG_FIXTURE[temp_dir]}"

sdks/python/apache_beam/yaml/integration_tests.py

@@ -26,15 +26,20 @@
 import os
 import random
 import secrets
+import shutil
 import sqlite3
 import string
 import struct
+import tempfile
+import time
 import unittest
 import uuid
 from datetime import datetime
 from datetime import timezone
 
 import mock
+import requests
+from testcontainers.core.container import DockerContainer
 
 from apache_beam.coders import Coder
 from apache_beam.coders.coder_impl import CoderImpl
@@ -556,6 +561,75 @@ def temp_oracle_database():
     yield f"jdbc:oracle:thin:system/oracle@localhost:{port}/XEPDB1"
 
 
+@contextlib.contextmanager
+def temp_iceberg_table_with_pk(table_data):
+
+  # Create a temp dir that will be shared between host and container.
+  # We use the exact same path on both to avoid path mapping issues.
+  # We create it in the current working directory (workspace) because
+  # Docker in GitHub Actions often cannot mount directories from /tmp.
+  temp_dir = tempfile.mkdtemp(dir=os.getcwd())
+  os.chmod(temp_dir, 0o777)
+
+  # Start the Iceberg REST catalog container
+  container = DockerContainer("tabulario/iceberg-rest:0.6.0")
+  container.with_exposed_ports(8181)
+  container.with_volume_mapping(temp_dir, temp_dir, mode='rw')
+  container.with_env("HADOOP_USER_NAME", "iceberg")
+  container.with_env("CATALOG_WAREHOUSE", temp_dir)
+  container.with_env(
+      "CATALOG_IO__IMPL", "org.apache.iceberg.hadoop.HadoopFileIO")
+
+  try:
+    container.start()
+
+    ip = container.get_container_host_ip()
+    port = container.get_exposed_port(8181)
+    api_url = f"http://{ip}:{port}"
+
+    # Poll the REST API until it is ready
+    for _ in range(30):
+      try:
+        response = requests.get(f"{api_url}/v1/config", timeout=5)
+        if response.status_code == 200:
+          break
+      except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
+        pass
+      time.sleep(1)
+    else:
+      raise RuntimeError("Iceberg REST catalog failed to start in time.")
+
+    # Create namespace 'db'
+    requests.post(
+        f"{api_url}/v1/namespaces",
+        json={"namespace": ["db"]},
+        headers={"Content-Type": "application/json"},
+        timeout=10)
+
+    # Create table with primary key
+    response = requests.post(
+        f"{api_url}/v1/namespaces/db/tables",
+        json=table_data,
+        headers={"Content-Type": "application/json"},
+        timeout=10)
+    if response.status_code != 200:
+      raise RuntimeError(f"Failed to create Iceberg table: {response.text}")
+
+    # Change permissions of the created directories inside the container
+    # so the host user can write to them.
+    container.get_wrapped_container().exec_run(f"chmod -R 777 {temp_dir}")
+
+    yield {
+        "api_url": api_url,
+        "temp_dir": temp_dir,
+        "table": f"db.{table_data['name']}"
+    }
+
+  finally:
+    container.stop()
+    shutil.rmtree(temp_dir, ignore_errors=True)
+
+
 @contextlib.contextmanager
 def temp_kafka_server():
   """Context manager to provide a temporary Kafka server for testing.