diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml index 816d9f7be90..d647afabf81 100644 --- a/.github/workflows/core.yml +++ b/.github/workflows/core.yml @@ -92,7 +92,7 @@ jobs: matrix: java: [ 11 ] env: - INTERPRETERS: 'hbase,jdbc,file,flink-cmd,cassandra,elasticsearch,bigquery,livy,groovy,java,neo4j,sparql,mongodb,influxdb,shell' + INTERPRETERS: 'hbase,jdbc,file,flink-cmd,cassandra,elasticsearch,bigquery,livy,groovy,java,neo4j,sparql,mongodb,influxdb,shell,spark-connect' steps: - name: Checkout uses: actions/checkout@v5 diff --git a/pom.xml b/pom.xml index a8e8a7f8f88..e9e2ee8669e 100644 --- a/pom.xml +++ b/pom.xml @@ -60,6 +60,7 @@ zeppelin-jupyter-interpreter-shaded groovy spark + spark-connect spark-submit markdown mongodb diff --git a/spark-connect/README.md b/spark-connect/README.md new file mode 100644 index 00000000000..afd4a6fb6e9 --- /dev/null +++ b/spark-connect/README.md @@ -0,0 +1,255 @@ +# Spark Connect Interpreter for Apache Zeppelin + +## What is Spark Connect? + +Spark Connect (Spark 3.5+) is a new client-server architecture for Apache Spark that decouples the Spark client from the Spark cluster. Unlike the traditional `spark` interpreter which requires running Spark in the same JVM process, Spark Connect is a **thin gRPC client** that communicates with a remote Spark cluster via the `sc://host:port` connection string. + +## Why Use Spark Connect? + +- **No local Spark installation** — Zeppelin doesn't need the full Spark distribution on its host +- **Remote cluster support** — Connect to any Spark 3.5+ cluster over the network +- **Token authentication** — Support for token-based auth and SSL +- **Multi-user isolation** — Per-user session quotas prevent resource exhaustion +- **Cleaner deployments** — Simpler Docker images, reduced memory footprint + +## Differences from the Legacy Spark Interpreter + +| Feature | Spark Interpreter | Spark Connect Interpreter | +|---------|-------------------|---------------------------| +| **Architecture** | In-process SparkContext | Remote gRPC client | +| **Installation** | Requires full Spark on host | Only Spark Connect client JAR needed | +| **Scala support** | Yes (via embedded Scala interpreter) | No (client-only protocol) | +| **R support** | Yes (SparkR) | No (not supported) | +| **ZeppelinContext** | Full support (`z.show()`, Angular) | Returns `null` | +| **Multi-user** | Global shared SparkContext | Isolated sessions per user | +| **Session quota** | Not enforced | Per-user quota (default: 5) | + +## Prerequisites + +1. **Spark 3.5.x cluster** running the Spark Connect server + ```bash + # Start a Spark Connect server on port 15002 + spark-shell --master --conf spark.connect.grpc.binding.port=15002 + ``` + +2. **Python 3.x** (for PySpark/IPySpark support) + +## Configuration Properties + +The Spark Connect interpreter supports the following configuration properties in the Zeppelin UI: + +### Connection Settings + +| Property | Type | Default | Description | +|----------|------|---------|-------------| +| `zeppelin.spark.connect.url` | string | `sc://localhost:15002` | Spark Connect server URL | +| `zeppelin.spark.connect.token` | string | `` | Optional token for authentication (redacted in logs) | +| `zeppelin.spark.connect.use_ssl` | checkbox | false | Enable SSL/TLS for connection | +| `zeppelin.spark.connect.user_id` | string | `` | User ID to report to Spark Connect server | + +### Session Management + +| Property | Type | Default | Description | +|----------|------|---------|-------------| +| `zeppelin.spark.connect.maxSessionsPerUser` | number | 5 | Maximum concurrent sessions per user | + +### Execution + +| Property | Type | Default | Description | +|----------|------|---------|-------------| +| `zeppelin.spark.maxResult` | number | 10000 | Maximum result rows to fetch | +| `zeppelin.spark.concurrentSQL` | checkbox | false | Allow concurrent SQL execution (within notebook) | +| `zeppelin.spark.concurrentSQL.max` | number | 10 | Max concurrent SQL threads | + +### PySpark + +| Property | Type | Default | Description | +|----------|------|---------|-------------| +| `zeppelin.python` | string | `python` | Python executable path | +| `zeppelin.pyspark.useIPython` | checkbox | true | Use IPython if available | + +## Usage + +### SQL Mode (Default) + +```sql +%spark-connect + +SELECT * FROM my_table LIMIT 10 +``` + +### SQL with Concurrent Execution + +```sql +%spark-connect.sql + +-- This uses the concurrentSQL scheduler +SELECT COUNT(*) FROM large_table +``` + +### PySpark + +```python +%spark-connect.pyspark + +df = spark.sql("SELECT * FROM my_table") +df.show() +``` + +### IPython PySpark + +```python +%spark-connect.ipyspark + +# Full IPython REPL with Spark +df = spark.sql("SELECT COUNT(*) FROM table") +df.collect() +``` + +## Architecture + +### SparkConnectInterpreter +- Core interpreter managing the remote Spark session +- Enforces per-user session quota via `ConcurrentHashMap` +- Uses `NotebookLockManager` for per-notebook sequential execution +- Delegates SQL parsing to `SqlSplitter` for multi-statement support + +### SparkConnectSqlInterpreter +- SQL-only frontend with optional concurrent scheduler +- Shares the Spark session with `SparkConnectInterpreter` + +### PySparkConnectInterpreter +- Uses PySpark's native Spark Connect client (`SparkSession.builder.remote(...)`) +- Forwards the connection URI to Python via the `SPARK_REMOTE` env var; no Py4j bridge to the Java SparkSession +- Python opens an **independent** Spark Connect session against the same gRPC server as the Java/SQL interpreter — cross-language sharing flows through catalog tables/views, not in-memory session state +- Supports the same Python executable resolution as Spark's own PySpark + +### IPySparkConnectInterpreter +- IPython variant of PySpark +- Same native-client model as `PySparkConnectInterpreter` + +### SparkConnectUtils +- Stateless utilities for: + - Building Spark Connect URIs with token/SSL/user_id params + - Formatting DataFrames as Zeppelin `%table` output + - Streaming large result sets to avoid memory overflow + +### NotebookLockManager +- Per-notebook `ReentrantLock` registry (fair FIFO ordering) +- Ensures sequential query execution within a single notebook +- Prevents concurrent modifications to shared notebook state + +## Session Isolation and Multi-User Support + +Each user gets **isolated Spark sessions** tracked in a global `ConcurrentHashMap`: +- Username extracted from Zeppelin auth (falls back to `"anonymous"`) +- Per-user quota enforced (`maxSessionsPerUser`, default 5) +- Prevents runaway session proliferation + +Within a notebook, a fair `ReentrantLock` ensures: +- Only one query executes at a time (even with `concurrentSQL=true`) +- FIFO ordering prevents starvation + +## Testing + +### Unit Tests (No Spark Server Required) + +Tests for `SparkConnectUtils` utility class: +```bash +mvn test -pl spark-connect -Dtest=SparkConnectUtilsTest +``` + +### Integration Tests (Requires Spark Connect Server) + +Full interpreter tests with a live Spark server: +```bash +SPARK_CONNECT_TEST_REMOTE=sc://localhost:15002 \ +mvn test -pl spark-connect +``` + +Only integration tests are executed when `SPARK_CONNECT_TEST_REMOTE` is set. + +## Limitations + +1. **No Scala interpreter** — Spark Connect is a client-only protocol; embedded Scala REPL not supported +2. **No R support** — `%spark.r` and `%spark.ir` not available +3. **No ZeppelinContext** — `z.show()`, Angular widgets, and other Zeppelin-specific features return `null` +4. **Spark 3.5.x only** — The gRPC protocol is version-locked to Spark 3.5 +5. **No progress tracking** — Job progress API always returns 0 +6. **Python and Java/SQL use separate Spark Connect sessions** — they connect to the same Spark Connect server, so data is shared through catalog tables (e.g. `createOrReplaceTempView` on a Hive/managed table, or persisting via `spark.sql("CREATE TABLE ...")`), but **session-local** temp views created in Python are not visible to `%spark-connect` SQL paragraphs, and vice versa + +## Dependency Shading + +The module uses Maven Shade Plugin to relocate conflicting dependencies: +- `io.netty` → `org.apache.zeppelin.spark.connect.io.netty` +- `com.google` → `org.apache.zeppelin.spark.connect.com.google` +- `io.grpc` → `org.apache.zeppelin.spark.connect.io.grpc` + +This prevents classpath conflicts with Zeppelin Server's own Netty and other interpreters. + +## Examples + +### Connect to a Remote Spark Cluster + +Configure in Zeppelin UI: +- **URL**: `sc://spark-server.example.com:15002` +- **Token**: `your-auth-token` (if required) +- **Use SSL**: Enable if cluster uses TLS + +### Run Multi-Statement SQL + +```sql +%spark-connect + +CREATE OR REPLACE TEMP VIEW my_view AS + SELECT * FROM source_table WHERE year = 2024; + +SELECT COUNT(*) FROM my_view; +``` + +### PySpark with Pandas + +```python +%spark-connect.pyspark + +# Create a Spark DataFrame and convert to Pandas +import pandas as pd +df_spark = spark.sql("SELECT * FROM my_table") +df_pandas = df_spark.toPandas() +print(df_pandas.head()) +``` + +### Inspect DataFrame Schema + +```python +%spark-connect.pyspark + +df = spark.sql("SELECT * FROM events LIMIT 1") +df.explain() # Logical and physical plan +df.printSchema() # Column names and types +``` + +## Troubleshooting + +### Connection Refused + +- Ensure Spark Connect server is running: `spark-shell --master --conf spark.connect.grpc.binding.port=15002` +- Verify network connectivity and firewall rules +- Check `zeppelin.spark.connect.url` configuration + +### Authentication Failures + +- Verify `zeppelin.spark.connect.token` matches server token requirements +- Enable `zeppelin.spark.connect.use_ssl` if cluster uses TLS + +### Out of Memory + +- Use `zeppelin.spark.maxResult` to limit rows fetched +- Use `spark.sql(...).limit(n)` in queries to reduce data transfer +- Enable `zeppelin.spark.connect.use_ssl` to stream results instead of collecting + +## References + +- [Apache Spark Connect Documentation](https://spark.apache.org/docs/latest/spark-connect-overview.html) +- [Spark Connect Protocol](https://spark.apache.org/docs/latest/spark-connect-introduction.html) +- [Zeppelin Interpreter Development](https://zeppelin.apache.org/docs/latest/usage/interpreter/interpreter_binding_mode.html) diff --git a/spark-connect/pom.xml b/spark-connect/pom.xml new file mode 100644 index 00000000000..cb2a5bd9dc3 --- /dev/null +++ b/spark-connect/pom.xml @@ -0,0 +1,130 @@ + + + + + 4.0.0 + + + zeppelin-interpreter-parent + org.apache.zeppelin + 0.11.2 + ../zeppelin-interpreter-parent/pom.xml + + + spark-connect-interpreter + jar + Zeppelin: Spark Connect Interpreter + Zeppelin Spark Connect support via gRPC client + + + spark-connect + 3.5.3 + 2.12 + + + + + org.apache.spark + spark-connect-client-jvm_${spark.scala.binary.version} + ${spark.connect.version} + + + + org.apache.zeppelin + zeppelin-python + ${project.version} + + + + org.apache.commons + commons-lang3 + + + + org.mockito + mockito-core + test + + + + + + + maven-resources-plugin + + + maven-shade-plugin + + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + reference.conf + + + + + org.apache.zeppelin:zeppelin-interpreter-shaded + + + + + io.netty + org.apache.zeppelin.spark.connect.io.netty + + + com.google + org.apache.zeppelin.spark.connect.com.google + + + io.grpc + org.apache.zeppelin.spark.connect.io.grpc + + + + + + org.apache.maven.plugins + maven-checkstyle-plugin + + + + + + + spark-connect-3.5 + + true + + + 3.5.3 + + + + + diff --git a/spark-connect/src/main/java/org/apache/zeppelin/spark/IPySparkConnectInterpreter.java b/spark-connect/src/main/java/org/apache/zeppelin/spark/IPySparkConnectInterpreter.java new file mode 100644 index 00000000000..73e5fb23211 --- /dev/null +++ b/spark-connect/src/main/java/org/apache/zeppelin/spark/IPySparkConnectInterpreter.java @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.python.IPythonInterpreter; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.Map; +import java.util.Properties; + +/** + * PySpark Connect Interpreter which uses IPython underlying. + * Uses PySpark's native Spark Connect client — IPython opens its own + * Spark Connect session pointed at the same gRPC server as the Java/SQL + * interpreter. No Py4j bridge. + */ +public class IPySparkConnectInterpreter extends IPythonInterpreter { + + private static final Logger LOGGER = LoggerFactory.getLogger(IPySparkConnectInterpreter.class); + + private PySparkConnectInterpreter pySparkConnectInterpreter; + private boolean opened = false; + private InterpreterContext curIntpContext; + + public IPySparkConnectInterpreter(Properties property) { + super(property); + } + + @Override + public synchronized void open() throws InterpreterException { + if (opened) { + return; + } + + this.pySparkConnectInterpreter = + getInterpreterInTheSameSessionByClassName(PySparkConnectInterpreter.class, false); + + if (pySparkConnectInterpreter != null) { + setProperty("zeppelin.python", pySparkConnectInterpreter.getPythonExec()); + } + setAdditionalPythonInitFile("python/zeppelin_sparkconnect.py"); + super.open(); + opened = true; + } + + @Override + protected Map setupKernelEnv() throws IOException { + Map envs = super.setupKernelEnv(); + String remote = SparkConnectUtils.buildConnectionString(getProperties(), getUserName()); + envs.put("SPARK_REMOTE", remote); + LOGGER.info("Set SPARK_REMOTE for IPython native Spark Connect client: {}", + remote.replaceAll("token=[^;]*", "token=[REDACTED]") + .replaceAll("user_id=[^;]*", "user_id=[REDACTED]")); + return envs; + } + + @Override + public org.apache.zeppelin.interpreter.InterpreterResult interpret(String st, + InterpreterContext context) throws InterpreterException { + InterpreterContext.set(context); + this.curIntpContext = context; + String setInptContextStmt = "intp.setInterpreterContextInPython()"; + org.apache.zeppelin.interpreter.InterpreterResult result = + super.interpret(setInptContextStmt, context); + if (result.code().equals(org.apache.zeppelin.interpreter.InterpreterResult.Code.ERROR)) { + return new org.apache.zeppelin.interpreter.InterpreterResult( + org.apache.zeppelin.interpreter.InterpreterResult.Code.ERROR, + "Fail to setCurIntpContext"); + } + + return super.interpret(st, context); + } + + public void setInterpreterContextInPython() { + InterpreterContext.set(curIntpContext); + } + + @Override + public void close() throws InterpreterException { + LOGGER.info("Close IPySparkConnectInterpreter (opened={})", opened); + try { + super.close(); + } finally { + opened = false; + pySparkConnectInterpreter = null; + LOGGER.info("IPySparkConnectInterpreter closed and state reset — ready for re-open"); + } + } + + @Override + public int getProgress(InterpreterContext context) throws InterpreterException { + return 0; + } +} diff --git a/spark-connect/src/main/java/org/apache/zeppelin/spark/NotebookLockManager.java b/spark-connect/src/main/java/org/apache/zeppelin/spark/NotebookLockManager.java new file mode 100644 index 00000000000..8f39e937cc9 --- /dev/null +++ b/spark-connect/src/main/java/org/apache/zeppelin/spark/NotebookLockManager.java @@ -0,0 +1,63 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Shared utility class for managing notebook-level locks. + * Ensures that only one query executes at a time per notebook, + * regardless of which interpreter (SparkConnectInterpreter or SparkConnectSqlInterpreter) is used. + */ +public class NotebookLockManager { + + // Locks per notebook to ensure one query at a time per notebook + private static final ConcurrentHashMap notebookLocks = + new ConcurrentHashMap<>(); + + /** + * Get or create a lock for the specified notebook. + * Uses fair locking to ensure FIFO ordering of query execution. + * + * @param noteId The notebook ID + * @return The lock for this notebook + */ + public static ReentrantLock getNotebookLock(String noteId) { + return notebookLocks.computeIfAbsent(noteId, + k -> new ReentrantLock(true)); // Fair lock for FIFO ordering + } + + /** + * Remove the lock for a notebook (cleanup when notebook is closed). + * + * @param noteId The notebook ID + */ + public static void removeNotebookLock(String noteId) { + notebookLocks.remove(noteId); + } + + /** + * Get the number of active notebook locks (for monitoring/debugging). + * + * @return The number of active locks + */ + public static int getActiveLockCount() { + return notebookLocks.size(); + } +} diff --git a/spark-connect/src/main/java/org/apache/zeppelin/spark/PySparkConnectInterpreter.java b/spark-connect/src/main/java/org/apache/zeppelin/spark/PySparkConnectInterpreter.java new file mode 100644 index 00000000000..672fe7678d9 --- /dev/null +++ b/spark-connect/src/main/java/org/apache/zeppelin/spark/PySparkConnectInterpreter.java @@ -0,0 +1,180 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.commons.lang3.StringUtils; +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.python.PythonInterpreter; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.Map; +import java.util.Properties; + +/** + * PySpark interpreter for Spark Connect. + * Uses PySpark's native Spark Connect client — Python opens its own + * Spark Connect session pointed at the same gRPC server as the Java/SQL + * interpreter. No Py4j bridge. + */ +public class PySparkConnectInterpreter extends PythonInterpreter { + + private static final Logger LOGGER = LoggerFactory.getLogger(PySparkConnectInterpreter.class); + + private InterpreterContext curIntpContext; + + public PySparkConnectInterpreter(Properties property) { + super(property); + } + + @Override + public void open() throws InterpreterException { + setProperty("zeppelin.python.useIPython", + getProperty("zeppelin.pyspark.connect.useIPython", "true")); + + String pythonExec = getPythonExec(); + LOGGER.info("Python executable resolved: {}", pythonExec); + + super.open(); + + if (!useIPython()) { + try { + bootstrapInterpreter("python/zeppelin_sparkconnect.py"); + } catch (IOException e) { + LOGGER.error("Fail to bootstrap spark connect", e); + throw new InterpreterException("Fail to bootstrap spark connect", e); + } + } + } + + @Override + public void close() throws InterpreterException { + LOGGER.info("Close PySparkConnectInterpreter"); + super.close(); + } + + @Override + protected org.apache.zeppelin.python.IPythonInterpreter getIPythonInterpreter() + throws InterpreterException { + return getInterpreterInTheSameSessionByClassName(IPySparkConnectInterpreter.class, false); + } + + @Override + public org.apache.zeppelin.interpreter.InterpreterResult interpret(String st, + InterpreterContext context) throws InterpreterException { + curIntpContext = context; + return super.interpret(st, context); + } + + @Override + protected void preCallPython(InterpreterContext context) { + callPython(new PythonInterpretRequest( + "intp.setInterpreterContextInPython()", false, false)); + } + + public void setInterpreterContextInPython() { + InterpreterContext.set(curIntpContext); + } + + @Override + protected Map setupPythonEnv() throws IOException { + Map env = super.setupPythonEnv(); + + String pythonExec = getPythonExec(); + env.put("PYSPARK_PYTHON", pythonExec); + LOGGER.info("Set PYSPARK_PYTHON: {}", pythonExec); + + String remote = SparkConnectUtils.buildConnectionString(getProperties(), getUserName()); + env.put("SPARK_REMOTE", remote); + LOGGER.info("Set SPARK_REMOTE for PySpark native Spark Connect client: {}", + remote.replaceAll("token=[^;]*", "token=[REDACTED]") + .replaceAll("user_id=[^;]*", "user_id=[REDACTED]")); + + setupCondaLibraryPath(env, pythonExec); + + LOGGER.info("LD_LIBRARY_PATH: {}", env.get("LD_LIBRARY_PATH")); + return env; + } + + /** + * Get Python executable following Spark's PySpark detection pattern exactly: + * 1. spark.pyspark.driver.python (from Spark Connect properties) + * 2. spark.pyspark.python (from Spark Connect properties) + * 3. PYSPARK_DRIVER_PYTHON (environment variable) + * 4. PYSPARK_PYTHON (environment variable) + * 5. zeppelin.python (Zeppelin property) - if set, validate it + * 6. Default to "python" (let system PATH handle it, just like Spark does) + */ + @Override + protected String getPythonExec() { + String driverPython = getProperty("spark.pyspark.driver.python", ""); + if (StringUtils.isNotBlank(driverPython)) { + LOGGER.info("Using Python executable from spark.pyspark.driver.python: {}", driverPython); + return driverPython; + } + + String pysparkPython = getProperty("spark.pyspark.python", ""); + if (StringUtils.isNotBlank(pysparkPython)) { + LOGGER.info("Using Python executable from spark.pyspark.python: {}", pysparkPython); + return pysparkPython; + } + + String envDriverPython = System.getenv("PYSPARK_DRIVER_PYTHON"); + if (StringUtils.isNotBlank(envDriverPython)) { + LOGGER.info("Using Python executable from PYSPARK_DRIVER_PYTHON: {}", envDriverPython); + return envDriverPython; + } + + String envPysparkPython = System.getenv("PYSPARK_PYTHON"); + if (StringUtils.isNotBlank(envPysparkPython)) { + LOGGER.info("Using Python executable from PYSPARK_PYTHON: {}", envPysparkPython); + return envPysparkPython; + } + + String zeppelinPython = getProperty("zeppelin.python", ""); + if (StringUtils.isNotBlank(zeppelinPython)) { + LOGGER.info("Using Python executable from zeppelin.python property: {}", zeppelinPython); + return zeppelinPython; + } + + LOGGER.info("No Python executable configured, defaulting to 'python' (will use system PATH)"); + return "python"; + } + + private void setupCondaLibraryPath(Map env, String pythonExec) { + if (pythonExec != null && pythonExec.contains("/conda/")) { + int binIndex = pythonExec.indexOf("/bin/"); + if (binIndex > 0) { + String condaBase = pythonExec.substring(0, binIndex); + String condaLib = condaBase + "/lib"; + java.io.File libDir = new java.io.File(condaLib); + if (libDir.exists() && libDir.isDirectory()) { + String ldLibraryPath = env.getOrDefault("LD_LIBRARY_PATH", ""); + if (ldLibraryPath.isEmpty()) { + env.put("LD_LIBRARY_PATH", condaLib); + } else if (!ldLibraryPath.contains(condaLib)) { + env.put("LD_LIBRARY_PATH", condaLib + ":" + ldLibraryPath); + } + LOGGER.info("Added conda lib directory to LD_LIBRARY_PATH: {}", condaLib); + } + } + } + } +} diff --git a/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectInterpreter.java b/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectInterpreter.java new file mode 100644 index 00000000000..b8ae477a0e8 --- /dev/null +++ b/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectInterpreter.java @@ -0,0 +1,343 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.exception.ExceptionUtils; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.zeppelin.interpreter.AbstractInterpreter; +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.interpreter.InterpreterResult; +import org.apache.zeppelin.interpreter.InterpreterResult.Code; +import org.apache.zeppelin.interpreter.ZeppelinContext; +import org.apache.zeppelin.interpreter.thrift.InterpreterCompletion; +import org.apache.zeppelin.interpreter.util.SqlSplitter; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Spark Connect interpreter for Zeppelin. + * Connects to a remote Spark cluster via the Spark Connect gRPC protocol. + * + * Session & Concurrency Model: + *
    + *
  • Max Spark Connect sessions per user is capped + * ({@code zeppelin.spark.connect.maxSessionsPerUser}, default 5). + * Each interpreter instance creates one session; Zeppelin's binding mode + * (per-user / per-note / scoped / isolated) controls how many instances exist.
    • + *
  • Notebooks are unlimited -- users may open as many as they like.
    • + *
  • Within a single notebook only one query executes at a time + * (per-notebook fair lock via {@link NotebookLockManager}).
    • + *
