fix: Fix ValueError in signal handling for Trino worker threads (#6428)

mbm-codes · jyejare · web-flow · commit 506d919f3aaa · 2026-07-01T16:39:59.000+05:30
* fix: fix signal.signal() ValueError in Trino worker threads

Signed-off-by: MBM &lt;mihir105@gmail.com&gt;

* test: patch signal.signal in main thread test to avoid handler leak

Signed-off-by: MBM &lt;mihir105@gmail.com&gt;

* Update sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/test_trino_queries.py

Apply reviewer suggestion: verify correct signals and handler in test

Co-authored-by: Jitendra Yejare &lt;jyejare@redhat.com&gt;
Signed-off-by: MBM &lt;mihir105@gmail.com&gt;

* test: address PR feedback for worker thread errors and signals

- Simplify worker thread assertion message to safely print errors list

- Add verification for signal handler registration details

Signed-off-by: MBM &lt;mihir105@gmail.com&gt;

---------

Signed-off-by: MBM &lt;mihir105@gmail.com&gt;
Co-authored-by: Jitendra Yejare &lt;jyejare@redhat.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/test_trino_queries.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/test_trino_queries.py
@@ -0,0 +1,48 @@
+import signal
+import threading
+from unittest.mock import MagicMock, patch
+
+from feast.infra.offline_stores.contrib.trino_offline_store.trino_queries import (
+    Query,
+)
+
+
+def test_query_init_in_main_thread_registers_signals():
+    """signal.signal() should work fine in main thread."""
+
+    # Should not raise any exception in main thread
+    cursor = MagicMock()
+
+    with patch("signal.signal") as mock_signal:
+        query = Query(query_text="SELECT 1", cursor=cursor)
+        assert query.query_text == "SELECT 1"
+
+        # Verify signal handlers are registered correctly
+        mock_signal.assert_any_call(signal.SIGINT, query.cancel)
+        mock_signal.assert_any_call(signal.SIGTERM, query.cancel)
+
+        # Expected signal.signal to be called twice for SIGINT and SIGTERM
+        assert mock_signal.call_count == 2
+
+
+def test_query_init_in_worker_thread_does_not_raise():
+    """Regression test: signal.signal() fails in non-main threads."""
+    # signal.signal() raises ValueError when called outside the main thread.
+    # This test verifies the fix guards against that by running Query.__init__
+    # in a worker thread and ensuring no exception is raised.
+
+    errors = []
+    cursor = MagicMock()
+
+    def create_query():
+        try:
+            query = Query(query_text="SELECT 1", cursor=cursor)
+            assert query.query_text == "SELECT 1"
+        except ValueError as e:
+            errors.append(e)
+
+    thread = threading.Thread(target=create_query)
+    thread.start()
+    thread.join()
+
+    assert not errors, f"Unexpected ValueError in worker thread: {errors}"
diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import signal
+import threading
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Dict, List, Optional
@@ -92,8 +93,9 @@ def __init__(self, query_text: str, cursor: Cursor):
         self.status = QueryStatus.PENDING
         self._cursor = cursor
 
-        signal.signal(signal.SIGINT, self.cancel)
-        signal.signal(signal.SIGTERM, self.cancel)
+        if threading.current_thread() is threading.main_thread():
+            signal.signal(signal.SIGINT, self.cancel)
+            signal.signal(signal.SIGTERM, self.cancel)
 
     def execute(self) -> Results:
         try:
