Add integration test for heartbeat

Author: azawlocki-sbdt

etc/config.properties

@@ -16,5 +16,6 @@ protocol.spooling.retrieval-mode=coordinator_proxy
 # Enable dynamic catalog management
 catalog.management=dynamic
 
-# Disable http request log
-http-server.log.enabled=false
+# Enable HTTP request log as it's grepped in integration tests for the heartbeat mechanism
+http-server.log.enabled=true
+http-server.log.immediate-flush=true

tests/development_server.py

@@ -16,6 +16,15 @@
 TRINO_VERSION = os.environ.get("TRINO_VERSION") or "latest"
 TRINO_HOST = "localhost"
 
+_trino_container: "DockerContainer | None" = None
+
+
+def get_trino_container() -> "DockerContainer | None":
+    """Return the Trino DockerContainer started by start_development_server(), or None
+    if Trino was not started by this module (e.g. TRINO_RUNNING_HOST points to an
+    existing server)."""
+    return _trino_container
+
 
 def create_bucket(s3_client):
     bucket_name = "spooling"
@@ -99,7 +108,12 @@ def start_development_server(port=None, trino_version=TRINO_VERSION):
         # Otherwise some tests fail with No nodes available
         time.sleep(2)
 
-        yield localstack, trino, network
+        global _trino_container
+        _trino_container = trino
+        try:
+            yield localstack, trino, network
+        finally:
+            _trino_container = None
     finally:
         # Stop containers when exiting the context
         if trino:

tests/integration/test_dbapi_integration.py

@@ -27,6 +27,7 @@
 from tzlocal import get_localzone_name  # type: ignore
 
 import trino
+from tests.development_server import get_trino_container
 from tests.integration.conftest import trino_version
 from trino import constants
 from trino.client import InlineSegment
@@ -1896,7 +1897,7 @@ def test_segments_cursor(trino_connection):
             assert isinstance(segment.segment.ack_uri, str), (
                 f"Expected string for ack_uri, got {type(segment.segment.ack_uri)}"
             )
-        total += len(list(SegmentIterator(segment, row_mapper)))
+        total += len(list(SegmentIterator(segment, row_mapper, cur._query._request)))
     assert total == 300875, f"Expected total rows 300875, got {total}"
 
 
@@ -2000,6 +2001,61 @@ def test_spooled_segments_lazy_description(trino_connection):
     assert len(cur.fetchall()) == 60175
 
 
+@pytest.mark.skipif(
+    trino_version() <= 466,
+    reason="spooling protocol was introduced in version 466"
+)
+def test_heartbeat_head_requests_during_spooled_download(run_trino):
+    """Verify that heartbeat HEAD requests are sent to the coordinator while
+    downloading spooled segments from external storage."""
+    trino_container = get_trino_container()
+    if trino_container is None:
+        pytest.skip("cannot read Trino HTTP request log (Trino not started by this test session)")
+
+    host, port = run_trino
+    conn = trino.dbapi.Connection(
+        host=host, port=port, user="test", source="test",
+        max_attempts=1, encoding="json", heartbeat_interval=0.1,
+    )
+
+    container = trino_container.get_wrapped_container()
+    log_path = "/data/trino/var/log/http-request.log"
+
+    # Capture the current size of the HTTP request log
+    exit_code, output = container.exec_run(["wc", "-l", log_path])
+    if exit_code != 0:
+        pytest.skip("cannot read Trino HTTP request log (log disabled or path changed)")
+    logfile_lines = int(output.decode().split()[0])
+
+    cur = conn.cursor()
+    cur.execute("""SELECT l.*
+    FROM tpch.tiny.lineitem l, TABLE(sequence(
+        start => 1,
+        stop => 5,
+        step => 1)) n""")
+    cur.fetchall()
+    cur.close()
+
+    for try_ in range(10):
+        if try_:
+            # Sometimes trino needs time to flush the logs
+            t.sleep(1.0)
+
+        _, output = container.exec_run(["tail", "-n", f"+{logfile_lines}", log_path])
+        loglines = output.decode().splitlines()
+        head_requests = [
+            line for line in loglines
+            if "HEAD" in line and "/v1/statement/executing/" in line
+        ]
+        if head_requests:
+            break
+
+    assert head_requests, (
+        "Expected heartbeat HEAD requests in http-request.log but found none.\n"
+        f"Log tail:\n{''.join(loglines)}"
+    )
+
+
 def get_cursor(legacy_prepared_statements, run_trino):
     host, port = run_trino