SS-441 Add OAUTH2 SERVER URL option to Iceberg catalog connections to support Unity Catalog

src/sql-lexer/src/keywords.txt

@@ -335,6 +335,7 @@ Notices
 Null
 Nullif
 Nulls
+Oauth2
 Objects
 Of
 Offset

src/sql-parser/src/ast/defs/ddl.rs

@@ -906,6 +906,7 @@ pub enum ConnectionOptionName {
     Endpoint,
     GcpConnection,
     Host,
+    Oauth2ServerUrl,
     Password,
     Port,
     ProgressTopic,
@@ -960,6 +961,7 @@ impl ConnectionOptionName {
             | ConnectionOptionName::Endpoint
             | ConnectionOptionName::GcpConnection
             | ConnectionOptionName::Host
+            | ConnectionOptionName::Oauth2ServerUrl
             | ConnectionOptionName::Port
             | ConnectionOptionName::ProgressTopic
             | ConnectionOptionName::ProgressTopicReplicationFactor
@@ -994,6 +996,7 @@ impl AstDisplay for ConnectionOptionName {
             ConnectionOptionName::Endpoint => "ENDPOINT",
             ConnectionOptionName::GcpConnection => "GCP CONNECTION",
             ConnectionOptionName::Host => "HOST",
+            ConnectionOptionName::Oauth2ServerUrl => "OAUTH2 SERVER URL",
             ConnectionOptionName::Password => "PASSWORD",
             ConnectionOptionName::Port => "PORT",
             ConnectionOptionName::ProgressTopic => "PROGRESS TOPIC",

src/sql-parser/src/parser.rs

@@ -2974,6 +2974,7 @@ impl<'a> Parser<'a> {
                 ENDPOINT,
                 GCP,
                 HOST,
+                OAUTH2,
                 PASSWORD,
                 PORT,
                 PUBLIC,
@@ -3031,6 +3032,10 @@ impl<'a> Parser<'a> {
                     ConnectionOptionName::GcpConnection
                 }
                 HOST => ConnectionOptionName::Host,
+                OAUTH2 => {
+                    self.expect_keywords(&[SERVER, URL])?;
+                    ConnectionOptionName::Oauth2ServerUrl
+                }
                 PASSWORD => ConnectionOptionName::Password,
                 PORT => ConnectionOptionName::Port,
                 PUBLIC => {

src/sql-parser/tests/testdata/ddl

@@ -578,6 +578,13 @@ CREATE CONNECTION awsconn TO AWS (ACCESS KEY ID = 'id', ENDPOINT = 'endpoint', R
 =>
 CreateConnection(CreateConnectionStatement { name: UnresolvedItemName([Ident("awsconn")]), connection_type: Aws, if_not_exists: false, values: [ConnectionOption { name: AccessKeyId, value: Some(Value(String("id"))) }, ConnectionOption { name: Endpoint, value: Some(Value(String("endpoint"))) }, ConnectionOption { name: Region, value: Some(Value(String("region"))) }, ConnectionOption { name: AssumeRoleArn, value: Some(Value(String("role-arn"))) }, ConnectionOption { name: SecretAccessKey, value: Some(Value(String("key"))) }, ConnectionOption { name: SessionToken, value: Some(Value(String("token"))) }], with_options: [] })
 
+parse-statement
+CREATE CONNECTION icebergconn TO ICEBERG CATALOG (CATALOG TYPE 'rest', URL 'https://example/api/catalog', CREDENTIAL 'id:secret', OAUTH2 SERVER URL 'https://example/oidc/v1/token', SCOPE 'all-apis', WAREHOUSE 'wh')
+----
+CREATE CONNECTION icebergconn TO ICEBERG CATALOG (CATALOG TYPE = 'rest', URL = 'https://example/api/catalog', CREDENTIAL = 'id:secret', OAUTH2 SERVER URL = 'https://example/oidc/v1/token', SCOPE = 'all-apis', WAREHOUSE = 'wh')
+=>
+CreateConnection(CreateConnectionStatement { name: UnresolvedItemName([Ident("icebergconn")]), connection_type: IcebergCatalog, if_not_exists: false, values: [ConnectionOption { name: CatalogType, value: Some(Value(String("rest"))) }, ConnectionOption { name: Url, value: Some(Value(String("https://example/api/catalog"))) }, ConnectionOption { name: Credential, value: Some(Value(String("id:secret"))) }, ConnectionOption { name: Oauth2ServerUrl, value: Some(Value(String("https://example/oidc/v1/token"))) }, ConnectionOption { name: Scope, value: Some(Value(String("all-apis"))) }, ConnectionOption { name: Warehouse, value: Some(Value(String("wh"))) }], with_options: [] })
+
 parse-statement
 CREATE CONNECTION privatelinkconn TO AWS PRIVATELINK (SERVICE NAME 'com.amazonaws.vpce.us-east-1.vpce-svc-0e123abc123198abc', AVAILABILITY ZONES ('use1-az1', 'use1-az4'))
 ----
@@ -3033,7 +3040,7 @@ CreateConnection(CreateConnectionStatement { name: UnresolvedItemName([Ident("my
 parse-statement
 CREATE CONNECTION my_ssh_tunnel FOR SSH TUNNEL (PUBLIC KEY 3 = nope)
 ----
-error: Expected one of ACCESS or ASSUME or AVAILABILITY or AWS or BROKER or BROKERS or CATALOG or CREDENTIAL or DATABASE or ENDPOINT or GCP or HOST or PASSWORD or PORT or PUBLIC or PROGRESS or REGION or REGISTRY or SASL or SCOPE or SECRET or SECURITY or SERVICE or SESSION or SSH or SSL or URL or USER or USERNAME or WAREHOUSE, found left parenthesis
+error: Expected one of ACCESS or ASSUME or AVAILABILITY or AWS or BROKER or BROKERS or CATALOG or CREDENTIAL or DATABASE or ENDPOINT or GCP or HOST or OAUTH2 or PASSWORD or PORT or PUBLIC or PROGRESS or REGION or REGISTRY or SASL or SCOPE or SECRET or SECURITY or SERVICE or SESSION or SSH or SSL or URL or USER or USERNAME or WAREHOUSE, found left parenthesis
 CREATE CONNECTION my_ssh_tunnel FOR SSH TUNNEL (PUBLIC KEY 3 = nope)
                                                ^
 

src/sql/src/plan/statement/ddl/connection.rs

@@ -62,6 +62,7 @@ generate_extracted_config!(
     (Endpoint, String),
     (GcpConnection, with_options::Object),
     (Host, String),
+    (Oauth2ServerUrl, String),
     (Password, with_options::Secret),
     (Port, u16),
     (ProgressTopic, String),
@@ -194,6 +195,7 @@ pub(super) fn validate_options_per_connection_type(
             CatalogType,
             Credential,
             GcpConnection,
+            Oauth2ServerUrl,
             Scope,
             Url,
             Warehouse,
@@ -697,6 +699,11 @@ impl ConnectionOptionExtracted {
                                 "invalid CONNECTION: ICEBERG s3tablesrest connections do not support GCP CONNECTION"
                             );
                         }
+                        if self.oauth2_server_url.is_some() {
+                            sql_bail!(
+                                "invalid CONNECTION: ICEBERG s3tablesrest connections do not support OAUTH2 SERVER URL"
+                            );
+                        }
                         let Some(warehouse) = warehouse else {
                             sql_bail!(
                                 "invalid CONNECTION: ICEBERG s3tablesrest connections must specify WAREHOUSE"
@@ -726,8 +733,14 @@ impl ConnectionOptionExtracted {
                             (Some(credential), None) => IcebergCatalogAuth::OAuth {
                                 credential,
                                 scope: self.scope.clone(),
+                                server_url: self.oauth2_server_url.clone(),
                             },
                             (None, Some(gcp_connection)) => {
+                                if self.oauth2_server_url.is_some() {
+                                    sql_bail!(
+                                        "invalid CONNECTION: OAUTH2 SERVER URL applies to CREDENTIAL auth, not GCP CONNECTION"
+                                    );
+                                }
                                 /// All BigLake Iceberg REST Catalogs use the same catalog URI.
                                 const BIGLAKE_CATALOG_URI: &str =
                                     "https://biglake.googleapis.com/iceberg/v1/restcatalog";

src/storage-types/src/connections.rs

@@ -93,6 +93,9 @@ pub mod string_or_secret;
 
 const REST_CATALOG_PROP_SCOPE: &str = "scope";
 const REST_CATALOG_PROP_CREDENTIAL: &str = "credential";
+/// Overrides the OAuth2 token endpoint. Spelled `uri` because that is the property name the
+/// Iceberg REST clients agree on, even though the SQL option says `URL`.
+const REST_CATALOG_PROP_OAUTH2_SERVER_URI: &str = "oauth2-server-uri";
 
 /// A credential loader that wraps an aws-sdk-rust credentials provider for use with
 /// iceberg/OpenDAL. This allows us to provide refreshable credentials from the AWS SDK
@@ -611,6 +614,13 @@ pub enum IcebergCatalogAuth<C: ConnectionAccess = InlinedConnection> {
         credential: StringOrSecret,
         /// OAuth2 scope
         scope: Option<String>,
+        /// Where to exchange `credential` for a bearer token.
+        ///
+        /// `None` uses the endpoint the Iceberg REST specification defines relative to the
+        /// catalog URL, `<url>/v1/oauth/tokens`. Catalogs that host their token endpoint
+        /// elsewhere, or behind an auth gateway that will not serve an unauthenticated
+        /// exchange, need this override.
+        server_url: Option<String>,
     },
     Gcp(GcpConnectionReference<C>),
 }
@@ -636,9 +646,15 @@ impl<R: ConnectionResolver> IntoInlineConnection<IcebergCatalogAuth, R>
     fn into_inline_connection(self, r: R) -> IcebergCatalogAuth {
         match self {
             IcebergCatalogAuth::Gcp(x) => IcebergCatalogAuth::Gcp(x.into_inline_connection(&r)),
-            IcebergCatalogAuth::OAuth { credential, scope } => {
-                IcebergCatalogAuth::OAuth { credential, scope }
-            }
+            IcebergCatalogAuth::OAuth {
+                credential,
+                scope,
+                server_url,
+            } => IcebergCatalogAuth::OAuth {
+                credential,
+                scope,
+                server_url,
+            },
         }
     }
 }
@@ -918,7 +934,11 @@ impl IcebergCatalogConnection<InlinedConnection> {
         // which happen at different stages of the [`RestCatalogBuilder`] -> [`RestCatalog`]
         // construction pipeline.
         let (storage_factory, custom_authenticator) = match &rest.auth {
-            IcebergCatalogAuth::OAuth { credential, scope } => {
+            IcebergCatalogAuth::OAuth {
+                credential,
+                scope,
+                server_url,
+            } => {
                 let credential = credential
                     .get_string(
                         in_task,
@@ -928,6 +948,13 @@ impl IcebergCatalogConnection<InlinedConnection> {
                     .map_err(|e| anyhow!("failed to read Iceberg catalog credential: {e}"))?;
                 props.insert(REST_CATALOG_PROP_CREDENTIAL.to_string(), credential);
 
+                if let Some(server_url) = server_url {
+                    props.insert(
+                        REST_CATALOG_PROP_OAUTH2_SERVER_URI.to_string(),
+                        server_url.clone(),
+                    );
+                }
+
                 if let Some(scope) = scope {
                     props.insert(REST_CATALOG_PROP_SCOPE.to_string(), scope.clone());
                 }

test/iceberg/mzcompose.py

@@ -95,6 +95,20 @@ def workflow_gcp_connection_validation(c: Composition) -> None:
     )
 
 
+def workflow_oauth2_server_url(c: Composition) -> None:
+    """OAUTH2 SERVER URL redirects a REST catalog connection's token exchange
+    away from the endpoint the Iceberg specification derives from the catalog
+    URL. Catalogs behind an auth gateway that will not serve an unauthenticated
+    exchange, such as Databricks Unity Catalog, need it. This exercises
+    connection planning only and needs no Iceberg backend."""
+    c.down(destroy_volumes=True)
+    c.up("materialized")
+
+    c.run_testdrive_files(
+        "oauth2-server-url.td",
+    )
+
+
 def workflow_mode_append(c: Composition) -> None:
     key = _setup(c)
 

test/iceberg/oauth2-server-url.td

@@ -0,0 +1,54 @@
+# Copyright Materialize, Inc. and contributors. All rights reserved.
+#
+# Use of this software is governed by the Business Source License
+# included in the LICENSE file at the root of this repository.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0.
+
+# OAUTH2 SERVER URL overrides where a REST catalog connection exchanges its
+# CREDENTIAL for a bearer token. Without it the endpoint is derived from the
+# catalog URL as `<url>/v1/oauth/tokens`, which the Iceberg REST specification
+# defines but not every catalog serves.
+
+# The accepted case below points at a host that does not exist, so skip the
+# connect-on-create validation. Everything asserted here is decided at plan time.
+$ postgres-execute connection=postgres://mz_system:materialize@${testdrive.materialize-internal-sql-addr}
+ALTER SYSTEM SET enable_default_connection_validation = false
+
+> CREATE SECRET oauth_cred AS 'client-id:client-secret'
+
+> CREATE CONNECTION oauth_override TO ICEBERG CATALOG (
+    CATALOG TYPE = 'REST',
+    URL = 'https://example.invalid/api/catalog',
+    CREDENTIAL = SECRET oauth_cred,
+    OAUTH2 SERVER URL = 'https://example.invalid/oidc/v1/token',
+    SCOPE = 'all-apis',
+    WAREHOUSE = 'wh'
+  );
+
+> SELECT name FROM mz_connections WHERE name = 'oauth_override'
+oauth_override
+
+# The option is only meaningful for the OAuth exchange, which s3tablesrest does
+# not perform: it authenticates with SigV4 through an AWS connection instead.
+! CREATE CONNECTION oauth_on_s3tables TO ICEBERG CATALOG (
+    CATALOG TYPE = 's3tablesrest',
+    URL = 'https://example.invalid/api/catalog',
+    OAUTH2 SERVER URL = 'https://example.invalid/oidc/v1/token',
+    WAREHOUSE = 'wh'
+  );
+contains:s3tablesrest connections do not support OAUTH2 SERVER URL
+
+# Omitting it stays valid, leaving the specification-derived endpoint in place.
+> CREATE CONNECTION oauth_default TO ICEBERG CATALOG (
+    CATALOG TYPE = 'REST',
+    URL = 'https://example.invalid/api/catalog',
+    CREDENTIAL = SECRET oauth_cred,
+    WAREHOUSE = 'wh'
+  );
+
+> DROP CONNECTION oauth_default
+> DROP CONNECTION oauth_override
+> DROP SECRET oauth_cred