+ */ +public class SparkConnectInterpreter extends AbstractInterpreter { + + private static final Logger LOGGER = LoggerFactory.getLogger(SparkConnectInterpreter.class); + + /** user -> number of active SparkSession instances owned by that user. */ + private static final ConcurrentHashMap userSessionCount = + new ConcurrentHashMap<>(); + + private static final int DEFAULT_MAX_SESSIONS_PER_USER = 5; + + private SparkSession sparkSession; + private SqlSplitter sqlSplitter; + private int maxResult; + private String currentUser; + private volatile boolean sessionSlotAcquired = false; + + public SparkConnectInterpreter(Properties properties) { + super(properties); + } + + @Override + public synchronized void open() throws InterpreterException { + if (sparkSession != null) { + LOGGER.warn("open() called but sparkSession is already active — skipping. " + + "Call close() first to restart the interpreter."); + return; + } + + try { + currentUser = getUserName(); + if (StringUtils.isBlank(currentUser)) { + currentUser = "anonymous"; + } + + int maxSessions = Integer.parseInt( + getProperty("zeppelin.spark.connect.maxSessionsPerUser", + String.valueOf(DEFAULT_MAX_SESSIONS_PER_USER))); + + if (!acquireSessionSlot(currentUser, maxSessions)) { + throw new InterpreterException( + String.format("User '%s' already has %d active Spark Connect sessions " + + "(max %d). Close an existing interpreter before opening a new one.", + currentUser, maxSessions, maxSessions)); + } + sessionSlotAcquired = true; + + LOGGER.info("Opening SparkConnectInterpreter for user: {} (session slot {}/{})", + currentUser, userSessionCount.getOrDefault(currentUser, 0), maxSessions); + + String remoteUrl = SparkConnectUtils.buildConnectionString(getProperties(), currentUser); + LOGGER.info("Connecting to Spark Connect server at: {}", + remoteUrl.replaceAll("token=[^;]*", "token=[REDACTED]") + .replaceAll("user_id=[^;]*", "user_id=[REDACTED]")); + + // Clear the thread-local active session on the Spark Connect client so that + // getOrCreate() creates a fresh remote session rather than reusing the previous + // closed one. We intentionally do NOT call clearDefaultSession() because that + // is a JVM-global operation and would disrupt other interpreter instances that + // are concurrently active in the same process. + try { + SparkSession.clearActiveSession(); + LOGGER.info("Cleared thread-local active Spark session (safe for multi-interpreter)"); + } catch (Exception e) { + LOGGER.warn("Could not clear active Spark session (non-fatal): {}", e.getMessage()); + } + + SparkSession.Builder builder = SparkSession.builder().remote(remoteUrl); + + String appName = getProperty("spark.app.name", "Zeppelin Spark Connect"); + if (StringUtils.isNotBlank(appName)) { + builder.appName(appName); + } + + String grpcMaxMsgSize = getProperty( + "spark.connect.grpc.maxMessageSize", "134217728"); + builder.config("spark.connect.grpc.maxMessageSize", grpcMaxMsgSize); + + for (Object key : getProperties().keySet()) { + String keyStr = key.toString(); + String value = getProperties().getProperty(keyStr); + if (StringUtils.isNotBlank(value) + && keyStr.startsWith("spark.") + && !keyStr.equals("spark.remote") + && !keyStr.equals("spark.connect.token") + && !keyStr.equals("spark.connect.use_ssl") + && !keyStr.equals("spark.app.name") + && !keyStr.equals("spark.connect.grpc.maxMessageSize")) { + builder.config(keyStr, value); + } + } + + sparkSession = builder.getOrCreate(); + LOGGER.info("Spark Connect session established for user: {}", currentUser); + + maxResult = Integer.parseInt(getProperty("zeppelin.spark.maxResult", "1000")); + sqlSplitter = new SqlSplitter(); + } catch (InterpreterException ie) { + throw ie; + } catch (Exception e) { + if (sessionSlotAcquired) { + releaseSessionSlot(currentUser); + sessionSlotAcquired = false; + } + LOGGER.error("Failed to connect to Spark Connect server", e); + throw new InterpreterException("Failed to connect to Spark Connect server: " + + e.getMessage(), e); + } + } + + @Override + public void close() throws InterpreterException { + LOGGER.info("Closing SparkConnectInterpreter for user: {} (sparkSession={})", + currentUser, sparkSession != null ? "active" : "null"); + if (sparkSession != null) { + try { + sparkSession.close(); + LOGGER.info("Spark Connect session closed for user: {}", currentUser); + } catch (Exception e) { + LOGGER.warn("Error closing Spark Connect session", e); + } finally { + sparkSession = null; + } + } else { + LOGGER.info("close() called but no active sparkSession — nothing to tear down"); + } + if (sessionSlotAcquired) { + releaseSessionSlot(currentUser); + sessionSlotAcquired = false; + } + } + + @Override + public ZeppelinContext getZeppelinContext() { + return null; + } + + @Override + public InterpreterResult internalInterpret(String st, InterpreterContext context) + throws InterpreterException { + if (sparkSession == null) { + return new InterpreterResult(Code.ERROR, + "Spark Connect session is not initialized. Check connection settings."); + } + + String noteId = context.getNoteId(); + if (StringUtils.isBlank(noteId)) { + return new InterpreterResult(Code.ERROR, + "Note ID is missing from interpreter context."); + } + + // Per-notebook lock: only one query at a time inside a notebook + ReentrantLock notebookLock = NotebookLockManager.getNotebookLock(noteId); + notebookLock.lock(); + try { + List sqls = sqlSplitter.splitSql(st); + int limit = Integer.parseInt(context.getLocalProperties().getOrDefault("limit", + String.valueOf(maxResult))); + + boolean useStreaming = Boolean.parseBoolean( + getProperty("zeppelin.spark.connect.streamResults", "false")); + + String curSql = null; + try { + for (String sql : sqls) { + curSql = sql; + if (StringUtils.isBlank(sql)) { + continue; + } + Dataset df = sparkSession.sql(sql); + if (useStreaming) { + SparkConnectUtils.streamDataFrame(df, limit, context.out); + } else { + String result = SparkConnectUtils.showDataFrame(df, limit); + context.out.write(result); + } + } + context.out.flush(); + } catch (Exception e) { + return handleSqlException(e, curSql, context); + } + + return new InterpreterResult(Code.SUCCESS); + } finally { + notebookLock.unlock(); + } + } + + // ---- session-slot helpers (static, shared across all instances) ---- + + /** + * Try to claim one session slot for the user. + * @return true if a slot was available and claimed + */ + private static synchronized boolean acquireSessionSlot(String user, int maxSessions) { + int current = userSessionCount.getOrDefault(user, 0); + if (current >= maxSessions) { + LOGGER.warn("User {} already has {} active Spark Connect sessions (max {})", + user, current, maxSessions); + return false; + } + userSessionCount.put(user, current + 1); + LOGGER.info("Acquired session slot for user {}. Active sessions: {}/{}", + user, current + 1, maxSessions); + return true; + } + + /** + * Release one session slot for the user. + */ + private static synchronized void releaseSessionSlot(String user) { + if (user == null) { + return; + } + int current = userSessionCount.getOrDefault(user, 0); + if (current <= 1) { + userSessionCount.remove(user); + } else { + userSessionCount.put(user, current - 1); + } + LOGGER.info("Released session slot for user {}. Remaining sessions: {}", + user, Math.max(0, current - 1)); + } + + /** Visible for testing. */ + static int getActiveSessionCount(String user) { + return userSessionCount.getOrDefault(user, 0); + } + + private InterpreterResult handleSqlException(Exception e, String sql, + InterpreterContext context) { + try { + LOGGER.error("Error executing SQL: {}", sql, e); + context.out.write("\nError in SQL: " + sql + "\n"); + if (Boolean.parseBoolean(getProperty("zeppelin.spark.sql.stacktrace", "true"))) { + if (e.getCause() != null) { + context.out.write(ExceptionUtils.getStackTrace(e.getCause())); + } else { + context.out.write(ExceptionUtils.getStackTrace(e)); + } + } else { + String msg = e.getCause() != null ? e.getCause().getMessage() : e.getMessage(); + context.out.write(msg + + "\nSet zeppelin.spark.sql.stacktrace = true to see full stacktrace"); + } + context.out.flush(); + } catch (IOException ex) { + LOGGER.error("Failed to write error output", ex); + } + return new InterpreterResult(Code.ERROR); + } + + @Override + public void cancel(InterpreterContext context) throws InterpreterException { + if (sparkSession != null) { + try { + sparkSession.interruptAll(); + } catch (Exception e) { + LOGGER.warn("Error interrupting Spark Connect session", e); + } + } + } + + @Override + public FormType getFormType() { + return FormType.SIMPLE; + } + + @Override + public int getProgress(InterpreterContext context) throws InterpreterException { + return 0; + } + + @Override + public List completion(String buf, int cursor, + InterpreterContext interpreterContext) throws InterpreterException { + return new ArrayList<>(); + } + + public SparkSession getSparkSession() { + return sparkSession; + } + + public int getMaxResult() { + return maxResult; + } +} diff --git a/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectSqlInterpreter.java b/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectSqlInterpreter.java new file mode 100644 index 00000000000..e642455d9aa --- /dev/null +++ b/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectSqlInterpreter.java @@ -0,0 +1,204 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.exception.ExceptionUtils; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.zeppelin.interpreter.AbstractInterpreter; +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.interpreter.InterpreterResult; +import org.apache.zeppelin.interpreter.InterpreterResult.Code; +import org.apache.zeppelin.interpreter.ZeppelinContext; +import org.apache.zeppelin.interpreter.thrift.InterpreterCompletion; +import org.apache.zeppelin.interpreter.util.SqlSplitter; +import org.apache.zeppelin.scheduler.Scheduler; +import org.apache.zeppelin.scheduler.SchedulerFactory; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Spark Connect SQL interpreter for Zeppelin. + * Delegates to SparkConnectInterpreter for the SparkSession, providing + * dedicated SQL execution with concurrent query support. + */ +public class SparkConnectSqlInterpreter extends AbstractInterpreter { + + private static final Logger LOGGER = LoggerFactory.getLogger(SparkConnectSqlInterpreter.class); + + private SparkConnectInterpreter sparkConnectInterpreter; + private SqlSplitter sqlSplitter; + + public SparkConnectSqlInterpreter(Properties properties) { + super(properties); + } + + @Override + public void open() throws InterpreterException { + this.sparkConnectInterpreter = + getInterpreterInTheSameSessionByClassName(SparkConnectInterpreter.class); + this.sqlSplitter = new SqlSplitter(); + } + + @Override + public void close() throws InterpreterException { + sparkConnectInterpreter = null; + } + + @Override + protected boolean isInterpolate() { + return Boolean.parseBoolean(getProperty("zeppelin.spark.sql.interpolation", "false")); + } + + @Override + public ZeppelinContext getZeppelinContext() { + return null; + } + + @Override + public InterpreterResult internalInterpret(String st, InterpreterContext context) + throws InterpreterException { + SparkSession sparkSession = sparkConnectInterpreter.getSparkSession(); + if (sparkSession == null) { + return new InterpreterResult(Code.ERROR, + "Spark Connect session is not initialized. Check connection settings."); + } + + // Get noteId from context for notebook-level synchronization + String noteId = context.getNoteId(); + if (StringUtils.isBlank(noteId)) { + return new InterpreterResult(Code.ERROR, + "Note ID is missing from interpreter context."); + } + + // Get or create lock for this notebook to ensure sequential execution + // This ensures one query at a time per notebook, even with concurrentSQL enabled + ReentrantLock notebookLock = NotebookLockManager.getNotebookLock(noteId); + + // Acquire lock to ensure only one query executes at a time for this notebook + notebookLock.lock(); + try { + List sqls = sqlSplitter.splitSql(st); + int maxResult = Integer.parseInt(context.getLocalProperties().getOrDefault("limit", + String.valueOf(sparkConnectInterpreter.getMaxResult()))); + + boolean useStreaming = Boolean.parseBoolean( + getProperty("zeppelin.spark.connect.streamResults", "false")); + + String curSql = null; + try { + for (String sql : sqls) { + curSql = sql; + if (StringUtils.isBlank(sql)) { + continue; + } + Dataset df = sparkSession.sql(sql); + if (useStreaming) { + SparkConnectUtils.streamDataFrame(df, maxResult, context.out); + } else { + String result = SparkConnectUtils.showDataFrame(df, maxResult); + context.out.write(result); + } + } + context.out.flush(); + } catch (Exception e) { + try { + LOGGER.error("Error executing SQL: {}", curSql, e); + context.out.write("\nError in SQL: " + curSql + "\n"); + if (Boolean.parseBoolean(getProperty("zeppelin.spark.sql.stacktrace", "true"))) { + if (e.getCause() != null) { + context.out.write(ExceptionUtils.getStackTrace(e.getCause())); + } else { + context.out.write(ExceptionUtils.getStackTrace(e)); + } + } else { + String msg = e.getCause() != null ? e.getCause().getMessage() : e.getMessage(); + context.out.write(msg + + "\nSet zeppelin.spark.sql.stacktrace = true to see full stacktrace"); + } + context.out.flush(); + } catch (IOException ex) { + LOGGER.error("Failed to write error output", ex); + } + return new InterpreterResult(Code.ERROR); + } + + return new InterpreterResult(Code.SUCCESS); + } finally { + notebookLock.unlock(); + } + } + + @Override + public void cancel(InterpreterContext context) throws InterpreterException { + SparkSession sparkSession = sparkConnectInterpreter.getSparkSession(); + if (sparkSession != null) { + try { + sparkSession.interruptAll(); + } catch (Exception e) { + LOGGER.warn("Error interrupting Spark Connect session", e); + } + } + } + + @Override + public FormType getFormType() { + return FormType.SIMPLE; + } + + @Override + public int getProgress(InterpreterContext context) throws InterpreterException { + return 0; + } + + @Override + public List completion(String buf, int cursor, + InterpreterContext interpreterContext) throws InterpreterException { + return new ArrayList<>(); + } + + @Override + public Scheduler getScheduler() { + if (concurrentSQL()) { + int maxConcurrency = Integer.parseInt( + getProperty("zeppelin.spark.concurrentSQL.max", "10")); + return SchedulerFactory.singleton().createOrGetParallelScheduler( + SparkConnectSqlInterpreter.class.getName() + this.hashCode(), maxConcurrency); + } else { + try { + return getInterpreterInTheSameSessionByClassName( + SparkConnectInterpreter.class, false).getScheduler(); + } catch (InterpreterException e) { + throw new RuntimeException("Failed to get scheduler", e); + } + } + } + + private boolean concurrentSQL() { + return Boolean.parseBoolean(getProperty("zeppelin.spark.concurrentSQL")); + } +} diff --git a/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectUtils.java b/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectUtils.java new file mode 100644 index 00000000000..ffad1fbb789 --- /dev/null +++ b/spark-connect/src/main/java/org/apache/zeppelin/spark/SparkConnectUtils.java @@ -0,0 +1,191 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.commons.lang3.StringUtils; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.util.Iterator; +import java.util.List; +import java.util.Properties; + +public class SparkConnectUtils { + + static final int DEFAULT_TRUNCATE_LENGTH = 256; + + private SparkConnectUtils() { + } + + /** + * Build the Spark Connect connection string from interpreter properties. + * Format: sc://hostname:port[/;param1=val1;param2=val2] + * + * Supports: token, use_ssl, user_id, and any extra params already in the URI. + * Examples: + * sc://localhost:15002 + * sc://localhost:15002/;use_ssl=true;token=abc123;user_id=alice + * sc://ranking-cluster-m:8080 + */ + public static String buildConnectionString(Properties properties) { + return buildConnectionString(properties, null); + } + + /** + * Build the Spark Connect connection string, including user_id so the Spark Connect + * server can attribute the session to the correct user in its own UI. + * + * @param properties interpreter properties + * @param userName the authenticated Zeppelin username; ignored if blank + */ + public static String buildConnectionString(Properties properties, String userName) { + String remote = properties.getProperty("spark.remote", "sc://localhost:15002"); + StringBuilder params = new StringBuilder(); + + String token = properties.getProperty("spark.connect.token", ""); + if (StringUtils.isNotBlank(token)) { + params.append(";token=").append(token); + } + + boolean useSsl = Boolean.parseBoolean( + properties.getProperty("spark.connect.use_ssl", "false")); + if (useSsl) { + params.append(";use_ssl=true"); + } + + if (StringUtils.isNotBlank(userName) && !remote.contains("user_id=")) { + params.append(";user_id=").append(userName); + } + + if (params.length() > 0) { + if (remote.contains(";")) { + remote = remote + params; + } else { + remote = remote + "/" + params; + } + } + return remote; + } + + /** + * Convert a Dataset to Zeppelin's %table format string. + * Applies limit before collecting to prevent OOM on the driver. + * Truncates cell values to avoid excessively wide output. + */ + public static String showDataFrame(Dataset df, int maxResult) { + return showDataFrame(df, maxResult, DEFAULT_TRUNCATE_LENGTH); + } + + public static String showDataFrame(Dataset df, int maxResult, int truncateLength) { + StructType schema = df.schema(); + StructField[] fields = schema.fields(); + + int effectiveLimit = Math.max(1, Math.min(maxResult, 100_000)); + + int estimatedRowSize = Math.max(fields.length * 20, 100); + int estimatedTotalBytes = estimatedRowSize * effectiveLimit; + StringBuilder sb = new StringBuilder(Math.min(estimatedTotalBytes, 10 * 1024 * 1024)); + sb.append("%table "); + + for (int i = 0; i < fields.length; i++) { + if (i > 0) { + sb.append('\t'); + } + sb.append(replaceReservedChars(fields[i].name())); + } + sb.append('\n'); + + List rows = df.limit(effectiveLimit).collectAsList(); + for (Row row : rows) { + for (int i = 0; i < row.length(); i++) { + if (i > 0) { + sb.append('\t'); + } + Object value = row.get(i); + String cellStr = value == null ? "null" : value.toString(); + if (truncateLength > 0 && cellStr.length() > truncateLength) { + cellStr = cellStr.substring(0, truncateLength) + "..."; + } + sb.append(replaceReservedChars(cellStr)); + } + sb.append('\n'); + } + + return sb.toString(); + } + + /** + * Stream a Dataset as Zeppelin %table format directly to an OutputStream, + * avoiding building the entire result in memory. + * Preferred for large result sets. + */ + public static void streamDataFrame(Dataset df, int maxResult, OutputStream out) + throws IOException { + streamDataFrame(df, maxResult, DEFAULT_TRUNCATE_LENGTH, out); + } + + public static void streamDataFrame(Dataset df, int maxResult, + int truncateLength, OutputStream out) throws IOException { + StructType schema = df.schema(); + StructField[] fields = schema.fields(); + + int effectiveLimit = Math.max(1, Math.min(maxResult, 100_000)); + + StringBuilder header = new StringBuilder("%table "); + for (int i = 0; i < fields.length; i++) { + if (i > 0) { + header.append('\t'); + } + header.append(replaceReservedChars(fields[i].name())); + } + header.append('\n'); + out.write(header.toString().getBytes(StandardCharsets.UTF_8)); + + Iterator it = df.limit(effectiveLimit).toLocalIterator(); + StringBuilder rowBuf = new StringBuilder(256); + while (it.hasNext()) { + rowBuf.setLength(0); + Row row = it.next(); + for (int i = 0; i < row.length(); i++) { + if (i > 0) { + rowBuf.append('\t'); + } + Object value = row.get(i); + String cellStr = value == null ? "null" : value.toString(); + if (truncateLength > 0 && cellStr.length() > truncateLength) { + cellStr = cellStr.substring(0, truncateLength) + "..."; + } + rowBuf.append(replaceReservedChars(cellStr)); + } + rowBuf.append('\n'); + out.write(rowBuf.toString().getBytes(StandardCharsets.UTF_8)); + } + out.flush(); + } + + static String replaceReservedChars(String str) { + if (str == null) { + return "null"; + } + return str.replace('\t', ' ').replace('\n', ' '); + } +} diff --git a/spark-connect/src/main/resources/interpreter-setting.json b/spark-connect/src/main/resources/interpreter-setting.json new file mode 100644 index 00000000000..78d68fda936 --- /dev/null +++ b/spark-connect/src/main/resources/interpreter-setting.json @@ -0,0 +1,187 @@ +[ + { + "group": "spark-connect", + "name": "spark-connect", + "className": "org.apache.zeppelin.spark.SparkConnectInterpreter", + "defaultInterpreter": true, + "properties": { + "spark.remote": { + "envName": "SPARK_REMOTE", + "propertyName": "spark.remote", + "defaultValue": "sc://localhost:15002", + "description": "Spark Connect server URI (e.g. sc://localhost:15002 or sc://dataproc-master:8080)", + "type": "string" + }, + "spark.app.name": { + "envName": null, + "propertyName": "spark.app.name", + "defaultValue": "Zeppelin Spark Connect", + "description": "Spark application name", + "type": "string" + }, + "zeppelin.spark.maxResult": { + "envName": null, + "propertyName": "zeppelin.spark.maxResult", + "defaultValue": "1000", + "description": "Max number of rows to display", + "type": "number" + }, + "zeppelin.spark.sql.stacktrace": { + "envName": null, + "propertyName": "zeppelin.spark.sql.stacktrace", + "defaultValue": true, + "description": "Show full exception stacktrace for SQL errors", + "type": "checkbox" + }, + "spark.connect.grpc.maxMessageSize": { + "envName": null, + "propertyName": "spark.connect.grpc.maxMessageSize", + "defaultValue": "134217728", + "description": "Max gRPC message size in bytes (default 128MB). Increase for large result sets.", + "type": "number" + }, + "zeppelin.spark.connect.streamResults": { + "envName": null, + "propertyName": "zeppelin.spark.connect.streamResults", + "defaultValue": false, + "description": "Stream query results row-by-row instead of building full result in memory. Recommended for large result sets.", + "type": "checkbox" + }, + "zeppelin.spark.connect.maxSessionsPerUser": { + "envName": null, + "propertyName": "zeppelin.spark.connect.maxSessionsPerUser", + "defaultValue": "5", + "description": "Maximum number of Spark Connect sessions (SparkSession instances) per user. Each interpreter instance creates one session.", + "type": "number" + } + }, + "editor": { + "language": "sql", + "editOnDblClick": false, + "completionKey": "TAB", + "completionSupport": false + } + }, + { + "group": "spark-connect", + "name": "sql", + "className": "org.apache.zeppelin.spark.SparkConnectSqlInterpreter", + "properties": { + "zeppelin.spark.concurrentSQL": { + "envName": null, + "propertyName": "zeppelin.spark.concurrentSQL", + "defaultValue": true, + "description": "Execute multiple SQL concurrently", + "type": "checkbox" + }, + "zeppelin.spark.concurrentSQL.max": { + "envName": null, + "propertyName": "zeppelin.spark.concurrentSQL.max", + "defaultValue": "10", + "description": "Max concurrent SQL executions", + "type": "number" + }, + "zeppelin.spark.sql.stacktrace": { + "envName": null, + "propertyName": "zeppelin.spark.sql.stacktrace", + "defaultValue": true, + "description": "Show full exception stacktrace for SQL errors", + "type": "checkbox" + } + }, + "editor": { + "language": "sql", + "editOnDblClick": false, + "completionKey": "TAB", + "completionSupport": false + } + }, + { + "group": "spark-connect", + "name": "pyspark", + "className": "org.apache.zeppelin.spark.PySparkConnectInterpreter", + "properties": { + "spark.remote": { + "envName": "SPARK_REMOTE", + "propertyName": "spark.remote", + "defaultValue": "sc://localhost:15002", + "description": "Spark Connect server URI (e.g. sc://localhost:15002 or sc://dataproc-master:8080)", + "type": "string" + }, + "spark.connect.token": { + "envName": "SPARK_CONNECT_TOKEN", + "propertyName": "spark.connect.token", + "defaultValue": "", + "description": "Authentication token for Spark Connect (optional)", + "type": "string" + }, + "spark.connect.use_ssl": { + "envName": "SPARK_CONNECT_USE_SSL", + "propertyName": "spark.connect.use_ssl", + "defaultValue": false, + "description": "Use SSL for Spark Connect connection", + "type": "checkbox" + }, + "spark.app.name": { + "envName": null, + "propertyName": "spark.app.name", + "defaultValue": "Zeppelin Spark Connect", + "description": "Spark application name", + "type": "string" + }, + "spark.pyspark.driver.python": { + "envName": "PYSPARK_DRIVER_PYTHON", + "propertyName": "spark.pyspark.driver.python", + "defaultValue": "", + "description": "Python executable to use for PySpark driver (highest priority). Follows Spark's PySpark detection pattern.", + "type": "string" + }, + "spark.pyspark.python": { + "envName": "PYSPARK_PYTHON", + "propertyName": "spark.pyspark.python", + "defaultValue": "", + "description": "Python executable to use for PySpark (second priority). Can also be set via PYSPARK_PYTHON environment variable.", + "type": "string" + }, + "zeppelin.python": { + "envName": "ZEPPELIN_PYTHON", + "propertyName": "zeppelin.python", + "defaultValue": "", + "description": "Python executable command (optional fallback). If not set, defaults to 'python' and uses system PATH (matching Spark's PySpark behavior).", + "type": "string" + }, + "zeppelin.pyspark.connect.useIPython": { + "envName": null, + "propertyName": "zeppelin.pyspark.connect.useIPython", + "defaultValue": true, + "description": "Use IPython if available", + "type": "checkbox" + }, + "zeppelin.spark.maxResult": { + "envName": null, + "propertyName": "zeppelin.spark.maxResult", + "defaultValue": "1000", + "description": "Max number of rows to display", + "type": "number" + } + }, + "editor": { + "language": "python", + "editOnDblClick": false, + "completionKey": "TAB", + "completionSupport": true + } + }, + { + "group": "spark-connect", + "name": "ipyspark", + "className": "org.apache.zeppelin.spark.IPySparkConnectInterpreter", + "properties": {}, + "editor": { + "language": "python", + "editOnDblClick": false, + "completionKey": "TAB", + "completionSupport": true + } + } +] diff --git a/spark-connect/src/main/resources/python/zeppelin_sparkconnect.py b/spark-connect/src/main/resources/python/zeppelin_sparkconnect.py new file mode 100644 index 00000000000..f73f6883b71 --- /dev/null +++ b/spark-connect/src/main/resources/python/zeppelin_sparkconnect.py @@ -0,0 +1,38 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# Bootstrap a native PySpark Spark Connect session. The Java interpreter +# populates SPARK_REMOTE with the connection URI (including any token / +# use_ssl / user_id params). Java/SQL and Python connect as independent +# sessions to the same Spark Connect server; cross-language sharing +# happens through catalog tables. + +import os +import warnings + +from pyspark.sql import SparkSession + +warnings.filterwarnings(action='ignore', module='pyspark.util') + +_remote = os.environ.get("SPARK_REMOTE") +if not _remote: + raise RuntimeError( + "SPARK_REMOTE env var not set. The Java interpreter is expected " + "to populate it from the 'spark.remote' interpreter property.") + +spark = SparkSession.builder.remote(_remote).getOrCreate() +sqlContext = sqlc = spark diff --git a/spark-connect/src/test/java/org/apache/zeppelin/spark/PySparkConnectInterpreterTest.java b/spark-connect/src/test/java/org/apache/zeppelin/spark/PySparkConnectInterpreterTest.java new file mode 100644 index 00000000000..f477229a1d5 --- /dev/null +++ b/spark-connect/src/test/java/org/apache/zeppelin/spark/PySparkConnectInterpreterTest.java @@ -0,0 +1,147 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.zeppelin.display.AngularObjectRegistry; +import org.apache.zeppelin.interpreter.Interpreter; +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.interpreter.InterpreterGroup; +import org.apache.zeppelin.interpreter.InterpreterOutput; +import org.apache.zeppelin.interpreter.InterpreterResult; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.apache.zeppelin.resource.LocalResourcePool; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +import java.io.IOException; +import java.util.LinkedList; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +/** + * Integration tests for PySparkConnectInterpreter. + * Requires a running Spark Connect server. + * Set SPARK_CONNECT_TEST_REMOTE env var (e.g. sc://localhost:15002) to enable. + */ +@EnabledIfEnvironmentVariable(named = "SPARK_CONNECT_TEST_REMOTE", matches = ".+") +public class PySparkConnectInterpreterTest { + + private static PySparkConnectInterpreter interpreter; + private static SparkConnectInterpreter sparkConnectInterpreter; + private static InterpreterGroup intpGroup; + + @BeforeAll + public static void setUp() throws Exception { + String remote = System.getenv("SPARK_CONNECT_TEST_REMOTE"); + Properties p = new Properties(); + p.setProperty("spark.remote", remote); + p.setProperty("spark.app.name", "ZeppelinPySparkConnectTest"); + p.setProperty("zeppelin.spark.maxResult", "100"); + p.setProperty("zeppelin.pyspark.connect.useIPython", "false"); + p.setProperty("zeppelin.python", "python"); + + intpGroup = new InterpreterGroup(); + + // Create SparkConnectInterpreter first (required dependency) + sparkConnectInterpreter = new SparkConnectInterpreter(p); + sparkConnectInterpreter.setInterpreterGroup(intpGroup); + intpGroup.put("session_1", new LinkedList()); + intpGroup.get("session_1").add(sparkConnectInterpreter); + sparkConnectInterpreter.open(); + + // Create PySparkConnectInterpreter + interpreter = new PySparkConnectInterpreter(p); + interpreter.setInterpreterGroup(intpGroup); + intpGroup.get("session_1").add(interpreter); + interpreter.open(); + } + + @AfterAll + public static void tearDown() throws InterpreterException { + if (interpreter != null) { + interpreter.close(); + } + if (sparkConnectInterpreter != null) { + sparkConnectInterpreter.close(); + } + } + + private static InterpreterContext getInterpreterContext() { + return InterpreterContext.builder() + .setNoteId("noteId") + .setParagraphId("paragraphId") + .setParagraphTitle("title") + .setAngularObjectRegistry(new AngularObjectRegistry(intpGroup.getId(), null)) + .setResourcePool(new LocalResourcePool("id")) + .setInterpreterOut(new InterpreterOutput()) + .setIntpEventClient(mock(RemoteInterpreterEventClient.class)) + .build(); + } + + @Test + void testSimpleQuery() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "df = spark.sql(\"SELECT 1 AS id, 'hello' AS message\")\ndf.show()", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("id") || output.contains("message") || output.contains("hello"), + "Output should contain query results: " + output); + } + + @Test + void testDataFrameVariable() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "df = spark.sql(\"SELECT 1 AS id, 'test' AS name\")\nprint(type(df))", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("DataFrame") || output.contains("pyspark.sql.connect"), + "Output should indicate native PySpark Connect DataFrame type: " + output); + } + + @Test + void testDeltaTableQuery() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + // Test the exact query from user + InterpreterResult result = interpreter.interpret( + "df = spark.sql(\"select * from gold.delta_test\")", context); + // This might fail if table doesn't exist, but should not crash + // We check that interpreter handled it gracefully + assertTrue(result.code() == InterpreterResult.Code.SUCCESS + || result.code() == InterpreterResult.Code.ERROR, + "Should handle query execution (success or error): " + result.code()); + } + + @Test + void testSparkVariableAvailable() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "print('Spark session:', type(spark))", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("SparkSession") || output.contains("spark"), + "Spark session should be available: " + output); + } +} diff --git a/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectInterpreterTest.java b/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectInterpreterTest.java new file mode 100644 index 00000000000..0b3cbc8213b --- /dev/null +++ b/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectInterpreterTest.java @@ -0,0 +1,161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.zeppelin.display.AngularObjectRegistry; +import org.apache.zeppelin.interpreter.Interpreter; +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.interpreter.InterpreterGroup; +import org.apache.zeppelin.interpreter.InterpreterOutput; +import org.apache.zeppelin.interpreter.InterpreterResult; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.apache.zeppelin.resource.LocalResourcePool; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +import java.io.IOException; +import java.util.LinkedList; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +/** + * Integration tests for SparkConnectInterpreter. + * Requires a running Spark Connect server. + * Set SPARK_CONNECT_TEST_REMOTE env var (e.g. sc://localhost:15002) to enable. + */ +@EnabledIfEnvironmentVariable(named = "SPARK_CONNECT_TEST_REMOTE", matches = ".+") +public class SparkConnectInterpreterTest { + + private static SparkConnectInterpreter interpreter; + private static InterpreterGroup intpGroup; + + @BeforeAll + public static void setUp() throws Exception { + String remote = System.getenv("SPARK_CONNECT_TEST_REMOTE"); + Properties p = new Properties(); + p.setProperty("spark.remote", remote); + p.setProperty("spark.app.name", "ZeppelinSparkConnectTest"); + p.setProperty("zeppelin.spark.maxResult", "100"); + p.setProperty("zeppelin.spark.sql.stacktrace", "true"); + + intpGroup = new InterpreterGroup(); + interpreter = new SparkConnectInterpreter(p); + interpreter.setInterpreterGroup(intpGroup); + intpGroup.put("session_1", new LinkedList()); + intpGroup.get("session_1").add(interpreter); + + interpreter.open(); + } + + @AfterAll + public static void tearDown() throws InterpreterException { + if (interpreter != null) { + interpreter.close(); + } + } + + private static InterpreterContext getInterpreterContext() { + return InterpreterContext.builder() + .setNoteId("noteId") + .setParagraphId("paragraphId") + .setParagraphTitle("title") + .setAngularObjectRegistry(new AngularObjectRegistry(intpGroup.getId(), null)) + .setResourcePool(new LocalResourcePool("id")) + .setInterpreterOut(new InterpreterOutput()) + .setIntpEventClient(mock(RemoteInterpreterEventClient.class)) + .build(); + } + + @Test + void testSparkSessionCreated() { + assertNotNull(interpreter.getSparkSession()); + } + + @Test + void testSimpleQuery() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "SELECT 1 AS id, 'hello' AS message", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("id")); + assertTrue(output.contains("message")); + assertTrue(output.contains("hello")); + } + + @Test + void testMultipleStatements() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "SELECT 1 AS a; SELECT 2 AS b", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + } + + @Test + void testInvalidSQL() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "SELECT FROM WHERE INVALID", context); + assertEquals(InterpreterResult.Code.ERROR, result.code()); + } + + @Test + void testMaxResultLimit() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + context.getLocalProperties().put("limit", "5"); + InterpreterResult result = interpreter.interpret( + "SELECT id FROM range(100)", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + String[] lines = output.split("\n"); + // header + 5 data rows + assertTrue(lines.length <= 7); + } + + @Test + void testDDL() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = interpreter.interpret( + "CREATE OR REPLACE TEMP VIEW test_view AS SELECT 1 AS id, 'test' AS name", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + + context = getInterpreterContext(); + result = interpreter.interpret("SELECT * FROM test_view", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("test")); + } + + @Test + void testFormType() { + assertEquals(Interpreter.FormType.SIMPLE, interpreter.getFormType()); + } + + @Test + void testProgress() throws InterpreterException { + InterpreterContext context = getInterpreterContext(); + assertEquals(0, interpreter.getProgress(context)); + } +} diff --git a/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectSqlInterpreterTest.java b/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectSqlInterpreterTest.java new file mode 100644 index 00000000000..acee3ed7a44 --- /dev/null +++ b/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectSqlInterpreterTest.java @@ -0,0 +1,164 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.apache.zeppelin.display.AngularObjectRegistry; +import org.apache.zeppelin.interpreter.Interpreter; +import org.apache.zeppelin.interpreter.InterpreterContext; +import org.apache.zeppelin.interpreter.InterpreterException; +import org.apache.zeppelin.interpreter.InterpreterGroup; +import org.apache.zeppelin.interpreter.InterpreterOutput; +import org.apache.zeppelin.interpreter.InterpreterResult; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.apache.zeppelin.resource.LocalResourcePool; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +import java.io.IOException; +import java.util.LinkedList; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +/** + * Integration tests for SparkConnectSqlInterpreter. + * Requires a running Spark Connect server. + * Set SPARK_CONNECT_TEST_REMOTE env var (e.g. sc://localhost:15002) to enable. + */ +@EnabledIfEnvironmentVariable(named = "SPARK_CONNECT_TEST_REMOTE", matches = ".+") +public class SparkConnectSqlInterpreterTest { + + private static SparkConnectInterpreter connectInterpreter; + private static SparkConnectSqlInterpreter sqlInterpreter; + private static InterpreterGroup intpGroup; + + @BeforeAll + public static void setUp() throws Exception { + String remote = System.getenv("SPARK_CONNECT_TEST_REMOTE"); + Properties p = new Properties(); + p.setProperty("spark.remote", remote); + p.setProperty("spark.app.name", "ZeppelinSparkConnectSqlTest"); + p.setProperty("zeppelin.spark.maxResult", "100"); + p.setProperty("zeppelin.spark.concurrentSQL", "true"); + p.setProperty("zeppelin.spark.concurrentSQL.max", "10"); + p.setProperty("zeppelin.spark.sql.stacktrace", "true"); + p.setProperty("zeppelin.spark.sql.interpolation", "false"); + + intpGroup = new InterpreterGroup(); + connectInterpreter = new SparkConnectInterpreter(p); + connectInterpreter.setInterpreterGroup(intpGroup); + + sqlInterpreter = new SparkConnectSqlInterpreter(p); + sqlInterpreter.setInterpreterGroup(intpGroup); + + intpGroup.put("session_1", new LinkedList()); + intpGroup.get("session_1").add(connectInterpreter); + intpGroup.get("session_1").add(sqlInterpreter); + + connectInterpreter.open(); + sqlInterpreter.open(); + } + + @AfterAll + public static void tearDown() throws InterpreterException { + if (sqlInterpreter != null) { + sqlInterpreter.close(); + } + if (connectInterpreter != null) { + connectInterpreter.close(); + } + } + + private static InterpreterContext getInterpreterContext() { + return InterpreterContext.builder() + .setNoteId("noteId") + .setParagraphId("paragraphId") + .setParagraphTitle("title") + .setAngularObjectRegistry(new AngularObjectRegistry(intpGroup.getId(), null)) + .setResourcePool(new LocalResourcePool("id")) + .setInterpreterOut(new InterpreterOutput()) + .setIntpEventClient(mock(RemoteInterpreterEventClient.class)) + .build(); + } + + @Test + void testSimpleQuery() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = sqlInterpreter.interpret( + "SELECT 1 AS id, 'hello' AS message", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("id")); + assertTrue(output.contains("message")); + assertTrue(output.contains("hello")); + } + + @Test + void testMultipleStatements() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = sqlInterpreter.interpret( + "SELECT 1 AS a; SELECT 2 AS b", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + assertEquals(2, context.out.toInterpreterResultMessage().size()); + } + + @Test + void testInvalidSQL() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = sqlInterpreter.interpret( + "SELECT FROM WHERE INVALID", context); + assertEquals(InterpreterResult.Code.ERROR, result.code()); + assertTrue(context.out.toInterpreterResultMessage().get(0).getData().length() > 0); + } + + @Test + void testMaxResultLimit() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + context.getLocalProperties().put("limit", "3"); + InterpreterResult result = sqlInterpreter.interpret( + "SELECT id FROM range(100)", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + String[] lines = output.split("\n"); + // header + 3 data rows + assertTrue(lines.length <= 5); + } + + @Test + void testCreateAndQuery() throws InterpreterException, IOException { + InterpreterContext context = getInterpreterContext(); + InterpreterResult result = sqlInterpreter.interpret( + "CREATE OR REPLACE TEMP VIEW sql_test AS SELECT 42 AS answer", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + + context = getInterpreterContext(); + result = sqlInterpreter.interpret("SELECT * FROM sql_test", context); + assertEquals(InterpreterResult.Code.SUCCESS, result.code()); + String output = context.out.toInterpreterResultMessage().get(0).getData(); + assertTrue(output.contains("42")); + } + + @Test + void testFormType() { + assertEquals(Interpreter.FormType.SIMPLE, sqlInterpreter.getFormType()); + } +} diff --git a/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectUtilsTest.java b/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectUtilsTest.java new file mode 100644 index 00000000000..ebafd1a9bef --- /dev/null +++ b/spark-connect/src/test/java/org/apache/zeppelin/spark/SparkConnectUtilsTest.java @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.zeppelin.spark; + +import org.junit.jupiter.api.Test; + +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class SparkConnectUtilsTest { + + @Test + void testBuildConnectionStringDefault() { + Properties props = new Properties(); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://localhost:15002", result); + } + + @Test + void testBuildConnectionStringCustomRemote() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://spark-server.example.com:15002"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://spark-server.example.com:15002", result); + } + + @Test + void testBuildConnectionStringWithToken() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://spark-server.example.com:15002"); + props.setProperty("spark.connect.token", "my-secret-token"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://spark-server.example.com:15002/;token=my-secret-token", result); + } + + @Test + void testBuildConnectionStringWithTokenAndExistingParams() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://host:15002/;use_ssl=true"); + props.setProperty("spark.connect.token", "tok123"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://host:15002/;use_ssl=true;token=tok123", result); + } + + @Test + void testBuildConnectionStringEmptyToken() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://host:15002"); + props.setProperty("spark.connect.token", ""); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://host:15002", result); + } + + @Test + void testBuildConnectionStringWithUseSsl() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://ranking-cluster-m:8080"); + props.setProperty("spark.connect.use_ssl", "true"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://ranking-cluster-m:8080/;use_ssl=true", result); + } + + @Test + void testBuildConnectionStringWithSslAndToken() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://cluster:8080"); + props.setProperty("spark.connect.use_ssl", "true"); + props.setProperty("spark.connect.token", "abc123"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://cluster:8080/;token=abc123;use_ssl=true", result); + } + + @Test + void testBuildConnectionStringDataprocTunnel() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://localhost:15002"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://localhost:15002", result); + } + + @Test + void testBuildConnectionStringDataprocDirect() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://ranking-cluster-m:8080"); + String result = SparkConnectUtils.buildConnectionString(props); + assertEquals("sc://ranking-cluster-m:8080", result); + } + + @Test + void testBuildConnectionStringWithUserName() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://cluster:8080"); + String result = SparkConnectUtils.buildConnectionString(props, "alice"); + assertEquals("sc://cluster:8080/;user_id=alice", result); + } + + @Test + void testBuildConnectionStringWithUserNameAndToken() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://cluster:8080"); + props.setProperty("spark.connect.token", "tok"); + String result = SparkConnectUtils.buildConnectionString(props, "bob"); + assertEquals("sc://cluster:8080/;token=tok;user_id=bob", result); + } + + @Test + void testBuildConnectionStringUserIdAlreadyInUrl() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://cluster:8080/;user_id=preexisting"); + String result = SparkConnectUtils.buildConnectionString(props, "alice"); + assertEquals("sc://cluster:8080/;user_id=preexisting", result); + } + + @Test + void testBuildConnectionStringNullUserName() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://cluster:8080"); + String result = SparkConnectUtils.buildConnectionString(props, null); + assertEquals("sc://cluster:8080", result); + } + + @Test + void testBuildConnectionStringBlankUserName() { + Properties props = new Properties(); + props.setProperty("spark.remote", "sc://cluster:8080"); + String result = SparkConnectUtils.buildConnectionString(props, " "); + assertEquals("sc://cluster:8080", result); + } + + @Test + void testReplaceReservedChars() { + assertEquals("hello world", SparkConnectUtils.replaceReservedChars("hello\tworld")); + assertEquals("hello world", SparkConnectUtils.replaceReservedChars("hello\nworld")); + assertEquals("null", SparkConnectUtils.replaceReservedChars(null)); + assertEquals("normal", SparkConnectUtils.replaceReservedChars("normal")); + assertEquals("a b c", SparkConnectUtils.replaceReservedChars("a\tb\nc")); + } +}