[YAML]: fix postcommit oracle bug and reorganize postcommit tests

.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": 1
+  "revision": 2
 }

.github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml

@@ -54,21 +54,22 @@ jobs:
       github.event_name == 'workflow_dispatch' ||
       github.event_name == 'pull_request_target' ||
       (github.event_name == 'schedule' && github.repository == 'apache/beam')
-    runs-on: [self-hosted, ubuntu-20.04, main]
+    runs-on: ubuntu-latest
     timeout-minutes: 100
     name: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
     strategy:
       matrix:
         job_name: ["beam_PostCommit_Yaml_Xlang_Direct"]
         job_phrase: ["Run Yaml_Xlang_Direct PostCommit"]
+        test_set: ["data", "databases", "messaging"]
     steps:
       - uses: actions/checkout@v4
       - name: Setup repository
         uses: ./.github/actions/setup-action
         with:
-          comment_phrase: ${{ matrix.job_phrase }}
+          comment_phrase: ${{ matrix.job_phrase }} ${{ matrix.test_set }}
           github_token: ${{ secrets.GITHUB_TOKEN }}
-          github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
+          github_job: ${{ matrix.job_name }} ${{ matrix.test_set }} (${{ matrix.job_phrase }} ${{ matrix.test_set }})
       - name: Setup environment
         uses: ./.github/actions/setup-environment-action
         with:
@@ -79,7 +80,7 @@ jobs:
       - name: run PostCommit Yaml Xlang Direct script
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
-          gradle-command: :sdks:python:postCommitYamlIntegrationTests -PbeamPythonExtra=ml_test
+          gradle-command: :sdks:python:postCommitYamlIntegrationTests -PyamlTestSet=${{ matrix.test_set }} -PbeamPythonExtra=ml_test
       - name: Archive Python Test Results
         uses: actions/upload-artifact@v4
         if: failure()
@@ -93,4 +94,5 @@ jobs:
           commit: '${{ env.prsha || env.GITHUB_SHA }}'
           comment_mode: ${{ github.event_name == 'issue_comment'  && 'always' || 'off' }}
           files: '**/pytest*.xml'
-          large_files: true
+          large_files: true
+

sdks/python/apache_beam/yaml/integration_tests.py

@@ -228,7 +228,7 @@ def temp_mysql_database():
                              with the MySQL database during setup.
       Exception: Any other exception encountered during the setup process.
   """
-  with MySqlContainer() as mysql_container:
+  with MySqlContainer(init=True) as mysql_container:
     try:
       # Make connection to temp database and create tmp table
       engine = sqlalchemy.create_engine(mysql_container.get_connection_url())
@@ -440,12 +440,12 @@ def temp_kafka_server():
       Exception: If there's an error starting the Kafka container or
                  interacting with the temporary Kafka server.
   """
-  try:
-    with KafkaContainer() as kafka_container:
+  with KafkaContainer() as kafka_container:
+    try:
       yield kafka_container.get_bootstrap_server()
-  except Exception as err:
-    logging.error("Error interacting with temporary Kakfa Server: %s", err)
-    raise err
+    except Exception as err:
+      logging.error("Error interacting with temporary Kakfa Server: %s", err)
+      raise err
 
 
 @contextlib.contextmanager

sdks/python/build.gradle

@@ -138,13 +138,13 @@ tasks.register("yamlIntegrationTests") {
   doLast {
     exec {
       executable 'sh'
-      args '-c', "${envdir}/bin/pytest -v apache_beam/yaml/integration_tests.py --deselect apache_beam/yaml/integration_tests.py::BigqueryTest::test_ReadFromBigQuery_ExternalJavaProvider_0"
+      args '-c', "${envdir}/bin/pytest -v apache_beam/yaml/integration_tests.py"
     }
   }
 }
 
 tasks.register("postCommitYamlIntegrationTests") {
-  description "Runs postcommit integration tests for yaml pipelines."
+  description "Runs postcommit integration tests for yaml pipelines - parameterized by yamlTestSet."
 
   dependsOn installGcpTest
   // Need to build all expansion services referenced in apache_beam/yaml/*.*
@@ -155,9 +155,28 @@ tasks.register("postCommitYamlIntegrationTests") {
   dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:build"
 
   doLast {
-    exec {
-      executable 'sh'
-      args '-c', "${envdir}/bin/pytest -v apache_beam/yaml/integration_tests.py --deselect apache_beam/yaml/integration_tests.py::BigqueryTest::test_ReadFromBigQuery_ExternalJavaProvider_0 --test_files_dir='extended_tests'"
+    def testSetInput = project.findProperty('yamlTestSet') ?: 'data,databases,messaging'
+    def testSetsToRun = testSetInput.tokenize(',').collect { it.trim() }.findAll { !it.isEmpty() }
+    testSetsToRun.each { currentTestSet ->
+      def test_files_dir
+
+      switch (currentTestSet) {
+        case 'data':
+          test_files_dir = 'extended_tests/data'
+          break
+        case 'databases':
+          test_files_dir = 'extended_tests/databases'
+          break
+        case 'messaging':
+          test_files_dir = 'extended_tests/messaging'
+          break
+        default:
+          throw StopExecutionException("Unknown yamlTestSet: ${testSet}. Must be one of 'data', 'databases', or 'messaging'.")
+      }
+      exec {
+        executable 'sh'
+        args '-c', "${envdir}/bin/pytest -v apache_beam/yaml/integration_tests.py --deselect apache_beam/yaml/integration_tests.py::BigqueryTest::test_ReadFromBigQuery_ExternalJavaProvider_0 --test_files_dir='${test_files_dir}'"
+      }
     }
   }
 }

sdks/python/setup.py

@@ -436,7 +436,8 @@ def get_portability_package_data():
               'virtualenv-clone>=0.5,<1.0',
               'mysql-connector-python>=9.3.0',
               'python-tds>=1.16.1',
-              'sqlalchemy-pytds>=1.0.2'
+              'sqlalchemy-pytds>=1.0.2',
+              'oracledb>=3.1.1'
           ],
           'gcp': [
               'cachetools>=3.1.0,<6',