diff --git a/build.sbt b/build.sbt index 393814ddcf..f9f7d867a5 100644 --- a/build.sbt +++ b/build.sbt @@ -185,6 +185,11 @@ lazy val client = (project in file("target/clients/java")) ), (Compile / compile) := ((Compile / compile) dependsOn generate).value, + // The javadoc tool crashes on the OpenAPI-generated sources (ClientCodeException wrapping a + // StringIndexOutOfBoundsException). The jar we consume needs no javadoc, so stop building it. + Compile / packageDoc / publishArtifact := false, + Compile / doc / sources := Seq.empty, + // OpenAPI generation specs openApiInputSpec := (file(".") / "api" / "all.yaml").toString, openApiGeneratorName := "java", @@ -556,6 +561,9 @@ lazy val spark = (project in file("connectors/spark")) "org.antlr" % "antlr4" % "4.9.3", "com.google.cloud.bigdataoss" % "util-hadoop" % "3.0.2" % Provided, "org.apache.hadoop" % "hadoop-azure" % "3.4.0" % Provided, + // S3VendedCredentialsProvider implements the SDK v2 AwsCredentialsProvider. Provided, like + // the GCS and ABFS SPIs above: hadoop-aws ships the SDK on the runtime classpath. + "software.amazon.awssdk" % "auth" % "2.24.0" % Provided, ), libraryDependencies ++= Seq( // Test dependencies diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/S3VendedCredentialsProvider.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/S3VendedCredentialsProvider.java new file mode 100644 index 0000000000..0d2f47d559 --- /dev/null +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/S3VendedCredentialsProvider.java @@ -0,0 +1,235 @@ +package io.unitycatalog.spark; + +import io.unitycatalog.client.ApiClient; +import io.unitycatalog.client.ApiException; +import io.unitycatalog.client.api.TemporaryCredentialsApi; +import io.unitycatalog.client.model.AwsCredentials; +import io.unitycatalog.client.model.GenerateTemporaryTableCredential; +import io.unitycatalog.client.model.TableOperation; +import io.unitycatalog.client.model.TemporaryCredentials; +import io.unitycatalog.spark.auth.AuthConfigUtils; +import io.unitycatalog.spark.auth.catalog.UCTokenProvider; +import java.net.URI; +import java.time.Instant; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.hadoop.conf.Configuration; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; + +/** + * S3A credentials provider backed by Unity Catalog vended credentials, renewing them before they + * expire. + * + *

Without this, the credentials UC vends at planning time are written verbatim into the plan as + * {@code fs.s3a.access.key}/{@code secret.key}/{@code session.token}. They are scoped to roughly an + * hour, so any stage still running past that point fails on S3 rather than on anything that looks + * like an auth problem. Azure and GCS already went through a provider; S3 was the one cloud left + * holding literal keys. + * + *

The provider re-requests credentials from UC on its own, which is why it is installed on the + * executors too: it needs the catalog coordinates and the auth configuration, both carried in the + * Hadoop configuration, and re-authenticates through {@link UCTokenProvider}. + * + *

