pyspark-client-wasm/pyspark_connect_web/jupyterlite/harness.html at main · HyukjinKwon/pyspark-client-wasm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
<!DOCTYPE html>
<!-- SPDX-License-Identifier: Apache-2.0 -->
<!--
  harness.html - a minimal standalone page that boots the pyspark-connect-web
  browser stack WITHOUT JupyterLite, for the headless e2e (tests/e2e):

    Pyodide (Web Worker) + micropip(pyspark-client==4.1.2 + the wheel)
      -> pcw.install()
      -> SparkSession.builder.remote("sc://.../;transport=grpcweb")
      -> SAB/Atomics blocking bridge (bridge.js on this main thread does the
         real cross-origin fetch to Envoy's grpc-web endpoint)

  window.__pcwRunPython(src) is exposed SYNCHRONOUSLY on load (so the e2e's
  bridgeAvailable() check passes immediately); each call internally awaits the
  worker boot + `spark` binding before dispatching. The page MUST be served
  cross-origin isolated (Envoy sets COOP/COEP on :8000) or SharedArrayBuffer is
  unavailable.
-->
<html lang="en">
  <head>
    <meta charset="utf-8" />
    <title>pyspark-connect-web e2e harness</title>
    <script type="module">
      import { installBridge } from "/worker/bridge.js";

      const SPARK_REMOTE =
        new URLSearchParams(location.search).get("remote") ||
        "sc://localhost:8081/;transport=grpcweb";

      const status = (s) => {
        const el = document.getElementById("status");
        if (el) el.textContent = s;
        // eslint-disable-next-line no-console
        console.log("[pcw harness]", s);
      };

      // MODULE worker: recent Pyodide refuses classic workers, and
      // worker_bootstrap.js loads Pyodide via dynamic import (ESM).
      const worker = new Worker("/worker/worker_bootstrap.js", {
        type: "module",
      });
      installBridge(worker); // attaches the SAB on pcw_sab, drives pcw_rpc

      // pcw_run request/response correlation.
      let seq = 0;
      const pending = new Map();
      worker.addEventListener("message", (ev) => {
        const m = ev.data || {};
        if (m.type === "pcw_result" && pending.has(m.id)) {
          pending.get(m.id).resolve(m.result);
          pending.delete(m.id);
        } else if (m.type === "pcw_run_error" && pending.has(m.id)) {
          pending.get(m.id).reject(new Error(m.message));
          pending.delete(m.id);
        }
      });
      function run(code) {
        const id = ++seq;
        return new Promise((resolve, reject) => {
          pending.set(id, { resolve, reject });
          worker.postMessage({ type: "pcw_run", id, code });
        });
      }

      // Readiness gate: resolved once Pyodide booted and `spark` is bound.
      let resolveReady, rejectReady;
      const readyPromise = new Promise((res, rej) => {
        resolveReady = res;
        rejectReady = rej;
      });

      // Exposed IMMEDIATELY (synchronously, on module eval) so the e2e's
      // bridgeAvailable() sees it; calls queue until the kernel is ready.
      window.__pcwRunPython = async (src) => {
        await readyPromise;
        return run(src);
      };

      const ready = new Promise((resolve, reject) => {
        worker.addEventListener("message", (ev) => {
          const m = ev.data || {};
          if (m.type === "pcw_ready") resolve();
          else if (m.type === "pcw_error") reject(new Error(m.message));
        });
      });

      async function boot() {
        if (self.crossOriginIsolated !== true) {
          throw new Error("not cross-origin isolated (need COOP/COEP)");
        }
        status("starting worker + Pyodide ...");
        worker.postMessage({ type: "pcw_boot" });
        await ready;
        status("binding SparkSession (" + SPARK_REMOTE + ") ...");
        await run(
          [
            "import pyspark_connect_web as pcw",
            "pcw.install()",
            "from pyspark.sql import SparkSession",
            "spark = SparkSession.builder.remote(" +
              JSON.stringify(SPARK_REMOTE) +
              ").getOrCreate()",
            "None",
          ].join("\n"),
        );
        status("ready: window.__pcwRunPython is live, spark is bound");
        resolveReady();
      }

      boot().catch((e) => {
        window.__pcwError = String((e && e.message) || e);
        status("FAILED: " + window.__pcwError);
        rejectReady(e);
        // eslint-disable-next-line no-console
        console.error("[pcw harness] failed:", e);
      });
    </script>
  </head>
  <body>
    <h1>pyspark-connect-web e2e harness</h1>
    <p id="status">booting...</p>
  </body>
</html>