Cluster environment variables not visible to Java Dataset.foreachPartition / mapPartitions

[capoorhimanshu]

We run a Spark JAR task (Java) on a standard access mode cluster. Credentials come from cluster environment
variables set as secret references in the cluster's Environment variables field, e.g.
MY_PASSWORD={{secrets/<scope>/<key>}}.

The variables are visible to the JAR task's driver code, but not to code running inside Dataset.foreachPartition /
mapPartitions on the workers (we only check presence, never values):

- DBR 16.4 LTS (Scala 2.12), standard access mode: inside foreachPartition the code runs as process "3@sandbox" and sees only 27 variables (DBR_VERSION, DB_*, HOME, LANG, PATH, PWD, SHLVL, payload_format, stageId, taskAttemptId). None of the cluster-configured variables are present, including a plain non-secret one. On the driver all are present.

- DBR 19 (Scala 2.13, Spark 4.2), standard access mode: driver sees MY_PASSWORD, mapPartitions code does not. The cluster edit page says "Databricks Runtime 19 and above restricts environment variables. The following
variables will be available in user code, but not loaded into the Spark session", listing our variables.

The standard compute limitations page says variables outside the predefined set "remain available to your user code,including UDFs", which does not match what we see.


Minimal reproduction (Java JAR task, MY_PASSWORD set as a cluster environment variable):

SparkSession spark = SparkSession.builder().getOrCreate();
boolean onDriver = System.getenv("MY_PASSWORD") != null; // true
List<String> onWorkers = spark.range(0, 2, 1, 2)
.mapPartitions((MapPartitionsFunction<Long, String>) rows -> Collections.singletonList(
String.valueOf(System.getenv("MY_PASSWORD") != null)).iterator(), Encoders.STRING())
.collectAsList();


Questions:
1. Is this expected for Scala/Java Dataset.foreachPartition / mapPartitions code on standard access mode?
2. Is there a supported way to make specific cluster environment variables visible to that code?

[SumeshKashyap]

Hello @capoorhimanshu 

On standard access mode this is expected. Scala/Java closures (foreachPartition, mapPartitions, UDFs) don't run in the executor JVM where the cluster env vars are set. They run in an isolated per-user sandbox process with a fixed, minimal environment. That's the 3@sandbox process and the 27 vars you're seeing. On DBR 19 spark.executorEnv.* is also restricted on standard, so you can't push the vars that way. The supported fix is to read the value on the driver and pass it into the closure, or to run the job on dedicated access mode.

Option 1: capture on the driver. This works on standard mode, DBR 16.4 and 19:

final String password = System.getenv("MY_PASSWORD");   // driver: present

ds.foreachPartition((ForeachPartitionFunction<Row>) rows -> {

    // 'password' is serialized with the task closure

    try (Connection c = DriverManager.getConnection(url, user, password)) {

        while (rows.hasNext()) { /* ... */ }

    }

});

Keep it in a local final variable (not a static field), or the sandbox will just see null. It travels inside the serialized task, so don't log or toString() the closure. Spark's redaction won't touch it there.

Option 2: dedicated (single user) access mode. Executors there run your code in their own JVM, which gets the cluster's environment variables, so System.getenv("MY_PASSWORD") works inside mapPartitions unchanged. This is the cleaner choice if the job runs as a service principal anyway.

About the docs line saying non-predefined variables "remain available to your user code, including UDFs": your repro shows that doesn't hold for JVM closures on standard. It's worth sending docs feedback on https://docs.databricks.com/aws/en/compute/standard-limitations so it gets clarified.

If this helps you and solves your problem , kindly accept them as answer so that community can be benefited

[capoorhimanshu]

Thanks for the response. Yes, that matches what I’m seeing and explains the sandbox process.

I was wondering if there’s a way to configure the shared cluster so that environment variables can be passed through to these sandbox workers.

Capturing the value on the driver is something I wanted to avoid, since on a shared cluster it would mean the secret/password could end up in the serialized task. The functionality already works on dedicated clusters; I was looking to extend it to shared clusters as well.

[saurabh18cs]

Hi  @capoorhimanshu 

the page does not explicitly say environment variables can be injected into or accessed from inside foreachPartition using standard compute.

[Isaac_18]

Hi @capoorhimanshu,

Your worry about the closure is fair. A secret that rides inside the serialized task is harder to control than one that never leaves a secret store.

As far as I can tell, there's no supported setting on standard access mode that pushes cluster environment variables into the sandbox where your Java closures run. The isolation is deliberate. On DBR 19, the docs also list spark.executorEnv.* among the Spark config properties you can't set on standard access mode, so that route is closed on purpose too. You're right that the "remain available to your user code, including UDFs" line on the limitations page doesn't match what you're seeing for JVM closures. It's worth sending docs feedback on that page so it gets clarified.

So the options are really about where the credential lives:

1. If it's a database password: use a JDBC Unity Catalog connection

The connection stores the URL and credentials in Unity Catalog, and users only need USE CONNECTION. Your code never sees the password, so nothing ends up in a closure. It works on standard access mode on DBR 17.3 LTS and above (Public Preview through 18, GA from DBR 19).

CREATE CONNECTION my_pg_conn TYPE JDBC
ENVIRONMENT (java_dependencies '["/Volumes/main/jars/drivers/postgresql-42.7.4.jar"]')
OPTIONS (
  url 'jdbc:postgresql://db-host:5432/appdb',
  user '<db_user>',
  password '<db_password>',
  externalOptionsAllowList 'dbtable,query'
);

GRANT USE CONNECTION ON CONNECTION my_pg_conn TO `data_engineers`;
-- users of the connection also need to read the driver JAR
GRANT READ VOLUME ON VOLUME main.jars.drivers TO `data_engineers`;

ds.write()
  .format("jdbc")
  .option("databricks.connection", "my_pg_conn")
  .option("dbtable", "public.target_table")
  .mode("append")
  .save();

The catch is that it's a DataFrame read/write (append or overwrite), not custom per-row logic. Row-level UPDATE/DELETE through the connection isn't supported. If your foreachPartition is doing something like upserts, a common pattern is to append into a staging table on the database side and run the upsert there. The connection owner sets the password once. After that, users can't read it back or override the connection's options at query time.

2. If you need custom code inside the partitions: dedicated compute assigned to a group

You already found that dedicated works. Dedicated access mode can also be assigned to a group (Public Preview, DBR 15.4+), so several people or service principals can share one cluster and the executors still get the cluster's environment variables. Everything on that cluster runs with the group's permissions, not each person's, so the group needs READ on the secret scope and the table grants. That's usually fine for a team-owned job cluster, but it's worth knowing before you move interactive users onto it.

3. Keep a narrow fallback

If neither fits, passing the value from the driver is still the only way into the sandbox. I'd at least keep the credential short-lived (for example, a token you fetch on the driver at job start, not a long-lived password) and limit who can read the job's logs and event logs.

Docs:

• Standard compute limitations: https://docs.databricks.com/aws/en/compute/standard-limitations
• JDBC connection: https://docs.databricks.com/aws/en/connect/jdbc-connection
• Dedicated compute group access: https://docs.databricks.com/aws/en/compute/group-access