This is an automated email from the ASF dual-hosted git repository.
roryqi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 4cccb1327e [#12569] feat(spark-connector): add token auth and key
caches on identity (#12570)
4cccb1327e is described below
commit 4cccb1327e1a082ecf59e0f0ae31e251a1e17a44
Author: Mark Hoerth <[email protected]>
AuthorDate: Sun Aug 23 04:32:34 2026 -0700
[#12569] feat(spark-connector): add token auth and key caches on identity
(#12570)
### What changes were proposed in this pull request?
Two changes to the Spark connector that only make sense together.
**1. Give the caches an eviction policy.** `GravitinoCatalogManager`'s
catalog cache was built with a bare `Caffeine.newBuilder().build()`, so
it had no bound and no expiry. It now expires by write, controlled by
`spark.sql.gravitino.catalogCacheTtlSec`.
**2. Add a `token` auth type, and key the caches on the identity that
token carries.** The single `GravitinoClient` built at driver init
becomes a per-identity cache of clients, bounded and expiring by access,
closing each client on eviction so it does not leak its HTTP connection
pool. The catalog cache key becomes `identity.key() + ":" +
catalogName`.
The token is resolved on every request rather than captured once,
reading the active Spark session's configuration in preference to the
application's, and `spark.sql.gravitino.token.file` in preference to
`spark.sql.gravitino.token.value`. Identity comes from the JWT claim
named by `spark.sql.gravitino.token.principalFields`, defaulting to
`sub`. The signature is not verified, because validating the token is
the server's job and this value only partitions a cache. Opaque, non-JWT
tokens fall back to a hash of the token.
`GravitinoCatalogManager` stays a singleton and
`getGravitinoCatalogInfo(String)` keeps its signature, so `BaseCatalog`
is untouched.
One note for anyone touching shading later: `GravitinoIdentity` parses
the JWT payload with Jackson, and it is the first class in
`spark-common` to do so. The `spark-runtime` jars exclude
`com.fasterxml.jackson` from relocation, so this resolves against the
Jackson that Spark puts on the driver classpath rather than the
relocated copy bundled via `client-java-runtime`. The surface used is
`ObjectMapper.readTree` plus `JsonNode.get/isNull/asText`, stable across
the Jackson versions shipped by Spark 3.3 through 3.5. I verified it by
running the class from the shaded 3.3 runtime jar against only Spark
3.3-era Jackson, with a negative control that fails with
`NoClassDefFoundError` when Jackson is removed.
### Why are the changes needed?
The half that bites today, under any auth type, is the missing eviction.
`loadCatalog` is a REST call the server authorizes; once it has
succeeded, the result is served for the remaining life of the driver and
no further request reaches Gravitino. A grant revoked in Gravitino
therefore keeps being honoured by a running Spark application.
The identity-blind cache key is not exploitable on main, because
`simple`, `basic`, `oauth2` and `kerberos` all resolve exactly one
credential per application, so there is never a second identity to
confuse. It arms the moment a per-user credential mode exists. The
cached value is not inert metadata: `GravitinoMetalake.loadCatalog`
returns `DTOConverters.toCatalog(name, dto, restClient)`, the returned
`RelationalCatalog` stores that client in `BaseSchemaCatalog`'s
`protected final RESTClient restClient`, `asTableCatalog()` returns
`this`, and `listTables` issues its request on that same field. So a
session served another session's cached entry would transmit under the
other session's credential.
That is the argument for landing both halves in one PR. Eviction without
an identity-aware key still mixes users once per-user tokens exist. An
identity-aware key without eviction still serves revoked grants.
There is also no existing way to present a bearer token obtained outside
Spark, which is the gap #11181 describes, and adding one is precisely
what makes the cache key defect live.
Fix: #12569
Related: #11181 and its PR #11182 ask for the bearer-token-file half of
this.
### Does this PR introduce _any_ user-facing change?
New property keys, all optional, no change to existing behaviour:
- `spark.sql.gravitino.authType=token`, a fifth valid value alongside
the existing four.
- `spark.sql.gravitino.token.value` and
`spark.sql.gravitino.token.file`, the file taking precedence.
- `spark.sql.gravitino.token.principalFields`, default `sub`, intended
to match the server's `gravitino.authenticator.oauth.principalFields`.
- `spark.sql.gravitino.clientCacheMaxSize` (100),
`spark.sql.gravitino.clientCacheTtlSec` (3600),
`spark.sql.gravitino.catalogCacheTtlSec` (300).
The property suffixes are added to `AuthProperties` next to
`basic.username` and `oauth2.serverUri`, namespaced by auth type, so the
Trino connector can reuse them.
`simple`, `basic`, `oauth2` and `kerberos` keep a single
application-wide identity and are unaffected, other than the catalog
cache now expiring, which is the intended fix.
Two limits worth stating rather than hiding, both documented in
`docs/spark-connector/spark-connector.md`:
- The catalog list registered with Spark at driver startup is resolved
with the application's identity, so a user may see a catalog name they
are not then allowed to open.
- This governs metadata resolution only. Executors read data with the
credentials the underlying catalog was built with.
### How was this patch tested?
`./gradlew :spark-connector:spark-common:test -PskipITs` is green, 124
tests.
New `TestGravitinoCatalogManager` covers seven cases: two subjects with
different `sub` build two clients and issue two loads; the same `sub`
shares one client and one load; opaque non-JWT tokens partition by token
value; `authType=simple` keeps one identity across two sessions;
`close()` closes every cached client; exceeding `clientCacheMaxSize`
evicts and closes; and a catalog entry past `catalogCacheTtlSec` is
reloaded rather than served stale.
The first of those is the regression test for the cache key. I checked
that it actually bites by reverting `cacheKey` to return the catalog
name alone, which fails four of the seven, then restored it.
New cases in `TestGravitinoDriverPlugin` cover the provider: `Bearer
<token>` as UTF-8 bytes, re-resolution per request so a changed value is
picked up, `token.file` taking precedence and being re-read after a
rewrite, and a missing token failing with a message naming both
properties.
Separately, to confirm that a cached `Catalog` really does carry the
client that loaded it rather than re-resolving one, I ran a throwaway
probe against the `client-java` MockServer harness: two clients with
distinct bearer tokens, catalog loaded by A, handed to B through a plain
map, `listTables()` called by B. The request went out as `Bearer
TOKEN-ALICE`, while the control of B calling through its own client went
out as `Bearer TOKEN-BOB`. That probe is not included here, since it
tests client behaviour that is working as designed.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_01WAqUbBPofTzWbFYD8EGu1R
Co-authored-by: Mark Hoerth <[email protected]>
Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
---
.../org/apache/gravitino/auth/AuthProperties.java | 20 ++
docs/spark-connector/spark-connector.md | 30 +++
.../spark/connector/GravitinoSparkConfig.java | 18 ++
.../connector/catalog/GravitinoCatalogManager.java | 176 +++++++++++--
.../spark/connector/catalog/GravitinoIdentity.java | 141 ++++++++++
.../connector/plugin/GravitinoDriverPlugin.java | 86 ++++++-
.../catalog/TestBaseCatalogAuthorization.java | 4 +-
.../catalog/TestGravitinoCatalogManager.java | 283 +++++++++++++++++++++
.../connector/glue/TestGravitinoGlueCatalog.java | 3 +-
.../plugin/TestGravitinoDriverPlugin.java | 89 +++++++
10 files changed, 828 insertions(+), 22 deletions(-)
diff --git
a/catalogs/catalog-common/src/main/java/org/apache/gravitino/auth/AuthProperties.java
b/catalogs/catalog-common/src/main/java/org/apache/gravitino/auth/AuthProperties.java
index 160b9d29b5..f2f6e0d069 100644
---
a/catalogs/catalog-common/src/main/java/org/apache/gravitino/auth/AuthProperties.java
+++
b/catalogs/catalog-common/src/main/java/org/apache/gravitino/auth/AuthProperties.java
@@ -28,6 +28,7 @@ public class AuthProperties {
public static final String BASIC_AUTH_TYPE = "basic";
public static final String OAUTH2_AUTH_TYPE = "oauth2";
public static final String KERBEROS_AUTH_TYPE = "kerberos";
+ public static final String TOKEN_AUTH_TYPE = "token";
/** The configuration key for the built-in IdP username used in Basic
authentication. */
public static final String GRAVITINO_BASIC_USERNAME = "basic.username";
@@ -48,6 +49,21 @@ public class AuthProperties {
/** The configuration key for the scope of the token. */
public static final String GRAVITINO_OAUTH2_SCOPE = "oauth2.scope";
+ // token
+ /** The configuration key for the bearer token presented in Token
authentication. */
+ public static final String GRAVITINO_TOKEN_VALUE = "token.value";
+
+ /**
+ * The configuration key for the file holding the bearer token presented in
Token authentication.
+ */
+ public static final String GRAVITINO_TOKEN_FILE = "token.file";
+
+ /**
+ * The configuration key for the ordered claim names identifying the caller
carried in the bearer
+ * token. Mirrors the server's {@code
gravitino.authenticator.oauth.principalFields}.
+ */
+ public static final String GRAVITINO_TOKEN_PRINCIPAL_FIELDS =
"token.principalFields";
+
public static boolean isKerberos(String authType) {
return KERBEROS_AUTH_TYPE.equalsIgnoreCase(authType);
}
@@ -64,5 +80,9 @@ public class AuthProperties {
return BASIC_AUTH_TYPE.equalsIgnoreCase(authType);
}
+ public static boolean isToken(String authType) {
+ return TOKEN_AUTH_TYPE.equalsIgnoreCase(authType);
+ }
+
private AuthProperties() {}
}
diff --git a/docs/spark-connector/spark-connector.md
b/docs/spark-connector/spark-connector.md
index 5a3d8562c1..704a62e469 100644
--- a/docs/spark-connector/spark-connector.md
+++ b/docs/spark-connector/spark-connector.md
@@ -34,6 +34,13 @@ The Apache Gravitino Spark connector leverages the Spark
DataSourceV2 interface
| spark.sql.gravitino.enableIcebergSupport | string | `false` | Set to
`true` to use Iceberg catalog.
| No |
| spark.sql.gravitino.enablePaimonSupport | string | `false` | Set to
`true` to use Paimon catalog.
| No |
| spark.sql.gravitino.client. | string | (none) | The
configuration key prefix for the Gravitino client config.
| No |
+| spark.sql.gravitino.authType | string | `simple` | The
authentication type, one of `simple`, `basic`, `oauth2`, `kerberos`, `token`.
| No |
+| spark.sql.gravitino.token.value | string | (none) | The
bearer token to present to Gravitino, used when `authType` is `token`.
| No |
+| spark.sql.gravitino.token.file | string | (none) | Path to
a file holding the bearer token. Takes precedence over `token.value` and is
re-read per request. | No |
+| spark.sql.gravitino.token.principalFields | string | `sub` | Comma
separated, ordered JWT claim names used to identify the caller in `token` mode.
| No |
+| spark.sql.gravitino.clientCacheMaxSize | int | `100` | The
maximum number of Gravitino clients cached, one per identity.
| No |
+| spark.sql.gravitino.clientCacheTtlSec | long | `3600` | Evicts a
cached Gravitino client this many seconds after it was last used.
| No |
+| spark.sql.gravitino.catalogCacheTtlSec | long | `300` | Evicts a
cached catalog this many seconds after it was loaded.
| No |
To configure the Gravitino client, use properties prefixed with
`spark.sql.gravitino.client.`. These properties will be passed to the Gravitino
client after removing the `spark.sql.` prefix.
@@ -41,6 +48,29 @@ To configure the Gravitino client, use properties prefixed
with `spark.sql.gravi
**Note:** Invalid configuration properties will result in exceptions. Please
see [Gravitino Java client
configurations](../how-to-use-gravitino-client.md#java-client-configuration)
for more support client configuration.
+### Per-user identity in `token` mode
+
+When `spark.sql.gravitino.authType` is `token`, the connector presents the
bearer token found in
+`spark.sql.gravitino.token.file`, or failing that
`spark.sql.gravitino.token.value`, resolving it again on
+every request and reading the active Spark session's configuration in
preference to the
+application's. The identity of the caller is taken from the JWT claim named by
+`spark.sql.gravitino.token.principalFields`, which should be set to match the
server's
+`gravitino.authenticator.oauth.principalFields` so that the connector and the
server agree on who
+is asking. The claim is read without verifying the signature, because
validating the token remains
+the server's job; the value is used only to partition the connector's caches.
An opaque,
+non-JWT token is partitioned by a hash of the token instead.
+
+Gravitino clients and the catalog metadata they load are cached per identity,
so two users sharing
+one Spark driver never see each other's cached metadata. All other
authentication types keep a
+single application-wide identity and are unaffected by any of this.
+
+Two limits are worth stating plainly:
+
+* The list of catalogs registered with Spark at driver startup is resolved
with the application's
+ identity, so a user may see a catalog name that they are not then allowed to
open.
+* This governs metadata resolution only. Data is read by the executors using
the credentials that
+ the underlying catalog was built with, not the end user's token.
+
```shell
./bin/spark-sql -v \
--conf
spark.plugins="org.apache.gravitino.spark.connector.plugin.GravitinoSparkPlugin"
\
diff --git
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/GravitinoSparkConfig.java
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/GravitinoSparkConfig.java
index be108c9af1..ed82ed23d9 100644
---
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/GravitinoSparkConfig.java
+++
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/GravitinoSparkConfig.java
@@ -49,6 +49,24 @@ public class GravitinoSparkConfig {
public static final String GRAVITINO_KERBEROS_PRINCIPAL =
"spark.kerberos.principal";
public static final String GRAVITINO_KERBEROS_KEYTAB_FILE_PATH =
"spark.kerberos.keytab";
+ public static final String GRAVITINO_TOKEN_VALUE =
+ GRAVITINO_PREFIX + AuthProperties.GRAVITINO_TOKEN_VALUE;
+ public static final String GRAVITINO_TOKEN_FILE =
+ GRAVITINO_PREFIX + AuthProperties.GRAVITINO_TOKEN_FILE;
+ public static final String GRAVITINO_TOKEN_PRINCIPAL_FIELDS =
+ GRAVITINO_PREFIX + AuthProperties.GRAVITINO_TOKEN_PRINCIPAL_FIELDS;
+ public static final String GRAVITINO_TOKEN_PRINCIPAL_FIELDS_DEFAULT = "sub";
+
+ public static final String GRAVITINO_CLIENT_CACHE_MAX_SIZE =
+ GRAVITINO_PREFIX + "clientCacheMaxSize";
+ public static final int GRAVITINO_CLIENT_CACHE_MAX_SIZE_DEFAULT = 100;
+ public static final String GRAVITINO_CLIENT_CACHE_TTL_SEC =
+ GRAVITINO_PREFIX + "clientCacheTtlSec";
+ public static final long GRAVITINO_CLIENT_CACHE_TTL_SEC_DEFAULT = 3600;
+ public static final String GRAVITINO_CATALOG_CACHE_TTL_SEC =
+ GRAVITINO_PREFIX + "catalogCacheTtlSec";
+ public static final long GRAVITINO_CATALOG_CACHE_TTL_SEC_DEFAULT = 300;
+
public static final String GRAVITINO_HIVE_METASTORE_URI = "metastore.uris";
public static final String SPARK_HIVE_METASTORE_URI = "hive.metastore.uris";
diff --git
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
index a484972656..21b7c520c8 100644
---
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
+++
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
@@ -20,34 +20,110 @@ package org.apache.gravitino.spark.connector.catalog;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
-import com.google.common.base.Supplier;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import java.time.Duration;
import java.util.Arrays;
+import java.util.List;
import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
import org.apache.gravitino.Catalog;
+import org.apache.gravitino.auth.AuthProperties;
import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.spark.connector.ConnectorConstants;
+import org.apache.gravitino.spark.connector.GravitinoSparkConfig;
+import org.apache.gravitino.spark.connector.plugin.GravitinoDriverPlugin;
+import org.apache.spark.SparkConf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-/** GravitinoCatalogManager is used to retrieve catalogs from Apache Gravitino
server. */
+/**
+ * GravitinoCatalogManager is used to retrieve catalogs from Apache Gravitino
server.
+ *
+ * <p>The manager itself stays a JVM-wide singleton, but the clients inside it
are cached per {@link
+ * GravitinoIdentity} so that a shared Spark driver does not serve one user's
catalog metadata to
+ * another. Every auth type other than {@code token} resolves to a single
application-wide identity,
+ * which keeps those deployments behaving exactly as before.
+ */
public class GravitinoCatalogManager {
private static final Logger LOG =
LoggerFactory.getLogger(GravitinoCatalogManager.class);
private static GravitinoCatalogManager gravitinoCatalogManager;
private volatile boolean isClosed = false;
+ private final SparkConf sparkConf;
+ private final String authType;
+ private final List<String> principalFields;
+ private final Function<GravitinoIdentity, GravitinoClient> clientBuilder;
+ private final Cache<GravitinoIdentity, GravitinoClient> clients;
private final Cache<String, Catalog> gravitinoCatalogs;
- private final GravitinoClient gravitinoClient;
+ private volatile Map<String, Catalog> applicationCatalogs =
ImmutableMap.of();
- private GravitinoCatalogManager(Supplier<GravitinoClient> clientBuilder) {
- this.gravitinoClient = clientBuilder.get();
- // Will not evict catalog by default
- this.gravitinoCatalogs = Caffeine.newBuilder().build();
+ private GravitinoCatalogManager(
+ SparkConf sparkConf, Function<GravitinoIdentity, GravitinoClient>
clientBuilder) {
+ this.sparkConf = sparkConf;
+ this.clientBuilder = clientBuilder;
+ this.authType =
+ sparkConf.get(GravitinoSparkConfig.GRAVITINO_AUTH_TYPE,
AuthProperties.SIMPLE_AUTH_TYPE);
+ this.principalFields =
+ Arrays.stream(
+ sparkConf
+ .get(
+ GravitinoSparkConfig.GRAVITINO_TOKEN_PRINCIPAL_FIELDS,
+
GravitinoSparkConfig.GRAVITINO_TOKEN_PRINCIPAL_FIELDS_DEFAULT)
+ .split(ConnectorConstants.COMMA))
+ .map(String::trim)
+ .filter(field -> !field.isEmpty())
+ .collect(ImmutableList.toImmutableList());
+ this.clients =
+ Caffeine.newBuilder()
+ .maximumSize(
+ sparkConf.getInt(
+ GravitinoSparkConfig.GRAVITINO_CLIENT_CACHE_MAX_SIZE,
+
GravitinoSparkConfig.GRAVITINO_CLIENT_CACHE_MAX_SIZE_DEFAULT))
+ .expireAfterAccess(
+ Duration.ofSeconds(
+ sparkConf.getLong(
+ GravitinoSparkConfig.GRAVITINO_CLIENT_CACHE_TTL_SEC,
+
GravitinoSparkConfig.GRAVITINO_CLIENT_CACHE_TTL_SEC_DEFAULT)))
+ // An evicted client still owns an HTTP connection pool, so it
must be closed.
+ .<GravitinoIdentity, GravitinoClient>removalListener(
+ (identity, client, cause) -> closeClient(identity, client))
+ .build();
+ this.gravitinoCatalogs =
+ Caffeine.newBuilder()
+ // A permission revoked in Gravitino must eventually stop being
served from cache.
+ .expireAfterWrite(
+ Duration.ofSeconds(
+ sparkConf.getLong(
+ GravitinoSparkConfig.GRAVITINO_CATALOG_CACHE_TTL_SEC,
+
GravitinoSparkConfig.GRAVITINO_CATALOG_CACHE_TTL_SEC_DEFAULT)))
+ .build();
}
- public static GravitinoCatalogManager create(Supplier<GravitinoClient>
clientBuilder) {
+ /**
+ * Creates the singleton GravitinoCatalogManager.
+ *
+ * @param sparkConf the application Spark configuration, used to read the
cache settings and, in
+ * {@code token} mode, to resolve the bearer token of the current request
+ * @param applicationUser the user the Spark application runs as, recorded
for diagnostics only
+ * and never used to derive an identity
+ * @param clientBuilder builds a Gravitino client for a given identity
+ * @return the created GravitinoCatalogManager
+ */
+ public static GravitinoCatalogManager create(
+ SparkConf sparkConf,
+ String applicationUser,
+ Function<GravitinoIdentity, GravitinoClient> clientBuilder) {
Preconditions.checkState(
gravitinoCatalogManager == null, "Should not create duplicate
GravitinoCatalogManager");
- gravitinoCatalogManager = new GravitinoCatalogManager(clientBuilder);
+ gravitinoCatalogManager = new GravitinoCatalogManager(sparkConf,
clientBuilder);
+ LOG.info(
+ "Created GravitinoCatalogManager for Spark user {} with auth type {}.",
+ applicationUser,
+ gravitinoCatalogManager.authType);
return gravitinoCatalogManager;
}
@@ -62,35 +138,97 @@ public class GravitinoCatalogManager {
public void close() {
Preconditions.checkState(!isClosed, "Gravitino Catalog is already closed");
isClosed = true;
- gravitinoClient.close();
+ // Caffeine dispatches the removal listener asynchronously, so shutdown
closes explicitly.
+ clients.asMap().forEach(GravitinoCatalogManager::closeClient);
+ clients.invalidateAll();
+ gravitinoCatalogs.invalidateAll();
+ applicationCatalogs = ImmutableMap.of();
gravitinoCatalogManager = null;
}
public Catalog getGravitinoCatalogInfo(String name) {
+ GravitinoIdentity identity = currentIdentity();
try {
- return gravitinoCatalogs.get(name, catalogName ->
loadCatalog(catalogName));
+ return gravitinoCatalogs.get(cacheKey(identity, name), key ->
loadCatalog(identity, name));
} catch (Exception e) {
LOG.error(String.format("Load catalog %s failed", name), e);
throw new RuntimeException(e);
}
}
+ /**
+ * Loads the relational catalogs visible to the application identity. This
runs at driver init,
+ * before any Spark session exists, so it never consults session state.
+ */
public void loadRelationalCatalogs() {
- Catalog[] catalogs = gravitinoClient.listCatalogsInfo();
- Arrays.stream(catalogs)
- .filter(catalog -> Catalog.Type.RELATIONAL.equals(catalog.type()))
- .forEach(catalog -> gravitinoCatalogs.put(catalog.name(), catalog));
+ GravitinoIdentity identity = applicationIdentity();
+ Catalog[] catalogs = getClient(identity).listCatalogsInfo();
+ Map<String, Catalog> relationalCatalogs =
+ Arrays.stream(catalogs)
+ .filter(catalog -> Catalog.Type.RELATIONAL.equals(catalog.type()))
+ .collect(
+ Collectors.toMap(Catalog::name, catalog -> catalog, (first,
second) -> second));
+ relationalCatalogs.forEach(
+ (name, catalog) -> gravitinoCatalogs.put(cacheKey(identity, name),
catalog));
+ this.applicationCatalogs = ImmutableMap.copyOf(relationalCatalogs);
}
+ /**
+ * Returns the catalogs loaded by {@link #loadRelationalCatalogs()}, that
is, the catalogs the
+ * application identity can see.
+ *
+ * @return the catalogs registered at driver startup, keyed by catalog name
+ */
public Map<String, Catalog> getCatalogs() {
- return gravitinoCatalogs.asMap();
+ return applicationCatalogs;
+ }
+
+ /**
+ * Resolves the identity of the caller. Outside {@code token} mode this is
the application
+ * identity and no session state is consulted at all. In {@code token} mode
the bearer token is
+ * read exactly as the request-time provider reads it, active session
configuration first and
+ * application configuration second, and the identity is derived from that
token alone.
+ *
+ * @return the identity of the current caller
+ */
+ @VisibleForTesting
+ GravitinoIdentity currentIdentity() {
+ if (!AuthProperties.isToken(authType)) {
+ return applicationIdentity();
+ }
+ return GravitinoIdentity.fromToken(
+ GravitinoDriverPlugin.resolveToken(sparkConf), principalFields);
+ }
+
+ @VisibleForTesting
+ GravitinoClient getClient(GravitinoIdentity identity) {
+ return clients.get(identity, clientBuilder);
+ }
+
+ private GravitinoIdentity applicationIdentity() {
+ return GravitinoIdentity.application(authType);
}
- private Catalog loadCatalog(String catalogName) {
- Catalog catalog = gravitinoClient.loadCatalog(catalogName);
+ private Catalog loadCatalog(GravitinoIdentity identity, String catalogName) {
+ Catalog catalog = getClient(identity).loadCatalog(catalogName);
Preconditions.checkArgument(
Catalog.Type.RELATIONAL.equals(catalog.type()), "Only support
relational catalog");
- LOG.info("Load catalog {} from Gravitino successfully.", catalogName);
+ LOG.info("Load catalog {} from Gravitino successfully for {}.",
catalogName, identity);
return catalog;
}
+
+ private static String cacheKey(GravitinoIdentity identity, String
catalogName) {
+ return identity.key() + ":" + catalogName;
+ }
+
+ private static void closeClient(GravitinoIdentity identity, GravitinoClient
client) {
+ if (client == null) {
+ return;
+ }
+ try {
+ client.close();
+ } catch (Exception e) {
+ LOG.warn("Failed to close the Gravitino client of {}.", identity, e);
+ }
+ }
}
diff --git
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoIdentity.java
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoIdentity.java
new file mode 100644
index 0000000000..050ef15825
--- /dev/null
+++
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoIdentity.java
@@ -0,0 +1,141 @@
+/*
+ * 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.gravitino.spark.connector.catalog;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.hash.Hashing;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.List;
+import java.util.Objects;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The identity a Gravitino request is made on behalf of. It exists only to
partition the client and
+ * catalog caches held by {@link GravitinoCatalogManager}; it is never
presented to the Gravitino
+ * server and never used to make a trust decision.
+ *
+ * <p>The safety property this class must uphold is that the cache key is at
least as fine-grained
+ * as the principal the Gravitino server derives from the same credential, and
never exactly equal
+ * to it in the sense of being coarser. A key finer than the server's
principal costs a redundant
+ * cache entry and nothing more. A key coarser than the server's principal
would let one user be
+ * served another user's cached metadata, which is precisely the bug this
class exists to prevent.
+ * When in doubt, be finer.
+ *
+ * <p>For that reason the claim read out of a JWT is taken without verifying
the signature.
+ * Validating the token is the server's job, and a value used only to split a
cache must not be
+ * mistaken for a verified assertion about who is calling.
+ */
+public final class GravitinoIdentity {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(GravitinoIdentity.class);
+
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ /** Prefix of every token-derived key, so it can never collide with an
auth-type key. */
+ private static final String TOKEN_KEY_PREFIX = "token:";
+
+ private static final String HASHED_KEY_PREFIX = TOKEN_KEY_PREFIX + "sha256:";
+
+ private static final int JWT_PART_COUNT = 3;
+
+ private final String key;
+
+ private GravitinoIdentity(String key) {
+ this.key = key;
+ }
+
+ /**
+ * Returns the single identity shared by every auth type that resolves one
credential for the
+ * whole Spark application, namely {@code simple}, {@code basic}, {@code
oauth2} and {@code
+ * kerberos}. Those deployments keep exactly one identity per application.
+ *
+ * @param authType the configured auth type, used verbatim as the key
+ * @return the application-wide identity for that auth type
+ */
+ public static GravitinoIdentity application(String authType) {
+ return new GravitinoIdentity(String.valueOf(authType));
+ }
+
+ /**
+ * Derives an identity from a bearer token.
+ *
+ * <p>When the token is a three-part JWT whose payload parses, the first of
{@code
+ * principalFields} that is present and non-blank supplies the key.
Otherwise, and for opaque
+ * access tokens generally, the key falls back to the hex SHA-256 of the
token. The signature is
+ * never verified.
+ *
+ * @param token the bearer token, never logged and never stored on the
returned identity
+ * @param principalFields the ordered claim names to try, mirroring the
server's {@code
+ * gravitino.authenticator.oauth.principalFields}
+ * @return an identity suitable for use as a cache key
+ */
+ public static GravitinoIdentity fromToken(String token, List<String>
principalFields) {
+ String[] parts = token.split("\\.");
+ if (parts.length == JWT_PART_COUNT) {
+ try {
+ JsonNode payload =
OBJECT_MAPPER.readTree(Base64.getUrlDecoder().decode(parts[1].trim()));
+ for (String field : principalFields) {
+ JsonNode value = payload.get(field);
+ if (value != null && !value.isNull() &&
StringUtils.isNotBlank(value.asText())) {
+ return new GravitinoIdentity(TOKEN_KEY_PREFIX + field + ":" +
value.asText());
+ }
+ }
+ LOG.debug("No principal field of {} found in the token payload.",
principalFields);
+ } catch (Exception e) {
+ LOG.debug("Cannot parse the token payload as a JWT, falling back to a
token hash.", e);
+ }
+ }
+ return new GravitinoIdentity(
+ HASHED_KEY_PREFIX + Hashing.sha256().hashString(token,
StandardCharsets.UTF_8));
+ }
+
+ /**
+ * Returns the cache key of this identity. It never is, and never contains,
a raw token.
+ *
+ * @return the cache key
+ */
+ public String key() {
+ return key;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof GravitinoIdentity)) {
+ return false;
+ }
+ return Objects.equals(key, ((GravitinoIdentity) o).key);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(key);
+ }
+
+ @Override
+ public String toString() {
+ return "GravitinoIdentity{key=" + key + "}";
+ }
+}
diff --git
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/plugin/GravitinoDriverPlugin.java
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/plugin/GravitinoDriverPlugin.java
index 90e7830bff..c27f78dce4 100644
---
a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/plugin/GravitinoDriverPlugin.java
+++
b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/plugin/GravitinoDriverPlugin.java
@@ -27,6 +27,11 @@ import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import java.io.File;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -40,6 +45,7 @@ import javax.annotation.Nullable;
import org.apache.commons.lang3.StringUtils;
import org.apache.gravitino.Catalog;
import org.apache.gravitino.auth.AuthProperties;
+import org.apache.gravitino.client.CustomTokenProvider;
import org.apache.gravitino.client.DefaultOAuth2TokenProvider;
import org.apache.gravitino.client.GravitinoClient;
import org.apache.gravitino.client.GravitinoClient.ClientBuilder;
@@ -55,9 +61,11 @@ import org.apache.spark.SparkConf;
import org.apache.spark.SparkContext;
import org.apache.spark.api.plugin.DriverPlugin;
import org.apache.spark.api.plugin.PluginContext;
+import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.internal.StaticSQLConf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import scala.Option;
/**
* GravitinoDriverPlugin creates GravitinoCatalogManager to fetch catalogs
from Apache Gravitino and
@@ -115,7 +123,9 @@ public class GravitinoDriverPlugin implements DriverPlugin {
this.catalogManager =
GravitinoCatalogManager.create(
- () ->
+ conf,
+ sc.sparkUser(),
+ identity ->
createGravitinoClient(
gravitinoUri, metalake, conf, sc.sparkUser(),
gravitinoClientConfig));
catalogManager.loadRelationalCatalogs();
@@ -249,6 +259,8 @@ public class GravitinoDriverPlugin implements DriverPlugin {
.withKeyTabFile(new File(keyTabFile))
.build();
builder.withKerberosAuth(kerberosTokenProvider);
+ } else if (AuthProperties.isToken(authType)) {
+ builder.withCustomTokenAuth(new DynamicBearerTokenProvider(sparkConf));
} else {
throw new UnsupportedOperationException("Unsupported auth type: " +
authType);
}
@@ -267,6 +279,78 @@ public class GravitinoDriverPlugin implements DriverPlugin
{
return sparkConf.get(configKey, null);
}
+ /**
+ * Resolves the bearer token to present to Gravitino for the operation
running on the current
+ * thread.
+ *
+ * <p>The active Spark session's configuration wins over the application
configuration, so a
+ * shared driver can carry a different end user's token per session, and
{@code tokenFile} wins
+ * over {@code token}, so a token refreshed on disk is picked up without
touching the
+ * configuration. Neither the token nor any prefix of it is ever logged.
+ *
+ * @param sparkConf the application Spark configuration, used when no
session supplies a token
+ * @return the bearer token, never blank
+ * @throws IllegalArgumentException if neither token property is set
+ * @throws UncheckedIOException if the configured token file cannot be read
+ */
+ public static String resolveToken(SparkConf sparkConf) {
+ String tokenFile =
+ getSessionOrApplicationConfig(sparkConf,
GravitinoSparkConfig.GRAVITINO_TOKEN_FILE);
+ String token =
+ StringUtils.isNotBlank(tokenFile)
+ ? readTokenFile(tokenFile)
+ : getSessionOrApplicationConfig(sparkConf,
GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE);
+ Preconditions.checkArgument(
+ StringUtils.isNotBlank(token),
+ String.format(
+ "Either %s or %s should be set when %s is %s",
+ GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE,
+ GravitinoSparkConfig.GRAVITINO_TOKEN_FILE,
+ GravitinoSparkConfig.GRAVITINO_AUTH_TYPE,
+ AuthProperties.TOKEN_AUTH_TYPE));
+ return token.trim();
+ }
+
+ private static String readTokenFile(String tokenFile) {
+ try {
+ return new String(Files.readAllBytes(Paths.get(tokenFile)),
StandardCharsets.UTF_8).trim();
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to read the Gravitino token file
" + tokenFile, e);
+ }
+ }
+
+ @Nullable
+ private static String getSessionOrApplicationConfig(SparkConf sparkConf,
String configKey) {
+ Option<SparkSession> activeSession = SparkSession.getActiveSession();
+ if (activeSession.isDefined()) {
+ Option<String> sessionValue =
activeSession.get().conf().getOption(configKey);
+ if (sessionValue.isDefined() &&
StringUtils.isNotBlank(sessionValue.get())) {
+ return sessionValue.get();
+ }
+ }
+ return getOptionalConfig(sparkConf, configKey);
+ }
+
+ /**
+ * Presents a bearer token that is resolved again on every request rather
than captured once. See
+ * {@link #resolveToken(SparkConf)} for where the token comes from.
+ */
+ @VisibleForTesting
+ static final class DynamicBearerTokenProvider extends CustomTokenProvider {
+
+ private final SparkConf sparkConf;
+
+ DynamicBearerTokenProvider(SparkConf sparkConf) {
+ this.sparkConf = sparkConf;
+ this.schemeName = "Bearer";
+ }
+
+ @Override
+ protected String getCustomTokenInfo() {
+ return resolveToken(sparkConf);
+ }
+ }
+
@VisibleForTesting
public static Map<String, String> extractGravitinoClientConfig(SparkConf
conf) {
return
Optional.ofNullable(conf.getAllWithPrefix(GRAVITINO_CLIENT_CONFIG_PREFIX))
diff --git
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestBaseCatalogAuthorization.java
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestBaseCatalogAuthorization.java
index 37ef2c501a..a15e4dfdd4 100644
---
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestBaseCatalogAuthorization.java
+++
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestBaseCatalogAuthorization.java
@@ -39,6 +39,7 @@ import
org.apache.gravitino.spark.connector.PropertiesConverter;
import org.apache.gravitino.spark.connector.SparkTransformConverter;
import org.apache.gravitino.spark.connector.SparkTypeConverter;
import org.apache.gravitino.spark.connector.authorization.AuthorizationTable;
+import org.apache.spark.SparkConf;
import org.apache.spark.sql.connector.catalog.Identifier;
import org.apache.spark.sql.connector.catalog.Table;
import org.apache.spark.sql.util.CaseInsensitiveStringMap;
@@ -59,7 +60,8 @@ public class TestBaseCatalogAuthorization {
@BeforeAll
void initCatalogManager() {
- GravitinoCatalogManager.create(() -> mock(GravitinoClient.class));
+ GravitinoCatalogManager.create(
+ new SparkConf(false), "user", identity -> mock(GravitinoClient.class));
}
@AfterAll
diff --git
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestGravitinoCatalogManager.java
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestGravitinoCatalogManager.java
new file mode 100644
index 0000000000..3a8eec8f84
--- /dev/null
+++
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/catalog/TestGravitinoCatalogManager.java
@@ -0,0 +1,283 @@
+/*
+ * 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.gravitino.spark.connector.catalog;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Base64;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BooleanSupplier;
+import java.util.function.Function;
+import org.apache.gravitino.Catalog;
+import org.apache.gravitino.auth.AuthProperties;
+import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.spark.connector.GravitinoSparkConfig;
+import org.apache.spark.SparkConf;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Verifies that GravitinoCatalogManager partitions its client and catalog
caches by the identity
+ * carried in the bearer token, and that no other auth type changes behaviour.
+ */
+public class TestGravitinoCatalogManager {
+
+ private static final String CATALOG_NAME = "test_catalog";
+
+ private ClientFactory clientFactory;
+
+ @AfterEach
+ void closeManager() {
+ try {
+ GravitinoCatalogManager.get().close();
+ } catch (IllegalStateException e) {
+ // The test closed the manager itself.
+ }
+ }
+
+ @Test
+ void testDifferentSubjectsDoNotShareCatalogCache() {
+ SparkConf sparkConf = tokenConf();
+ GravitinoCatalogManager manager = createManager(sparkConf);
+
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("alice"));
+ Catalog aliceCatalog = manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("bob"));
+ Catalog bobCatalog = manager.getGravitinoCatalogInfo(CATALOG_NAME);
+
+ assertEquals(2, clientFactory.clientCount());
+ assertEquals(2, clientFactory.loadCount());
+ assertNotSame(aliceCatalog, bobCatalog);
+ }
+
+ @Test
+ void testSameSubjectSharesClientAndCatalogCache() {
+ SparkConf sparkConf = tokenConf();
+ GravitinoCatalogManager manager = createManager(sparkConf);
+
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("alice",
"first"));
+ Catalog first = manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("alice",
"second"));
+ Catalog second = manager.getGravitinoCatalogInfo(CATALOG_NAME);
+
+ assertEquals(1, clientFactory.clientCount());
+ assertEquals(1, clientFactory.loadCount());
+ assertSame(first, second);
+ }
+
+ @Test
+ void testOpaqueTokensArePartitionedByTokenValue() {
+ SparkConf sparkConf = tokenConf();
+ GravitinoCatalogManager manager = createManager(sparkConf);
+
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE,
"opaque-token-one");
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE,
"opaque-token-one");
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+
+ assertEquals(1, clientFactory.clientCount());
+ assertEquals(1, clientFactory.loadCount());
+
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE,
"opaque-token-two");
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+
+ assertEquals(2, clientFactory.clientCount());
+ assertEquals(2, clientFactory.loadCount());
+ }
+
+ @Test
+ void testSimpleAuthKeepsOneApplicationIdentity() {
+ SparkConf sparkConf = new SparkConf(false);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_AUTH_TYPE,
AuthProperties.SIMPLE_AUTH_TYPE);
+ GravitinoCatalogManager manager = createManager(sparkConf);
+
+ // Tokens are irrelevant outside token mode: both sessions must resolve to
one identity.
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("alice"));
+ Catalog first = manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("bob"));
+ Catalog second = manager.getGravitinoCatalogInfo(CATALOG_NAME);
+
+ assertEquals(1, clientFactory.clientCount());
+ assertEquals(1, clientFactory.loadCount());
+ assertSame(first, second);
+ }
+
+ @Test
+ void testCloseClosesEveryCachedClient() {
+ SparkConf sparkConf = tokenConf();
+ GravitinoCatalogManager manager = createManager(sparkConf);
+
+ for (String user : new String[] {"alice", "bob", "carol"}) {
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt(user));
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ }
+ assertEquals(3, clientFactory.clientCount());
+
+ manager.close();
+
+ assertEquals(3, clientFactory.closedCount());
+ }
+
+ @Test
+ void testClientCacheEvictsAndClosesEvictedClient() {
+ SparkConf sparkConf = tokenConf();
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_CLIENT_CACHE_MAX_SIZE, "2");
+ GravitinoCatalogManager manager = createManager(sparkConf);
+
+ for (String user : new String[] {"alice", "bob", "carol"}) {
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt(user));
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ }
+
+ assertEquals(3, clientFactory.clientCount());
+ // Caffeine evicts and dispatches the removal listener asynchronously, and
a read keeps the
+ // cache draining its buffers while we wait.
+ assertTrue(
+ await(
+ () -> {
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ return clientFactory.closedCount() >= 1;
+ }),
+ "Exceeding the client cache size should evict and close a client");
+ }
+
+ @Test
+ void testCatalogCacheEntryExpires() throws InterruptedException {
+ SparkConf sparkConf = tokenConf();
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_CATALOG_CACHE_TTL_SEC, "1");
+ GravitinoCatalogManager manager = createManager(sparkConf);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, jwt("alice"));
+
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+ assertEquals(1, clientFactory.loadCount());
+
+ Thread.sleep(1500);
+ manager.getGravitinoCatalogInfo(CATALOG_NAME);
+
+ assertEquals(2, clientFactory.loadCount(), "A stale catalog entry must be
reloaded");
+ // The client is cached independently of the catalog entry.
+ assertEquals(1, clientFactory.clientCount());
+ }
+
+ private GravitinoCatalogManager createManager(SparkConf sparkConf) {
+ clientFactory = new ClientFactory();
+ return GravitinoCatalogManager.create(sparkConf, "spark-user",
clientFactory);
+ }
+
+ private static SparkConf tokenConf() {
+ SparkConf sparkConf = new SparkConf(false);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_AUTH_TYPE,
AuthProperties.TOKEN_AUTH_TYPE);
+ return sparkConf;
+ }
+
+ /** Builds an unsigned three part JWT. Nothing verifies it, so no signing
key is needed. */
+ private static String jwt(String subject) {
+ return jwt(subject, "unused");
+ }
+
+ private static String jwt(String subject, String jwtId) {
+ return base64Url("{\"alg\":\"none\",\"typ\":\"JWT\"}")
+ + "."
+ + base64Url(String.format("{\"sub\":\"%s\",\"jti\":\"%s\"}", subject,
jwtId))
+ + ".signature";
+ }
+
+ private static String base64Url(String value) {
+ return Base64.getUrlEncoder()
+ .withoutPadding()
+ .encodeToString(value.getBytes(StandardCharsets.UTF_8));
+ }
+
+ private static boolean await(BooleanSupplier condition) {
+ for (int i = 0; i < 100; i++) {
+ if (condition.getAsBoolean()) {
+ return true;
+ }
+ try {
+ Thread.sleep(50);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+ return condition.getAsBoolean();
+ }
+
+ /** Hands out a distinct mock client per identity and counts what the
manager asks of it. */
+ private static class ClientFactory implements Function<GravitinoIdentity,
GravitinoClient> {
+
+ private final List<AtomicBoolean> closedFlags = new ArrayList<>();
+ private final AtomicInteger clients = new AtomicInteger();
+ private final AtomicInteger loads = new AtomicInteger();
+
+ @Override
+ public GravitinoClient apply(GravitinoIdentity identity) {
+ clients.incrementAndGet();
+ GravitinoClient client = mock(GravitinoClient.class);
+ when(client.loadCatalog(anyString()))
+ .thenAnswer(
+ invocation -> {
+ loads.incrementAndGet();
+ Catalog catalog = mock(Catalog.class);
+ when(catalog.type()).thenReturn(Catalog.Type.RELATIONAL);
+ when(catalog.name()).thenReturn(invocation.getArgument(0));
+ return catalog;
+ });
+ // Closing twice must not be counted twice: the shutdown path closes
explicitly and the
+ // removal listener may then fire for the same client.
+ AtomicBoolean closed = new AtomicBoolean(false);
+ synchronized (closedFlags) {
+ closedFlags.add(closed);
+ }
+ doAnswer(
+ invocation -> {
+ closed.set(true);
+ return null;
+ })
+ .when(client)
+ .close();
+ return client;
+ }
+
+ int clientCount() {
+ return clients.get();
+ }
+
+ int loadCount() {
+ return loads.get();
+ }
+
+ int closedCount() {
+ synchronized (closedFlags) {
+ return (int) closedFlags.stream().filter(AtomicBoolean::get).count();
+ }
+ }
+ }
+}
diff --git
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/glue/TestGravitinoGlueCatalog.java
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/glue/TestGravitinoGlueCatalog.java
index af980f7309..364dcc74e2 100644
---
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/glue/TestGravitinoGlueCatalog.java
+++
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/glue/TestGravitinoGlueCatalog.java
@@ -43,6 +43,7 @@ import
org.apache.gravitino.spark.connector.SparkTransformConverter;
import org.apache.gravitino.spark.connector.SparkTypeConverter;
import org.apache.gravitino.spark.connector.catalog.GravitinoCatalogManager;
import org.apache.iceberg.spark.SparkCatalog;
+import org.apache.spark.SparkConf;
import org.apache.spark.sql.catalyst.analysis.NoSuchFunctionException;
import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
import org.apache.spark.sql.connector.catalog.Identifier;
@@ -66,7 +67,7 @@ public class TestGravitinoGlueCatalog {
// GravitinoGlueCatalog extends BaseCatalog which calls
GravitinoCatalogManager.get()
// in its constructor, so we must initialize the manager first.
GravitinoClient mockClient = mock(GravitinoClient.class);
- GravitinoCatalogManager.create(() -> mockClient);
+ GravitinoCatalogManager.create(new SparkConf(false), "user", identity ->
mockClient);
}
@AfterAll
diff --git
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/plugin/TestGravitinoDriverPlugin.java
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/plugin/TestGravitinoDriverPlugin.java
index 7da747c98d..da92012e3d 100644
---
a/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/plugin/TestGravitinoDriverPlugin.java
+++
b/spark-connector/spark-common/src/test/java/org/apache/gravitino/spark/connector/plugin/TestGravitinoDriverPlugin.java
@@ -21,12 +21,21 @@ package org.apache.gravitino.spark.connector.plugin;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import com.google.common.collect.ImmutableMap;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import org.apache.gravitino.auth.AuthProperties;
+import org.apache.gravitino.spark.connector.GravitinoSparkConfig;
import
org.apache.gravitino.spark.connector.authorization.GravitinoAuthorizationSparkSessionExtensions;
+import
org.apache.gravitino.spark.connector.plugin.GravitinoDriverPlugin.DynamicBearerTokenProvider;
import org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions;
import org.apache.spark.SparkConf;
import org.apache.spark.sql.internal.StaticSQLConf;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
public class TestGravitinoDriverPlugin {
@@ -58,4 +67,84 @@ public class TestGravitinoDriverPlugin {
assertEquals(extension,
sparkConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS().key()));
}
+
+ @Test
+ void testTokenAuthTypeBuildsClient() {
+ SparkConf sparkConf = tokenAuthConf();
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, "a-token");
+
+ // The client cannot reach a server here, but it must get past auth
configuration first: an
+ // unsupported auth type or a missing token would fail before any
connection is attempted.
+ Exception e =
+ Assertions.assertThrows(
+ Exception.class,
+ () ->
+ GravitinoDriverPlugin.createGravitinoClient(
+ "http://127.0.0.1:1";, "metalake", sparkConf, "user",
ImmutableMap.of()));
+ Assertions.assertFalse(e instanceof UnsupportedOperationException,
e.toString());
+ Assertions.assertFalse(e instanceof IllegalArgumentException,
e.toString());
+ }
+
+ @Test
+ void testTokenProviderReturnsBearerToken() {
+ SparkConf sparkConf = tokenAuthConf();
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, "a-token");
+
+ assertEquals(
+ "Bearer a-token",
+ new String(
+ new DynamicBearerTokenProvider(sparkConf).getTokenData(),
StandardCharsets.UTF_8));
+ }
+
+ @Test
+ void testTokenIsResolvedOnEveryRequest() {
+ SparkConf sparkConf = tokenAuthConf();
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, "first-token");
+ DynamicBearerTokenProvider provider = new
DynamicBearerTokenProvider(sparkConf);
+
+ assertEquals("Bearer first-token", new String(provider.getTokenData(),
StandardCharsets.UTF_8));
+
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, "second-token");
+
+ assertEquals(
+ "Bearer second-token", new String(provider.getTokenData(),
StandardCharsets.UTF_8));
+ }
+
+ @Test
+ void testTokenFileTakesPrecedenceAndIsRereadEveryRequest(@TempDir Path
tempDir)
+ throws IOException {
+ Path tokenFile = tempDir.resolve("token");
+ Files.write(tokenFile, "file-token\n".getBytes(StandardCharsets.UTF_8));
+ SparkConf sparkConf = tokenAuthConf();
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE, "conf-token");
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_TOKEN_FILE,
tokenFile.toString());
+ DynamicBearerTokenProvider provider = new
DynamicBearerTokenProvider(sparkConf);
+
+ assertEquals("Bearer file-token", new String(provider.getTokenData(),
StandardCharsets.UTF_8));
+
+ Files.write(tokenFile, "rotated-token\n".getBytes(StandardCharsets.UTF_8));
+
+ assertEquals(
+ "Bearer rotated-token", new String(provider.getTokenData(),
StandardCharsets.UTF_8));
+ }
+
+ @Test
+ void testTokenAuthTypeWithoutTokenFails() {
+ SparkConf sparkConf = tokenAuthConf();
+
+ IllegalArgumentException e =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ GravitinoDriverPlugin.createGravitinoClient(
+ "http://127.0.0.1:1";, "metalake", sparkConf, "user",
ImmutableMap.of()));
+
Assertions.assertTrue(e.getMessage().contains(GravitinoSparkConfig.GRAVITINO_TOKEN_VALUE));
+
Assertions.assertTrue(e.getMessage().contains(GravitinoSparkConfig.GRAVITINO_TOKEN_FILE));
+ }
+
+ private static SparkConf tokenAuthConf() {
+ SparkConf sparkConf = new SparkConf(false);
+ sparkConf.set(GravitinoSparkConfig.GRAVITINO_AUTH_TYPE,
AuthProperties.TOKEN_AUTH_TYPE);
+ return sparkConf;
+ }
}