Backport of the 0.3.x {@code AwsVendedTokenProvider} and its {@code GenericCredentialProvider} + * base, collapsed into one class: the 0.3.x split exists to share a cache across four clouds and + * five credential scopes, none of which applies to a table-scoped S3-only path. + */ +public class S3VendedCredentialsProvider implements AwsCredentialsProvider { + + /** Bootstrap credentials vended on the driver, so the first S3 call needs no round trip. */ + protected static final String INIT_ACCESS_KEY = "fs.s3a.uc.init.access.key"; + + protected static final String INIT_SECRET_KEY = "fs.s3a.uc.init.secret.key"; + protected static final String INIT_SESSION_TOKEN = "fs.s3a.uc.init.session.token"; + protected static final String INIT_EXPIRATION_TIME = "fs.s3a.uc.init.expiration.time"; + + /** Coordinates the provider needs to re-request credentials for the same scope. */ + protected static final String UC_URI = "fs.s3a.uc.uri"; + + protected static final String UC_TABLE_ID = "fs.s3a.uc.table.id"; + protected static final String UC_TABLE_OPERATION = "fs.s3a.uc.table.operation"; + + /** Auth configuration for the UC client, flattened under this prefix. */ + protected static final String UC_AUTH_PREFIX = "fs.s3a.uc.auth."; + + /** How long before expiry a credential is renewed. */ + protected static final String RENEWAL_LEAD_TIME_MILLIS = "fs.s3a.uc.renewal.leadTimeMillis"; + + protected static final long DEFAULT_RENEWAL_LEAD_TIME_MILLIS = 60_000L; + + /** + * Keyed by credential scope, and static because {@code generateCredentialProps} sets {@code + * fs.s3a.impl.disable.cache=true}: Spark builds a fresh S3AFileSystem, and so a fresh provider, + * on every resolution. An instance-level cache would never be reused, and each new FileSystem + * would re-authenticate and re-vend. + */ + private static final Map CACHE = new ConcurrentHashMap<>(); + + private final Configuration conf; + private final CredentialScope scope; + private final long renewalLeadTimeMillis; + + /** Constructor signature required by Hadoop's S3A provider instantiation. */ + public S3VendedCredentialsProvider(URI uri, Configuration conf) { + this.conf = conf; + this.scope = + new CredentialScope( + conf.get(UC_URI), conf.get(UC_TABLE_ID), conf.get(UC_TABLE_OPERATION)); + this.renewalLeadTimeMillis = + conf.getLong(RENEWAL_LEAD_TIME_MILLIS, DEFAULT_RENEWAL_LEAD_TIME_MILLIS); + + VendedCredentials bootstrap = readBootstrapCredentials(conf); + if (bootstrap != null) { + CACHE.putIfAbsent(scope, bootstrap); + } + } + + @Override + public software.amazon.awssdk.auth.credentials.AwsCredentials resolveCredentials() { + VendedCredentials credentials = + CACHE.compute( + scope, + (key, cached) -> { + if (cached != null && !cached.readyToRenew(renewalLeadTimeMillis)) { + return cached; + } + return vend(key); + }); + + return AwsSessionCredentials.builder() + .accessKeyId(credentials.accessKeyId) + .secretAccessKey(credentials.secretAccessKey) + .sessionToken(credentials.sessionToken) + .build(); + } + + private VendedCredentials vend(CredentialScope key) { + if (key.ucUri == null || key.tableId == null) { + throw new IllegalStateException( + "Cannot renew Unity Catalog credentials: " + + UC_URI + + " and " + + UC_TABLE_ID + + " are missing from the Hadoop configuration"); + } + + URI uri = URI.create(key.ucUri); + ApiClient apiClient = + new ApiClient().setHost(uri.getHost()).setPort(uri.getPort()).setScheme(uri.getScheme()); + + Map authConfigs = AuthConfigUtils.buildAuthConfigs(prefixedAuthOptions()); + if (!authConfigs.isEmpty()) { + UCTokenProvider tokenProvider = UCTokenProvider.create(authConfigs); + apiClient = + apiClient.setRequestInterceptor( + request -> request.header("Authorization", "Bearer " + tokenProvider.accessToken())); + } + + TemporaryCredentials vended; + try { + vended = + new TemporaryCredentialsApi(apiClient) + .generateTemporaryTableCredentials( + new GenerateTemporaryTableCredential() + .tableId(key.tableId) + .operation(TableOperation.fromValue(key.operation))); + } catch (ApiException e) { + throw new RuntimeException( + "Failed to renew Unity Catalog credentials for table " + key.tableId, e); + } + + AwsCredentials aws = vended.getAwsTempCredentials(); + if (aws == null) { + throw new IllegalStateException( + "Unity Catalog returned no AWS credentials for table " + key.tableId); + } + return new VendedCredentials( + aws.getAccessKeyId(), + aws.getSecretAccessKey(), + aws.getSessionToken(), + vended.getExpirationTime()); + } + + /** + * Re-keys the auth options from {@link #UC_AUTH_PREFIX} to the bare names {@link AuthConfigUtils} + * expects, so the catalog and the provider share one normalization path. + */ + private Map prefixedAuthOptions() { + Map options = new HashMap<>(); + conf.getPropsWithPrefix(UC_AUTH_PREFIX).forEach(options::put); + return options; + } + + private static VendedCredentials readBootstrapCredentials(Configuration conf) { + String accessKey = conf.get(INIT_ACCESS_KEY); + String secretKey = conf.get(INIT_SECRET_KEY); + String sessionToken = conf.get(INIT_SESSION_TOKEN); + if (accessKey == null || secretKey == null || sessionToken == null) { + return null; + } + // Absent expiry means "never renew", matching how 0.3.x treats a static credential. UC only + // omits it for non-expiring credentials. + long expiration = conf.getLong(INIT_EXPIRATION_TIME, Long.MAX_VALUE); + return new VendedCredentials(accessKey, secretKey, sessionToken, expiration); + } + + /** Test-only: the cache is static, so it outlives a single test. */ + static void clearCache() { + CACHE.clear(); + } + + /** The resource a vended credential grants access to; two calls sharing it can share a credential. */ + private static final class CredentialScope { + private final String ucUri; + private final String tableId; + private final String operation; + + CredentialScope(String ucUri, String tableId, String operation) { + this.ucUri = ucUri; + this.tableId = tableId; + this.operation = operation == null ? TableOperation.READ_WRITE.getValue() : operation; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + CredentialScope that = (CredentialScope) o; + return Objects.equals(ucUri, that.ucUri) + && Objects.equals(tableId, that.tableId) + && Objects.equals(operation, that.operation); + } + + @Override + public int hashCode() { + return Objects.hash(ucUri, tableId, operation); + } + } + + private static final class VendedCredentials { + private final String accessKeyId; + private final String secretAccessKey; + private final String sessionToken; + private final Long expirationTime; + + VendedCredentials( + String accessKeyId, String secretAccessKey, String sessionToken, Long expirationTime) { + this.accessKeyId = accessKeyId; + this.secretAccessKey = secretAccessKey; + this.sessionToken = sessionToken; + this.expirationTime = expirationTime; + } + + boolean readyToRenew(long leadTimeMillis) { + return expirationTime != null + && expirationTime <= Instant.now().toEpochMilli() + leadTimeMillis; + } + } +} diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/UCSingleCatalog.scala b/connectors/spark/src/main/scala/io/unitycatalog/spark/UCSingleCatalog.scala index 3b0aeda260..57154c8eed 100644 --- a/connectors/spark/src/main/scala/io/unitycatalog/spark/UCSingleCatalog.scala +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/UCSingleCatalog.scala @@ -3,7 +3,8 @@ package io.unitycatalog.spark import io.unitycatalog.client.{ApiClient, ApiException} import io.unitycatalog.client.api.{SchemasApi, TablesApi, TemporaryCredentialsApi} import io.unitycatalog.client.model.{ColumnInfo, ColumnTypeName, CreateSchema, CreateTable, DataSourceFormat, GenerateTemporaryPathCredential, GenerateTemporaryTableCredential, ListTablesResponse, PathOperation, SchemaInfo, TableOperation, TableType, TemporaryCredentials} -import io.unitycatalog.spark.auth.catalog.UCTokenProvider +import io.unitycatalog.spark.auth.AuthConfigUtils +import io.unitycatalog.spark.auth.catalog.{AuthConfigs, UCTokenProvider} import java.net.URI import java.util @@ -43,24 +44,13 @@ class UCSingleCatalog extends TableCatalog with SupportsNamespaces with Logging .setHost(url.getHost) .setPort(url.getPort) .setScheme(url.getScheme) - // Backport (0.3.x formalism): resolve auth from options. Supports a static `token` - // or the OAuth 2.0 client-credentials keys `oauth.uri`/`oauth.clientId`/`oauth.clientSecret` - // (machine-to-machine). The interceptor is dynamic: `accessToken()` is called on every - // request, so `OAuthUCTokenProvider` refreshes the token transparently for long sessions. - val hasAuthConfig = options.get(UCTokenProvider.TOKEN) != null || - options.get(UCTokenProvider.OAUTH_URI) != null || - options.get(UCTokenProvider.OAUTH_CLIENT_ID) != null || - options.get(UCTokenProvider.OAUTH_CLIENT_SECRET) != null - if (hasAuthConfig) { - // `options` is a CaseInsensitiveStringMap (keys lowercased), so read each key via `get` - // and rebuild a map with the exact keys the factory expects (e.g. `oauth.clientId`). - val authOptions = new util.HashMap[String, String] - Seq(UCTokenProvider.TOKEN, UCTokenProvider.OAUTH_URI, UCTokenProvider.OAUTH_CLIENT_ID, - UCTokenProvider.OAUTH_CLIENT_SECRET).foreach { key => - val value = options.get(key) - if (value != null) authOptions.put(key, value) - } - val tokenProvider = UCTokenProvider.create(authOptions, "") + // Backport (0.3.x formalism): `AuthConfigUtils` normalizes the options into a flat, `type`-keyed + // map, accepting both the new `auth.*` keys and the legacy un-prefixed ones this connector + // shipped. The interceptor is dynamic: `accessToken()` is called on every request, so the + // providers refresh the token transparently for long sessions. + val authConfigs = AuthConfigUtils.buildAuthConfigs(options.asCaseSensitiveMap()) + if (authConfigs.containsKey(AuthConfigs.TYPE)) { + val tokenProvider = UCTokenProvider.create(authConfigs) apiClient = apiClient.setRequestInterceptor { request => request.header("Authorization", "Bearer " + tokenProvider.accessToken()) } @@ -72,7 +62,12 @@ class UCSingleCatalog extends TableCatalog with SupportsNamespaces with Logging // where external credential vending is disabled but the caller already has direct storage access. skipCredentialVending = java.lang.Boolean.parseBoolean( options.getOrDefault("skipCredentialVending", "false")) - val proxy = new UCProxy(apiClient, temporaryCredentialsApi, skipCredentialVending) + val proxy = new UCProxy( + apiClient, + temporaryCredentialsApi, + skipCredentialVending, + urlStr, + authConfigs.asScala.toMap) proxy.initialize(name, options) if (UCSingleCatalog.LOAD_DELTA_CATALOG.get()) { try { @@ -188,12 +183,47 @@ object UCSingleCatalog { val LOAD_DELTA_CATALOG = ThreadLocal.withInitial[Boolean](() => true) val DELTA_CATALOG_LOADED = ThreadLocal.withInitial[Boolean](() => false) + /** + * What `S3VendedCredentialsProvider` needs to re-request credentials for the same scope: where UC + * lives, which table, which operation, and how to authenticate. Carried through the Hadoop + * configuration so the executors can renew too, since the driver's provider instance does not + * travel with the plan. + */ + case class S3CredentialRenewal( + ucUri: String, + tableId: String, + operation: String, + authConfigs: Map[String, String]) { + + def hadoopProps(temporaryCredentials: TemporaryCredentials): Map[String, String] = { + val bootstrap = Map( + S3VendedCredentialsProvider.INIT_ACCESS_KEY -> + temporaryCredentials.getAwsTempCredentials.getAccessKeyId, + S3VendedCredentialsProvider.INIT_SECRET_KEY -> + temporaryCredentials.getAwsTempCredentials.getSecretAccessKey, + S3VendedCredentialsProvider.INIT_SESSION_TOKEN -> + temporaryCredentials.getAwsTempCredentials.getSessionToken, + S3VendedCredentialsProvider.UC_URI -> ucUri, + S3VendedCredentialsProvider.UC_TABLE_ID -> tableId, + S3VendedCredentialsProvider.UC_TABLE_OPERATION -> operation, + "fs.s3a.aws.credentials.provider" -> classOf[S3VendedCredentialsProvider].getName + ) + // UC omits the expiry only for non-expiring credentials; the provider then never renews. + val expiry = Option(temporaryCredentials.getExpirationTime) + .map(e => S3VendedCredentialsProvider.INIT_EXPIRATION_TIME -> e.toString) + bootstrap ++ expiry ++ authConfigs.map { + case (k, v) => (S3VendedCredentialsProvider.UC_AUTH_PREFIX + k) -> v + } + } + } + def generateCredentialProps( scheme: String, - temporaryCredentials: TemporaryCredentials): Map[String, String] = { + temporaryCredentials: TemporaryCredentials, + renewal: Option[S3CredentialRenewal] = None): Map[String, String] = { if (scheme == "s3") { val awsCredentials = temporaryCredentials.getAwsTempCredentials - Map( + val base = Map( // TODO: how to support s3:// properly? "fs.s3a.access.key" -> awsCredentials.getAccessKeyId, "fs.s3a.secret.key" -> awsCredentials.getSecretAccessKey, @@ -202,6 +232,10 @@ object UCSingleCatalog { "fs.s3.impl.disable.cache" -> "true", "fs.s3a.impl.disable.cache" -> "true" ) + // Vended credentials last about an hour, and the keys above are literals frozen into the + // plan. Where the table coordinates are known, hand S3A a provider that can re-vend instead, + // so a stage outliving them does not fail on S3. Azure and GCS already work this way. + renewal.map(base ++ _.hadoopProps(temporaryCredentials)).getOrElse(base) } else if (scheme == "gs") { val gcsCredentials = temporaryCredentials.getGcpOauthToken Map( @@ -232,7 +266,9 @@ object UCSingleCatalog { private class UCProxy( apiClient: ApiClient, temporaryCredentialsApi: TemporaryCredentialsApi, - skipCredentialVending: Boolean) extends TableCatalog with SupportsNamespaces { + skipCredentialVending: Boolean, + ucUri: String, + authConfigs: Map[String, String]) extends TableCatalog with SupportsNamespaces { private[this] var name: String = null private[this] var tablesApi: TablesApi = null private[this] var schemasApi: SchemasApi = null @@ -300,7 +336,11 @@ private class UCProxy( ) } } - UCSingleCatalog.generateCredentialProps(uri.getScheme, temporaryCredentials) + UCSingleCatalog.generateCredentialProps( + uri.getScheme, + temporaryCredentials, + Some(UCSingleCatalog.S3CredentialRenewal( + ucUri, tableId, TableOperation.READ_WRITE.getValue, authConfigs))) } val sparkTable = CatalogTable( identifier, diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/AuthConfigUtils.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/AuthConfigUtils.java new file mode 100644 index 0000000000..7018887ea7 --- /dev/null +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/AuthConfigUtils.java @@ -0,0 +1,110 @@ +package io.unitycatalog.spark.auth; + +import io.unitycatalog.spark.auth.catalog.AuthConfigs; +import java.util.HashMap; +import java.util.Map; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.sparkproject.guava.base.Preconditions; + +/** + * Normalizes catalog options into the flat, {@code type}-keyed configuration map that {@code + * UCTokenProvider.create} expects. + * + *

Backport of the 0.3.x {@code AuthConfigUtils}, extended with the legacy {@code oauth.*} shape: + * this connector shipped those keys before the {@code auth.} prefix existed, so both are accepted. + */ +public class AuthConfigUtils { + private static final String AUTH_PREFIX = "auth."; + + private AuthConfigUtils() {} + + public static Map buildAuthConfigs(Map configs) { + Map newConfigs = new HashMap<>(); + + for (Map.Entry e : configs.entrySet()) { + if (e.getKey().startsWith(AUTH_PREFIX)) { + // Remove the 'auth.' prefix from the key and add the normalized key-value pair. + String newKey = e.getKey().substring(AUTH_PREFIX.length()).trim(); + if (!newKey.isEmpty() && isSet(e.getValue())) { + newConfigs.put(newKey, e.getValue()); + } + } + } + + // Unity Catalog versions 0.3.0 and earlier did not use the 'auth.token' key. To maintain + // backward compatibility, we also copy the legacy 'token' key directly into the new config map. + String token = configs.get(AuthConfigs.STATIC_TOKEN); + if (isSet(token)) { + Preconditions.checkArgument( + !newConfigs.containsKey(AuthConfigs.STATIC_TOKEN), + "Static token was configured twice, choose only one: 'token' (legacy) or 'auth.token' (new-style)."); + + newConfigs.put(AuthConfigs.TYPE, AuthConfigs.STATIC_TYPE_VALUE); + newConfigs.put(AuthConfigs.STATIC_TOKEN, token); + } + + // Same treatment for the un-prefixed oauth keys, which this fork shipped before adopting the + // 0.3.x formalism. Infers the type so existing catalog configs keep working untouched. + copyLegacyGroup( + configs, + newConfigs, + AuthConfigs.OAUTH_TYPE_VALUE, + AuthConfigs.OAUTH_URI, + AuthConfigs.OAUTH_CLIENT_ID, + AuthConfigs.OAUTH_CLIENT_SECRET); + + copyLegacyGroup( + configs, + newConfigs, + AuthConfigs.OIDC_TYPE_VALUE, + AuthConfigs.OIDC_URI, + AuthConfigs.OIDC_CLIENT_ID, + AuthConfigs.OIDC_TOKEN_FILE_PATH); + + return new CaseInsensitiveStringMap(newConfigs); + } + + private static void copyLegacyGroup( + Map configs, + Map newConfigs, + String typeValue, + String... keys) { + boolean anyLegacy = false; + for (String key : keys) { + if (isSet(configs.get(key))) { + anyLegacy = true; + break; + } + } + if (!anyLegacy) { + return; + } + + for (String key : keys) { + String value = configs.get(key); + if (!isSet(value)) { + continue; + } + Preconditions.checkArgument( + !newConfigs.containsKey(key), + "'%s' was configured twice, choose only one: '%s' (legacy) or 'auth.%s' (new-style).", + key, + key, + key); + newConfigs.put(key, value); + } + + if (!newConfigs.containsKey(AuthConfigs.TYPE)) { + newConfigs.put(AuthConfigs.TYPE, typeValue); + } + } + + /** + * An empty value counts as unset. Spark hands over the keys a session declared even when their + * value is blank, and a blank credential must leave the catalog unauthenticated rather than + * select a provider that then rejects it. + */ + private static boolean isSet(String value) { + return value != null && !value.trim().isEmpty(); + } +} diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/AuthConfigs.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/AuthConfigs.java new file mode 100644 index 0000000000..7bd385e2ef --- /dev/null +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/AuthConfigs.java @@ -0,0 +1,38 @@ +package io.unitycatalog.spark.auth.catalog; + +/** + * Configuration keys for {@link UCTokenProvider} implementations. + * + *

Backport of the 0.3.x {@code io.unitycatalog.client.auth.AuthConfigs} to the 0.2.x connector. + * The 0.3.x client is generated at build time here, so the constants live in the Spark module + * instead; the key names are kept identical so a future rebase is a package move. Public rather + * than package-private for the same reason: {@link io.unitycatalog.spark.auth.AuthConfigUtils} sits + * in the package upstream puts it in, and reads these instead of redeclaring them. + */ +public final class AuthConfigs { + private AuthConfigs() {} + + // Define the authentication type, where the type can be: + // 1. static: which uses the FixedUCTokenProvider to return a pre-configured token. + // 2. oauth: which uses the OAuthUCTokenProvider to run the client credentials flow. + // 3. oidc: which uses the FileOidcUCTokenProvider to exchange a federated OIDC token. + // 4. fully qualified class name of a custom UCTokenProvider implementation. + public static final String TYPE = "type"; + + // Configure keys for the static token provider. + public static final String STATIC_TYPE_VALUE = "static"; + public static final String STATIC_TOKEN = "token"; + + // Configure keys for the oauth token provider. + public static final String OAUTH_TYPE_VALUE = "oauth"; + public static final String OAUTH_URI = "oauth.uri"; + public static final String OAUTH_CLIENT_ID = "oauth.clientId"; + public static final String OAUTH_CLIENT_SECRET = "oauth.clientSecret"; + + // Configure keys for the OIDC federation token provider. lbc addition with no 0.3.x counterpart + // yet, so the naming follows the oauth block rather than inventing a new shape. + public static final String OIDC_TYPE_VALUE = "oidc"; + public static final String OIDC_URI = "oidc.uri"; + public static final String OIDC_CLIENT_ID = "oidc.clientId"; + public static final String OIDC_TOKEN_FILE_PATH = "oidc.tokenFilePath"; +} diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FileOidcUCTokenProvider.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FileOidcUCTokenProvider.java new file mode 100644 index 0000000000..740d8b61e7 --- /dev/null +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FileOidcUCTokenProvider.java @@ -0,0 +1,205 @@ +package io.unitycatalog.spark.auth.catalog; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.IOException; +import java.net.URI; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.time.Instant; +import java.util.Map; +import java.util.function.Supplier; +import org.sparkproject.guava.base.Preconditions; + +/** + * Internal class - not intended for direct use. + * + *

Token provider for OIDC workload identity federation: reads an OIDC token from a file and + * exchanges it for a Unity Catalog access token (RFC 8693 token exchange). Nothing secret is + * stored, so there is no client secret to rotate; the caller is identified by its {@code + * oidc.clientId} plus the trust the server places in the token issuer. + * + *

Intended for Kubernetes workloads, where the file is a projected service account token. + * Cached and renewed {@link #DEFAULT_LEAD_RENEWAL_TIME_SECONDS} seconds before expiration, + * thread-safe via double-checked locking, like {@link OAuthUCTokenProvider}. + * + *

lbc addition, with no 0.3.x counterpart yet, so the option naming follows the {@code oauth.*} + * block rather than inventing a new shape. + */ +class FileOidcUCTokenProvider implements UCTokenProvider { + private static final long DEFAULT_LEAD_RENEWAL_TIME_SECONDS = 30L; + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + private static final String GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"; + private static final String SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt"; + + private String oidcUri; + private String clientId; + private String tokenFilePath; + private long leadRenewalTimeSeconds; + private HttpClient httpClient; + private Supplier clock; + + private volatile TempToken tempToken; + + FileOidcUCTokenProvider() {} + + // Package-private constructor for testing with custom dependencies. + FileOidcUCTokenProvider( + String oidcUri, + String clientId, + String tokenFilePath, + long leadRenewalTimeSeconds, + HttpClient httpClient, + Supplier clock) { + Preconditions.checkNotNull(oidcUri, "OIDC URI must not be null"); + Preconditions.checkNotNull(clientId, "OIDC client ID must not be null"); + Preconditions.checkNotNull(tokenFilePath, "OIDC token file path must not be null"); + Preconditions.checkArgument( + leadRenewalTimeSeconds >= 0, + "Lead renewal time must be non-negative, but got %s", + leadRenewalTimeSeconds); + Preconditions.checkNotNull(httpClient, "HTTP client must not be null"); + Preconditions.checkNotNull(clock, "Clock must not be null"); + + this.oidcUri = oidcUri; + this.clientId = clientId; + this.tokenFilePath = tokenFilePath; + this.leadRenewalTimeSeconds = leadRenewalTimeSeconds; + this.httpClient = httpClient; + this.clock = clock; + } + + @Override + public void initialize(Map configs) { + String oidcUri = configs.get(AuthConfigs.OIDC_URI); + Preconditions.checkArgument( + oidcUri != null && !oidcUri.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.OIDC_URI); + this.oidcUri = oidcUri; + + String clientId = configs.get(AuthConfigs.OIDC_CLIENT_ID); + Preconditions.checkArgument( + clientId != null && !clientId.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.OIDC_CLIENT_ID); + this.clientId = clientId; + + String tokenFilePath = configs.get(AuthConfigs.OIDC_TOKEN_FILE_PATH); + Preconditions.checkArgument( + tokenFilePath != null && !tokenFilePath.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.OIDC_TOKEN_FILE_PATH); + this.tokenFilePath = tokenFilePath; + + this.leadRenewalTimeSeconds = DEFAULT_LEAD_RENEWAL_TIME_SECONDS; + this.httpClient = HttpClient.newHttpClient(); + this.clock = Instant::now; + } + + @Override + public String accessToken() { + if (tempToken == null || tempToken.isReadyToRenew()) { + synchronized (this) { + if (tempToken == null || tempToken.isReadyToRenew()) { + tempToken = renewToken(); + } + } + } + return tempToken.token(); + } + + private TempToken renewToken() { + try { + // Re-read on every renewal: kubelet rotates the projected token, so a token read once at + // initialization would expire mid-session. + String subjectToken = readSubjectToken(); + + String formData = + "grant_type=" + + encode(GRANT_TYPE) + + "&subject_token_type=" + + encode(SUBJECT_TOKEN_TYPE) + + "&subject_token=" + + encode(subjectToken) + + "&client_id=" + + encode(clientId) + + "&scope=all-apis"; + + // No Authorization header: the subject token is the proof, there is no secret to present. + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(oidcUri)) + .header("Content-Type", "application/x-www-form-urlencoded") + .POST(HttpRequest.BodyPublishers.ofString(formData)) + .build(); + + HttpResponse response = + httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + + if (response.statusCode() != 200) { + throw new IOException( + String.format( + "Failed to exchange OIDC token. HTTP status: %d, Response: %s", + response.statusCode(), response.body())); + } + + JsonNode jsonNode = OBJECT_MAPPER.readTree(response.body()); + String accessToken = jsonNode.get("access_token").asText(); + long expiresInSeconds = jsonNode.get("expires_in").asLong(); + + return new TempToken(accessToken, clock.get().plusSeconds(expiresInSeconds)); + } catch (Exception e) { + throw new RuntimeException("Failed to renew OIDC federated token", e); + } + } + + private String readSubjectToken() throws IOException { + Path path = Paths.get(tokenFilePath); + if (!Files.isRegularFile(path)) { + throw new IOException( + String.format( + "OIDC token file %s does not exist. Expected a projected service account token", + tokenFilePath)); + } + String token = new String(Files.readAllBytes(path), StandardCharsets.UTF_8).trim(); + if (token.isEmpty()) { + throw new IOException(String.format("OIDC token file %s is empty", tokenFilePath)); + } + return token; + } + + private static String encode(String value) { + try { + return URLEncoder.encode(value, StandardCharsets.UTF_8.name()); + } catch (java.io.UnsupportedEncodingException e) { + throw new IllegalStateException("UTF-8 is always supported", e); + } + } + + private class TempToken { + private final String token; + private final Instant expirationTime; + + TempToken(String token, Instant expirationTime) { + this.token = token; + this.expirationTime = expirationTime; + } + + String token() { + return token; + } + + boolean isReadyToRenew() { + Instant renewalTime = expirationTime.minusSeconds(leadRenewalTimeSeconds); + return clock.get().isAfter(renewalTime); + } + } +} diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FixedUCTokenProvider.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FixedUCTokenProvider.java index 3d75186933..ea1eaacd32 100644 --- a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FixedUCTokenProvider.java +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/FixedUCTokenProvider.java @@ -1,13 +1,25 @@ package io.unitycatalog.spark.auth.catalog; +import java.util.Map; import org.sparkproject.guava.base.Preconditions; -/** A {@link UCTokenProvider} that always returns a pre-configured static token. */ -public class FixedUCTokenProvider implements UCTokenProvider { - private final String token; +/** + * Internal class - not intended for direct use. + * + *

A {@link UCTokenProvider} that always returns a pre-configured static token. + */ +class FixedUCTokenProvider implements UCTokenProvider { + private String token; - public FixedUCTokenProvider(String token) { - Preconditions.checkNotNull(token, "Token must not be null"); + FixedUCTokenProvider() {} + + @Override + public void initialize(Map configs) { + String token = configs.get(AuthConfigs.STATIC_TOKEN); + Preconditions.checkArgument( + token != null && !token.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.STATIC_TOKEN); this.token = token; } diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/OAuthUCTokenProvider.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/OAuthUCTokenProvider.java index b9d4d3edd8..1c4945d0a9 100644 --- a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/OAuthUCTokenProvider.java +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/OAuthUCTokenProvider.java @@ -10,38 +10,37 @@ import java.nio.charset.StandardCharsets; import java.time.Instant; import java.util.Base64; +import java.util.Map; import java.util.function.Supplier; import org.sparkproject.guava.base.Preconditions; /** - * OAuth-based token provider that fetches and automatically renews access tokens using the - * OAuth 2.0 client credentials flow (machine-to-machine). + * Internal class - not intended for direct use. * - *

Backport of the 0.3.x {@code OAuthUCTokenProvider} to the 0.2.x connector. The 0.3.x - * version relies on {@code RetryingApiClient}/{@code ApiClientConf}/injectable {@code Clock}, - * none of which exist in 0.2.x; here the request uses the JDK {@link HttpClient} and - * {@link Instant#now()} directly. Behaviour is otherwise identical: token is cached and renewed - * {@link #DEFAULT_LEAD_RENEWAL_TIME_SECONDS} seconds before expiration, thread-safe via - * double-checked locking. The connector calls {@link #accessToken()} on every request so the - * refreshed token propagates without re-initializing the catalog. + *

OAuth-based token provider that fetches and automatically renews access tokens using the OAuth + * 2.0 client credentials flow (machine-to-machine). + * + *

Backport of the 0.3.x {@code OAuthTokenProvider} to the 0.2.x connector. The 0.3.x version + * relies on {@code RetryingApiClient}/{@code Clock}, neither of which exists in 0.2.x; here the + * request uses the JDK {@link HttpClient} and {@link Instant#now()} directly, so the exchange is not + * retried. Behaviour is otherwise identical: the token is cached and renewed {@link + * #DEFAULT_LEAD_RENEWAL_TIME_SECONDS} seconds before expiration, thread-safe via double-checked + * locking. */ -public class OAuthUCTokenProvider implements UCTokenProvider { +class OAuthUCTokenProvider implements UCTokenProvider { private static final long DEFAULT_LEAD_RENEWAL_TIME_SECONDS = 30L; private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - private final String oauthUri; - private final String oauthClientId; - private final String oauthClientSecret; - private final long leadRenewalTimeSeconds; - private final HttpClient httpClient; - private final Supplier clock; + private String oauthUri; + private String oauthClientId; + private String oauthClientSecret; + private long leadRenewalTimeSeconds; + private HttpClient httpClient; + private Supplier clock; private volatile TempToken tempToken; - public OAuthUCTokenProvider(String oauthUri, String oauthClientId, String oauthClientSecret) { - this(oauthUri, oauthClientId, oauthClientSecret, DEFAULT_LEAD_RENEWAL_TIME_SECONDS, - HttpClient.newHttpClient(), Instant::now); - } + OAuthUCTokenProvider() {} // Package-private constructor for testing with custom dependencies. OAuthUCTokenProvider( @@ -54,8 +53,10 @@ public OAuthUCTokenProvider(String oauthUri, String oauthClientId, String oauthC Preconditions.checkNotNull(oauthUri, "OAuth URI must not be null"); Preconditions.checkNotNull(oauthClientId, "OAuth client ID must not be null"); Preconditions.checkNotNull(oauthClientSecret, "OAuth client secret must not be null"); - Preconditions.checkArgument(leadRenewalTimeSeconds >= 0, - "Lead renewal time must be non-negative, but got %s", leadRenewalTimeSeconds); + Preconditions.checkArgument( + leadRenewalTimeSeconds >= 0, + "Lead renewal time must be non-negative, but got %s", + leadRenewalTimeSeconds); Preconditions.checkNotNull(httpClient, "HTTP client must not be null"); Preconditions.checkNotNull(clock, "Clock must not be null"); @@ -67,6 +68,34 @@ public OAuthUCTokenProvider(String oauthUri, String oauthClientId, String oauthC this.clock = clock; } + @Override + public void initialize(Map configs) { + String oauthUri = configs.get(AuthConfigs.OAUTH_URI); + Preconditions.checkArgument( + oauthUri != null && !oauthUri.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.OAUTH_URI); + this.oauthUri = oauthUri; + + String oauthClientId = configs.get(AuthConfigs.OAUTH_CLIENT_ID); + Preconditions.checkArgument( + oauthClientId != null && !oauthClientId.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.OAUTH_CLIENT_ID); + this.oauthClientId = oauthClientId; + + String oauthClientSecret = configs.get(AuthConfigs.OAUTH_CLIENT_SECRET); + Preconditions.checkArgument( + oauthClientSecret != null && !oauthClientSecret.isEmpty(), + "Configuration key '%s' is missing or empty", + AuthConfigs.OAUTH_CLIENT_SECRET); + this.oauthClientSecret = oauthClientSecret; + + this.leadRenewalTimeSeconds = DEFAULT_LEAD_RENEWAL_TIME_SECONDS; + this.httpClient = HttpClient.newHttpClient(); + this.clock = Instant::now; + } + @Override public String accessToken() { if (tempToken == null || tempToken.isReadyToRenew()) { @@ -83,25 +112,27 @@ private TempToken renewToken() { try { // Basic auth header from clientId:clientSecret. String credentials = String.format("%s:%s", oauthClientId, oauthClientSecret); - String encodedCredentials = Base64.getEncoder() - .encodeToString(credentials.getBytes(StandardCharsets.UTF_8)); + String encodedCredentials = + Base64.getEncoder().encodeToString(credentials.getBytes(StandardCharsets.UTF_8)); String formData = "grant_type=client_credentials&scope=all-apis"; - HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create(oauthUri)) - .header("Authorization", "Basic " + encodedCredentials) - .header("Content-Type", "application/x-www-form-urlencoded") - .POST(HttpRequest.BodyPublishers.ofString(formData)) - .build(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(oauthUri)) + .header("Authorization", "Basic " + encodedCredentials) + .header("Content-Type", "application/x-www-form-urlencoded") + .POST(HttpRequest.BodyPublishers.ofString(formData)) + .build(); - HttpResponse response = httpClient.send(request, - HttpResponse.BodyHandlers.ofString()); + HttpResponse response = + httpClient.send(request, HttpResponse.BodyHandlers.ofString()); if (response.statusCode() != 200) { - throw new IOException(String.format( - "Failed to fetch OAuth token. HTTP status: %d, Response: %s", - response.statusCode(), response.body())); + throw new IOException( + String.format( + "Failed to fetch OAuth token. HTTP status: %d, Response: %s", + response.statusCode(), response.body())); } JsonNode jsonNode = OBJECT_MAPPER.readTree(response.body()); diff --git a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/UCTokenProvider.java b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/UCTokenProvider.java index 969c4d4e24..da024680fd 100644 --- a/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/UCTokenProvider.java +++ b/connectors/spark/src/main/scala/io/unitycatalog/spark/auth/catalog/UCTokenProvider.java @@ -1,62 +1,91 @@ package io.unitycatalog.spark.auth.catalog; -import static org.sparkproject.guava.base.Preconditions.checkArgument; - +import io.unitycatalog.spark.auth.AuthConfigUtils; import java.util.Map; +import org.sparkproject.guava.base.Preconditions; /** * Interface for providing access tokens to authenticate with Unity Catalog. * *

Implementations: + * *

* - *

Backport of the 0.3.x {@code UCTokenProvider} to the 0.2.x connector: same option - * formalism ({@code token} or {@code oauth.uri}/{@code oauth.clientId}/{@code oauth.clientSecret}), - * but the OAuth implementation relies only on the JDK HTTP client available in 0.2.x. + *

Backport of the 0.3.x {@code TokenProvider} to the 0.2.x connector, including the {@code + * type}-based dispatch. The OAuth implementation relies only on the JDK HTTP client available in + * 0.2.x, as {@code RetryingApiClient} does not exist here. + * + *

The 0.3.x {@code configs()} accessor is deliberately left out: it exists there so an executor + * can rebuild a provider from the Hadoop configuration and renew vended credentials, which the + * 0.2.x connector never does. Everything here runs on the driver. */ public interface UCTokenProvider { - String OAUTH_URI = "oauth.uri"; - String OAUTH_CLIENT_ID = "oauth.clientId"; - String OAUTH_CLIENT_SECRET = "oauth.clientSecret"; - String TOKEN = "token"; + /** + * Initializes the token provider with configuration parameters. + * + * @param configs configuration map with authentication settings, keys without prefix + * @throws IllegalArgumentException if required parameters are missing or invalid + */ + void initialize(Map configs); /** Returns the access token for Unity Catalog authentication, refreshing it when needed. */ String accessToken(); /** - * Creates a token provider from catalog options (keys without prefix). Returns a - * {@link FixedUCTokenProvider} when {@code token} is set, otherwise an - * {@link OAuthUCTokenProvider} when the three {@code oauth.*} keys are set. + * Creates a token provider from a configuration map. * - * @param optionKeyPrefix prefix prepended to option keys in error messages, e.g. - * {@code "spark.sql.catalog.."} - * @throws IllegalArgumentException if no complete authentication configuration is found + *

Dispatches on the required {@code type} key: {@code static}, {@code oauth}, {@code oidc}, or + * the fully qualified class name of a custom {@link UCTokenProvider} implementation. Legacy + * option shapes are normalized to a {@code type} upstream by {@link AuthConfigUtils}. + * + * @throws IllegalArgumentException if {@code type} is missing, or if the parameters required by + * the selected type are missing or invalid + * @throws RuntimeException if a custom provider class cannot be instantiated */ - static UCTokenProvider create(Map options, String optionKeyPrefix) { - String token = options.get(TOKEN); - if (token != null && !token.isEmpty()) { - return new FixedUCTokenProvider(token); - } + static UCTokenProvider create(Map configs) { + String authType = configs.get(AuthConfigs.TYPE); + Preconditions.checkArgument( + authType != null && !authType.trim().isEmpty(), + "Required configuration key '%s' is missing or empty. " + + "Must be 'static', 'oauth', 'oidc', or a fully qualified UCTokenProvider class name.", + AuthConfigs.TYPE); + + UCTokenProvider tokenProvider; + switch (authType) { + case AuthConfigs.STATIC_TYPE_VALUE: + tokenProvider = new FixedUCTokenProvider(); + break; + + case AuthConfigs.OAUTH_TYPE_VALUE: + tokenProvider = new OAuthUCTokenProvider(); + break; + + case AuthConfigs.OIDC_TYPE_VALUE: + tokenProvider = new FileOidcUCTokenProvider(); + break; - String oauthUri = options.get(OAUTH_URI); - String oauthClientId = options.get(OAUTH_CLIENT_ID); - String oauthClientSecret = options.get(OAUTH_CLIENT_SECRET); - if (oauthUri != null || oauthClientId != null || oauthClientSecret != null) { - checkArgument(oauthUri != null && oauthClientId != null && oauthClientSecret != null, - "Incomplete OAuth configuration detected. All of the keys are required: " - + "%soauth.uri, %soauth.clientId, %soauth.clientSecret. Please ensure they are " - + "all set.", optionKeyPrefix, optionKeyPrefix, optionKeyPrefix); - return new OAuthUCTokenProvider(oauthUri, oauthClientId, oauthClientSecret); + default: + try { + tokenProvider = + (UCTokenProvider) Class.forName(authType).getDeclaredConstructor().newInstance(); + } catch (Exception e) { + throw new RuntimeException( + String.format( + "Failed to instantiate custom UCTokenProvider '%s'. Ensure the class exists, " + + "implements UCTokenProvider, and has a public no-arg constructor.", + authType), + e); + } + break; } - throw new IllegalArgumentException(String.format("Cannot determine UC authentication " - + "configuration from options, please set %stoken for static token authentication or " - + "%soauth.uri, %soauth.clientId, %soauth.clientSecret for OAuth 2.0 authentication " - + "(all three required)", - optionKeyPrefix, optionKeyPrefix, optionKeyPrefix, optionKeyPrefix)); + tokenProvider.initialize(configs); + return tokenProvider; } } diff --git a/connectors/spark/src/test/java/io/unitycatalog/spark/S3VendedCredentialsProviderTest.java b/connectors/spark/src/test/java/io/unitycatalog/spark/S3VendedCredentialsProviderTest.java new file mode 100644 index 0000000000..be3313086e --- /dev/null +++ b/connectors/spark/src/test/java/io/unitycatalog/spark/S3VendedCredentialsProviderTest.java @@ -0,0 +1,115 @@ +package io.unitycatalog.spark; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.net.URI; +import java.time.Instant; +import org.apache.hadoop.conf.Configuration; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; + +public class S3VendedCredentialsProviderTest { + + private static final URI S3_URI = URI.create("s3://bucket/table"); + + @BeforeEach + public void setUp() { + // The cache is static, so it survives across tests. + S3VendedCredentialsProvider.clearCache(); + } + + @Test + public void servesTheBootstrapCredentialsWithoutCallingUc() { + Configuration conf = bootstrapConf(farFuture()); + + S3VendedCredentialsProvider provider = new S3VendedCredentialsProvider(S3_URI, conf); + + // No UC coordinates are needed while the bootstrap credentials are still valid: had it tried to + // renew, the missing table id would have thrown. + AwsSessionCredentials credentials = (AwsSessionCredentials) provider.resolveCredentials(); + assertThat(credentials.accessKeyId()).isEqualTo("init-access-key"); + assertThat(credentials.secretAccessKey()).isEqualTo("init-secret-key"); + assertThat(credentials.sessionToken()).isEqualTo("init-session-token"); + } + + @Test + public void treatsAMissingExpiryAsNonExpiring() { + Configuration conf = bootstrapConf(null); + + S3VendedCredentialsProvider provider = new S3VendedCredentialsProvider(S3_URI, conf); + + // UC omits the expiry only for credentials that do not expire, so the provider must not try to + // renew them; that attempt would fail on the absent UC coordinates. + assertThat(provider.resolveCredentials()).isNotNull(); + } + + @Test + public void renewsOnceTheBootstrapCredentialsAreWithinTheLeadTime() { + Configuration conf = bootstrapConf(Instant.now().toEpochMilli() + 5_000L); + conf.setLong(S3VendedCredentialsProvider.RENEWAL_LEAD_TIME_MILLIS, 60_000L); + // Deliberately no UC coordinates: reaching the renewal path is what the failure proves. + conf.unset(S3VendedCredentialsProvider.UC_TABLE_ID); + + S3VendedCredentialsProvider provider = new S3VendedCredentialsProvider(S3_URI, conf); + + assertThatThrownBy(provider::resolveCredentials) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Cannot renew Unity Catalog credentials"); + } + + @Test + public void sharesACachedCredentialAcrossProviderInstancesOfTheSameScope() { + Configuration conf = bootstrapConf(farFuture()); + conf.set(S3VendedCredentialsProvider.UC_URI, "https://uc.example.com"); + conf.set(S3VendedCredentialsProvider.UC_TABLE_ID, "table-1"); + + new S3VendedCredentialsProvider(S3_URI, conf).resolveCredentials(); + + // fs.s3a.impl.disable.cache=true means Spark rebuilds the FileSystem, and the provider with it. + // A second instance for the same scope must reuse the cached credential rather than re-vend, + // which it could not do here: these credentials carry no UC-reachable endpoint. + Configuration second = new Configuration(false); + second.set(S3VendedCredentialsProvider.UC_URI, "https://uc.example.com"); + second.set(S3VendedCredentialsProvider.UC_TABLE_ID, "table-1"); + second.set(S3VendedCredentialsProvider.UC_TABLE_OPERATION, "READ_WRITE"); + + AwsSessionCredentials credentials = + (AwsSessionCredentials) + new S3VendedCredentialsProvider(S3_URI, second).resolveCredentials(); + + assertThat(credentials.accessKeyId()).isEqualTo("init-access-key"); + } + + @Test + public void isolatesDifferentScopes() { + Configuration first = bootstrapConf(farFuture()); + first.set(S3VendedCredentialsProvider.UC_URI, "https://uc.example.com"); + first.set(S3VendedCredentialsProvider.UC_TABLE_ID, "table-1"); + new S3VendedCredentialsProvider(S3_URI, first).resolveCredentials(); + + // Another table must not be served table-1's credential; with no bootstrap of its own, the only + // way forward is a renewal, which fails on the absent endpoint. + Configuration other = new Configuration(false); + other.set(S3VendedCredentialsProvider.UC_TABLE_ID, "table-2"); + + assertThatThrownBy(() -> new S3VendedCredentialsProvider(S3_URI, other).resolveCredentials()) + .isInstanceOf(IllegalStateException.class); + } + + private static Long farFuture() { + return Instant.now().toEpochMilli() + 3_600_000L; + } + + private static Configuration bootstrapConf(Long expirationTime) { + Configuration conf = new Configuration(false); + conf.set(S3VendedCredentialsProvider.INIT_ACCESS_KEY, "init-access-key"); + conf.set(S3VendedCredentialsProvider.INIT_SECRET_KEY, "init-secret-key"); + conf.set(S3VendedCredentialsProvider.INIT_SESSION_TOKEN, "init-session-token"); + if (expirationTime != null) { + conf.setLong(S3VendedCredentialsProvider.INIT_EXPIRATION_TIME, expirationTime); + } + return conf; + } +} diff --git a/connectors/spark/src/test/java/io/unitycatalog/spark/auth/AuthConfigUtilsTest.java b/connectors/spark/src/test/java/io/unitycatalog/spark/auth/AuthConfigUtilsTest.java new file mode 100644 index 0000000000..c0ee342614 --- /dev/null +++ b/connectors/spark/src/test/java/io/unitycatalog/spark/auth/AuthConfigUtilsTest.java @@ -0,0 +1,165 @@ +package io.unitycatalog.spark.auth; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import io.unitycatalog.spark.auth.catalog.AuthConfigs; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; + +public class AuthConfigUtilsTest { + + @Test + public void stripsTheAuthPrefix() { + Map options = new HashMap<>(); + options.put("auth.type", AuthConfigs.OIDC_TYPE_VALUE); + options.put("auth.oidc.uri", "https://example.com/oidc/v1/token"); + options.put("auth.oidc.clientId", "client-id"); + options.put("auth.oidc.tokenFilePath", "/var/run/secrets/token"); + options.put("uri", "https://example.com"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs) + .containsEntry(AuthConfigs.TYPE, AuthConfigs.OIDC_TYPE_VALUE) + .containsEntry(AuthConfigs.OIDC_URI, "https://example.com/oidc/v1/token") + .containsEntry(AuthConfigs.OIDC_CLIENT_ID, "client-id") + .containsEntry(AuthConfigs.OIDC_TOKEN_FILE_PATH, "/var/run/secrets/token") + .doesNotContainKey("uri"); + } + + @Test + public void infersStaticTypeFromTheLegacyTokenKey() { + Map options = new HashMap<>(); + options.put(AuthConfigs.STATIC_TOKEN, "static-token"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs) + .containsEntry(AuthConfigs.TYPE, AuthConfigs.STATIC_TYPE_VALUE) + .containsEntry(AuthConfigs.STATIC_TOKEN, "static-token"); + } + + @Test + public void infersOAuthTypeFromTheLegacyUnprefixedKeys() { + Map options = new HashMap<>(); + options.put(AuthConfigs.OAUTH_URI, "https://example.com/oidc/v1/token"); + options.put(AuthConfigs.OAUTH_CLIENT_ID, "client-id"); + options.put(AuthConfigs.OAUTH_CLIENT_SECRET, "client-secret"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs) + .containsEntry(AuthConfigs.TYPE, AuthConfigs.OAUTH_TYPE_VALUE) + .containsEntry(AuthConfigs.OAUTH_CLIENT_ID, "client-id"); + } + + @Test + public void infersOidcTypeFromTheLegacyUnprefixedKeys() { + Map options = new HashMap<>(); + options.put(AuthConfigs.OIDC_URI, "https://example.com/oidc/v1/token"); + options.put(AuthConfigs.OIDC_CLIENT_ID, "client-id"); + options.put(AuthConfigs.OIDC_TOKEN_FILE_PATH, "/var/run/secrets/token"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs) + .containsEntry(AuthConfigs.TYPE, AuthConfigs.OIDC_TYPE_VALUE) + .containsEntry(AuthConfigs.OIDC_TOKEN_FILE_PATH, "/var/run/secrets/token"); + } + + @Test + public void anExplicitTypeWinsOverTheInferredOne() { + Map options = new HashMap<>(); + options.put("auth.type", AuthConfigs.OIDC_TYPE_VALUE); + options.put("auth.oidc.uri", "https://example.com/oidc/v1/token"); + options.put("auth.oidc.clientId", "client-id"); + options.put("auth.oidc.tokenFilePath", "/var/run/secrets/token"); + options.put(AuthConfigs.OAUTH_URI, "https://example.com/oidc/v1/token"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs).containsEntry(AuthConfigs.TYPE, AuthConfigs.OIDC_TYPE_VALUE); + } + + @Test + public void rejectsATokenConfiguredTwice() { + Map options = new HashMap<>(); + options.put(AuthConfigs.STATIC_TOKEN, "legacy-token"); + options.put("auth.token", "new-style-token"); + + assertThatThrownBy(() -> AuthConfigUtils.buildAuthConfigs(options)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Static token was configured twice"); + } + + @Test + public void rejectsAnOidcKeyConfiguredTwice() { + Map options = new HashMap<>(); + options.put(AuthConfigs.OIDC_CLIENT_ID, "legacy-client-id"); + options.put("auth.oidc.clientId", "new-style-client-id"); + + assertThatThrownBy(() -> AuthConfigUtils.buildAuthConfigs(options)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("was configured twice"); + } + + @Test + public void yieldsNoTypeWhenNoAuthIsConfigured() { + Map options = new HashMap<>(); + options.put("uri", "https://example.com"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + // UCSingleCatalog keys off the absence of `type` to skip installing an interceptor, which is + // how an unauthenticated local metastore keeps working. + assertThat(configs).doesNotContainKey(AuthConfigs.TYPE); + } + + @Test + public void keysAreCaseInsensitive() { + Map options = new HashMap<>(); + options.put("auth.type", AuthConfigs.OIDC_TYPE_VALUE); + options.put("auth.oidc.clientid", "client-id"); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + // Spark lowercases catalog option keys, so the returned map must not be case-sensitive. + assertThat(configs.get(AuthConfigs.OIDC_CLIENT_ID)).isEqualTo("client-id"); + } + + @Test + public void treatsAnEmptyValueAsUnset() { + Map options = new HashMap<>(); + // Spark hands over a declared key even when its value is blank, e.g. an unauthenticated local + // metastore. That must leave the catalog unauthenticated, not select a provider. + options.put(AuthConfigs.STATIC_TOKEN, ""); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs).doesNotContainKey(AuthConfigs.TYPE); + } + + @Test + public void treatsAnEmptyPrefixedValueAsUnset() { + Map options = new HashMap<>(); + options.put("auth.token", " "); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs).doesNotContainKey(AuthConfigs.TYPE); + } + + @Test + public void blankLegacyKeysDoNotInferAType() { + Map options = new HashMap<>(); + options.put(AuthConfigs.OAUTH_URI, ""); + options.put(AuthConfigs.OAUTH_CLIENT_ID, ""); + options.put(AuthConfigs.OAUTH_CLIENT_SECRET, ""); + + Map configs = AuthConfigUtils.buildAuthConfigs(options); + + assertThat(configs).doesNotContainKey(AuthConfigs.TYPE); + } +} diff --git a/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/FileOidcUCTokenProviderTest.java b/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/FileOidcUCTokenProviderTest.java new file mode 100644 index 0000000000..1568294d67 --- /dev/null +++ b/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/FileOidcUCTokenProviderTest.java @@ -0,0 +1,193 @@ +package io.unitycatalog.spark.auth.catalog; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Instant; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.Flow; +import java.util.function.Supplier; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class FileOidcUCTokenProviderTest { + + private static final String URI = "https://example.com/oidc/v1/token"; + private static final String CLIENT_ID = "client-id"; + + @TempDir Path tempDir; + + private HttpClient httpClient; + private Instant now; + + @BeforeEach + public void setUp() { + httpClient = mock(HttpClient.class); + now = Instant.parse("2026-01-01T00:00:00Z"); + } + + @Test + public void exchangesTheTokenReadFromFile() throws Exception { + Path tokenFile = writeToken("header.payload.signature"); + stubResponse(200, "{\"access_token\":\"exchanged\",\"expires_in\":3600}"); + + FileOidcUCTokenProvider provider = provider(tokenFile); + + assertThat(provider.accessToken()).isEqualTo("exchanged"); + + String body = capturedBody(); + assertThat(body) + .contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange"); + assertThat(body).contains("subject_token_type=urn%3Aietf%3Aparams%3Aoauth%3Atoken-type%3Ajwt"); + assertThat(body).contains("subject_token=header.payload.signature"); + assertThat(body).contains("client_id=client-id"); + assertThat(body).contains("scope=all-apis"); + assertThat(capturedRequest().headers().firstValue("Authorization")).isEmpty(); + } + + @Test + public void cachesTheTokenWhileItIsStillValid() throws Exception { + Path tokenFile = writeToken("header.payload.signature"); + stubResponse(200, "{\"access_token\":\"exchanged\",\"expires_in\":3600}"); + + FileOidcUCTokenProvider provider = provider(tokenFile); + + provider.accessToken(); + provider.accessToken(); + + verify(httpClient, times(1)).send(any(HttpRequest.class), any()); + } + + @Test + public void rereadsTheFileAndExchangesAgainOnceExpired() throws Exception { + Path tokenFile = writeToken("first.jwt.value"); + stubResponse(200, "{\"access_token\":\"first\",\"expires_in\":3600}"); + + FileOidcUCTokenProvider provider = provider(tokenFile); + assertThat(provider.accessToken()).isEqualTo("first"); + + // kubelet rotated the projected token, and the cached access token is past its renewal lead. + writeToken("second.jwt.value"); + stubResponse(200, "{\"access_token\":\"second\",\"expires_in\":3600}"); + now = now.plusSeconds(3600); + + assertThat(provider.accessToken()).isEqualTo("second"); + assertThat(capturedBody()).contains("subject_token=second.jwt.value"); + verify(httpClient, times(2)).send(any(HttpRequest.class), any()); + } + + @Test + public void failsWhenTheResponseIsNotSuccessful() throws Exception { + Path tokenFile = writeToken("header.payload.signature"); + stubResponse(400, "{\"error\":\"invalid_grant\"}"); + + FileOidcUCTokenProvider provider = provider(tokenFile); + + assertThatThrownBy(provider::accessToken) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("Failed to renew OIDC federated token") + .hasRootCauseMessage( + "Failed to exchange OIDC token. HTTP status: 400, Response: {\"error\":\"invalid_grant\"}"); + } + + @Test + public void failsWhenTheTokenFileIsMissing() { + FileOidcUCTokenProvider provider = provider(tempDir.resolve("absent")); + + assertThatThrownBy(provider::accessToken) + .isInstanceOf(RuntimeException.class) + .hasRootCauseInstanceOf(IOException.class) + .rootCause() + .hasMessageContaining("does not exist"); + } + + @Test + public void failsWhenTheTokenFileIsEmpty() throws Exception { + Path tokenFile = writeToken(" \n"); + + FileOidcUCTokenProvider provider = provider(tokenFile); + + assertThatThrownBy(provider::accessToken) + .isInstanceOf(RuntimeException.class) + .hasRootCauseInstanceOf(IOException.class) + .rootCause() + .hasMessageContaining("is empty"); + } + + private FileOidcUCTokenProvider provider(Path tokenFile) { + Supplier clock = () -> now; + return new FileOidcUCTokenProvider( + URI, CLIENT_ID, tokenFile.toString(), 30L, httpClient, clock); + } + + private Path writeToken(String content) throws IOException { + Path tokenFile = tempDir.resolve("token"); + Files.write(tokenFile, content.getBytes(StandardCharsets.UTF_8)); + return tokenFile; + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private void stubResponse(int statusCode, String body) throws Exception { + HttpResponse response = mock(HttpResponse.class); + when(response.statusCode()).thenReturn(statusCode); + when(response.body()).thenReturn(body); + when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))) + .thenReturn((HttpResponse) response); + } + + private HttpRequest capturedRequest() throws Exception { + org.mockito.ArgumentCaptor captor = + org.mockito.ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClient, org.mockito.Mockito.atLeastOnce()).send(captor.capture(), any()); + List requests = captor.getAllValues(); + return requests.get(requests.size() - 1); + } + + private String capturedBody() throws Exception { + Optional publisher = capturedRequest().bodyPublisher(); + assertThat(publisher).isPresent(); + return readBody(publisher.get()); + } + + private static String readBody(HttpRequest.BodyPublisher publisher) { + ByteArrayOutputStream collected = new ByteArrayOutputStream(); + publisher.subscribe( + new Flow.Subscriber() { + @Override + public void onSubscribe(Flow.Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(java.nio.ByteBuffer item) { + byte[] bytes = new byte[item.remaining()]; + item.get(bytes); + collected.write(bytes, 0, bytes.length); + } + + @Override + public void onError(Throwable throwable) { + throw new IllegalStateException(throwable); + } + + @Override + public void onComplete() {} + }); + return new String(collected.toByteArray(), StandardCharsets.UTF_8); + } +} diff --git a/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/UCTokenProviderTest.java b/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/UCTokenProviderTest.java index 18d490930c..232f24f72b 100644 --- a/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/UCTokenProviderTest.java +++ b/connectors/spark/src/test/java/io/unitycatalog/spark/auth/catalog/UCTokenProviderTest.java @@ -9,59 +9,108 @@ public class UCTokenProviderTest { - private static final String PREFIX = "spark.sql.catalog.cat."; - @Test - public void createReturnsFixedProviderWhenTokenIsSet() { - Map options = new HashMap<>(); - options.put(UCTokenProvider.TOKEN, "static-token"); + public void createReturnsFixedProviderForStaticType() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.TYPE, AuthConfigs.STATIC_TYPE_VALUE); + configs.put(AuthConfigs.STATIC_TOKEN, "static-token"); - UCTokenProvider provider = UCTokenProvider.create(options, PREFIX); + UCTokenProvider provider = UCTokenProvider.create(configs); assertThat(provider).isInstanceOf(FixedUCTokenProvider.class); assertThat(provider.accessToken()).isEqualTo("static-token"); } @Test - public void createReturnsOAuthProviderWhenAllOAuthKeysAreSet() { - Map options = new HashMap<>(); - options.put(UCTokenProvider.OAUTH_URI, "https://example.com/oidc/v1/token"); - options.put(UCTokenProvider.OAUTH_CLIENT_ID, "client-id"); - options.put(UCTokenProvider.OAUTH_CLIENT_SECRET, "client-secret"); + public void createReturnsOAuthProviderForOAuthType() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.TYPE, AuthConfigs.OAUTH_TYPE_VALUE); + configs.put(AuthConfigs.OAUTH_URI, "https://example.com/oidc/v1/token"); + configs.put(AuthConfigs.OAUTH_CLIENT_ID, "client-id"); + configs.put(AuthConfigs.OAUTH_CLIENT_SECRET, "client-secret"); - UCTokenProvider provider = UCTokenProvider.create(options, PREFIX); + UCTokenProvider provider = UCTokenProvider.create(configs); assertThat(provider).isInstanceOf(OAuthUCTokenProvider.class); } @Test - public void fixedProviderTakesPrecedenceOverOAuth() { - Map options = new HashMap<>(); - options.put(UCTokenProvider.TOKEN, "static-token"); - options.put(UCTokenProvider.OAUTH_URI, "https://example.com/oidc/v1/token"); - options.put(UCTokenProvider.OAUTH_CLIENT_ID, "client-id"); - options.put(UCTokenProvider.OAUTH_CLIENT_SECRET, "client-secret"); + public void createReturnsFileOidcProviderForOidcType() { + UCTokenProvider provider = UCTokenProvider.create(oidcConfigs()); - UCTokenProvider provider = UCTokenProvider.create(options, PREFIX); + assertThat(provider).isInstanceOf(FileOidcUCTokenProvider.class); + } - assertThat(provider).isInstanceOf(FixedUCTokenProvider.class); + @Test + public void createInstantiatesACustomProviderByClassName() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.TYPE, CustomTokenProvider.class.getName()); + + UCTokenProvider provider = UCTokenProvider.create(configs); + + assertThat(provider).isInstanceOf(CustomTokenProvider.class); + assertThat(provider.accessToken()).isEqualTo("custom-token"); } @Test - public void createFailsOnIncompleteOAuthConfig() { - Map options = new HashMap<>(); - options.put(UCTokenProvider.OAUTH_URI, "https://example.com/oidc/v1/token"); - options.put(UCTokenProvider.OAUTH_CLIENT_ID, "client-id"); + public void createFailsWhenTypeIsMissing() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.STATIC_TOKEN, "static-token"); - assertThatThrownBy(() -> UCTokenProvider.create(options, PREFIX)) + assertThatThrownBy(() -> UCTokenProvider.create(configs)) .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Incomplete OAuth configuration"); + .hasMessageContaining("Required configuration key 'type' is missing or empty"); + } + + @Test + public void createFailsWhenACustomClassCannotBeInstantiated() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.TYPE, "com.example.NoSuchProvider"); + + assertThatThrownBy(() -> UCTokenProvider.create(configs)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("Failed to instantiate custom UCTokenProvider"); } @Test - public void createFailsWhenNoAuthConfig() { - assertThatThrownBy(() -> UCTokenProvider.create(new HashMap<>(), PREFIX)) + public void createFailsOnIncompleteOidcConfig() { + Map configs = oidcConfigs(); + configs.remove(AuthConfigs.OIDC_TOKEN_FILE_PATH); + + assertThatThrownBy(() -> UCTokenProvider.create(configs)) .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Cannot determine UC authentication configuration"); + .hasMessageContaining("Configuration key 'oidc.tokenFilePath' is missing or empty"); + } + + @Test + public void createFailsOnIncompleteOAuthConfig() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.TYPE, AuthConfigs.OAUTH_TYPE_VALUE); + configs.put(AuthConfigs.OAUTH_URI, "https://example.com/oidc/v1/token"); + configs.put(AuthConfigs.OAUTH_CLIENT_ID, "client-id"); + + assertThatThrownBy(() -> UCTokenProvider.create(configs)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Configuration key 'oauth.clientSecret' is missing or empty"); + } + + private static Map oidcConfigs() { + Map configs = new HashMap<>(); + configs.put(AuthConfigs.TYPE, AuthConfigs.OIDC_TYPE_VALUE); + configs.put(AuthConfigs.OIDC_URI, "https://example.com/oidc/v1/token"); + configs.put(AuthConfigs.OIDC_CLIENT_ID, "client-id"); + configs.put(AuthConfigs.OIDC_TOKEN_FILE_PATH, "/var/run/secrets/token"); + return configs; + } + + /** Public so {@code Class.forName(...).getDeclaredConstructor().newInstance()} can reach it. */ + public static class CustomTokenProvider implements UCTokenProvider { + @Override + public void initialize(Map configs) {} + + @Override + public String accessToken() { + return "custom-token"; + } } }