[SPARK-53759][PYTHON][4.0] Fix missing close in simple-worker path

python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py

@@ -96,4 +96,7 @@ def process(df_id, batch_id):  # type: ignore[no-untyped-def]
     sock.settimeout(None)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/connect/streaming/worker/listener_worker.py

@@ -112,4 +112,7 @@ def process(listener_event_str, listener_event_type):  # type: ignore[no-untyped
     sock.settimeout(None)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/streaming/python_streaming_source_runner.py

@@ -211,4 +211,7 @@ def main(infile: IO, outfile: IO) -> None:
     sock.settimeout(None)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/streaming/transform_with_state_driver_worker.py

@@ -99,4 +99,7 @@ def process(
     (sock_file, sock) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/analyze_udtf.py

@@ -279,4 +279,7 @@ def invalid_analyze_result_field(field_name: str, expected_field: str) -> PySpar
     # TODO: Remove the following two lines and use `Process.pid()` when we drop JDK 8.
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/commit_data_source_write.py

@@ -124,4 +124,7 @@ def main(infile: IO, outfile: IO) -> None:
     (sock_file, _) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/create_data_source.py

@@ -189,4 +189,7 @@ def main(infile: IO, outfile: IO) -> None:
     (sock_file, _) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/lookup_data_sources.py

@@ -109,4 +109,7 @@ def main(infile: IO, outfile: IO) -> None:
     (sock_file, _) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/plan_data_source_read.py

@@ -382,4 +382,7 @@ def data_source_read_func(iterator: Iterable[pa.RecordBatch]) -> Iterable[pa.Rec
     (sock_file, _) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/python_streaming_sink_runner.py

@@ -153,4 +153,7 @@ def main(infile: IO, outfile: IO) -> None:
     (sock_file, _) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/sql/worker/write_into_data_source.py

@@ -260,4 +260,7 @@ def batch_to_rows() -> Iterator[Row]:
     (sock_file, _) = local_connect_and_auth(java_port, auth_secret)
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()

python/pyspark/worker.py

@@ -2081,4 +2081,7 @@ def process():
     # TODO: Remove the following two lines and use `Process.pid()` when we drop JDK 8.
     write_int(os.getpid(), sock_file)
     sock_file.flush()
-    main(sock_file, sock_file)
+    try:
+        main(sock_file, sock_file)
+    finally:
+        sock_file.close()