bigquery: QueryJob.result() waits forever if a jobs.getQueryResults response is lost (no default request timeout) #18857

"""Show that QueryJob.result() waits forever when one polling response is lost.

A local HTTP server mimics the BigQuery REST API. It never answers the first request
to the chosen endpoint, which is what a lost response looks like to the client.
No GCP access is needed.

    python repro.py <default|none|SECONDS> <jobs|queries>

    default | none | SECONDS   job.result() | job.result(timeout=None) | job.result(timeout=SECONDS)
    jobs | queries             silent endpoint: jobs.get | jobs.getQueryResults
    ALARM=<seconds>            give up waiting after this long (default 20)
"""
import json
import os
import signal
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import google.cloud.bigquery
from google.auth.credentials import AnonymousCredentials
from google.cloud import bigquery

timeout_arg, silent = sys.argv[1], sys.argv[2]
alarm = int(os.environ.get("ALARM", "20"))

REF = {"projectId": "p", "jobId": "job1", "location": "US"}
JOB = {
    "kind": "bigquery#job",
    "id": "p:US.job1",
    "jobReference": REF,
    "configuration": {"query": {"query": "SELECT 1", "useLegacySql": False}},
    "statistics": {},
}
requests_seen = {"jobs": 0, "queries": 0}
start = time.time()


class Handler(BaseHTTPRequestHandler):
    def log_message(self, *_):
        pass

    def _send(self, body):
        data = json.dumps(body).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_POST(self):  # jobs.insert
        self.rfile.read(int(self.headers.get("Content-Length", 0)))
        self._send({**JOB, "status": {"state": "RUNNING"}})

    def do_GET(self):
        kind = "queries" if "/queries/" in self.path else "jobs"
        requests_seen[kind] += 1
        path = self.path.split("?")[0].split("/v2")[-1]
        if kind == silent and requests_seen[kind] == 1:
            print(f"  t={time.time() - start:5.1f}s GET {path} -> no response", flush=True)
            time.sleep(3600)
            return
        print(f"  t={time.time() - start:5.1f}s GET {path} -> ok", flush=True)
        if kind == "queries":
            self._send({"kind": "bigquery#getQueryResultsResponse", "jobReference": REF,
                        "jobComplete": True, "totalRows": "0", "schema": {"fields": []}})
        else:
            self._send({**JOB, "status": {"state": "DONE"}})


server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
threading.Thread(target=server.serve_forever, daemon=True).start()

client = bigquery.Client(
    project="p", location="US", credentials=AnonymousCredentials(),
    client_options={"api_endpoint": f"http://127.0.0.1:{server.server_address[1]}"},
)
print(f"google-cloud-bigquery {google.cloud.bigquery.__version__}, "
      f"result({timeout_arg}), no response from {silent}", flush=True)


def give_up(*_):
    print(f"RESULT: still blocked after {alarm}s", flush=True)
    os._exit(0)


signal.signal(signal.SIGALRM, give_up)
signal.alarm(alarm)

try:
    job = client.query("SELECT 1")
    kwargs = {} if timeout_arg == "default" else {"timeout": None if timeout_arg == "none" else float(timeout_arg)}
    job.result(**kwargs)
    print(f"RESULT: completed after {time.time() - start:.1f}s", flush=True)
except Exception as exc:
    print(f"RESULT: {type(exc).__name__} after {time.time() - start:.1f}s: {str(exc)[:100]}", flush=True)
os._exit(0)