From b56a8cf87c2edb5fa8a0a3a87ac7cf2b2a550064 Mon Sep 17 00:00:00 2001
From: Howie Wang <hongyiw@wepay.com>
Date: Tue, 26 Mar 2024 14:11:45 -0700
Subject: [PATCH] feat: Make OAuth token server configurable (#305)

---
 crates/catalog/rest/src/catalog.rs | 45 ++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs
index 52731c840..a98153fa3 100644
--- a/crates/catalog/rest/src/catalog.rs
+++ b/crates/catalog/rest/src/catalog.rs
@@ -71,7 +71,11 @@ impl RestCatalogConfig {
     }
 
     fn get_token_endpoint(&self) -> String {
-        [&self.uri, PATH_V1, "oauth", "tokens"].join("/")
+        if let Some(auth_url) = self.props.get("rest.authorization-url") {
+            auth_url.to_string()
+        } else {
+            [&self.uri, PATH_V1, "oauth", "tokens"].join("/")
+        }
     }
 
     fn namespaces_endpoint(&self) -> String {
@@ -865,8 +869,12 @@ mod tests {
     }
 
     async fn create_oauth_mock(server: &mut ServerGuard) -> Mock {
+        create_oauth_mock_with_path(server, "/v1/oauth/tokens").await
+    }
+
+    async fn create_oauth_mock_with_path(server: &mut ServerGuard, path: &str) -> Mock {
         server
-            .mock("POST", "/v1/oauth/tokens")
+            .mock("POST", path)
             .with_status(200)
             .with_body(
                 r#"{
@@ -955,6 +963,39 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_oauth_with_auth_url() {
+        let mut server = Server::new_async().await;
+        let config_mock = create_config_mock(&mut server).await;
+
+        let mut auth_server = Server::new_async().await;
+        let auth_server_path = "/some/path";
+        let oauth_mock = create_oauth_mock_with_path(&mut auth_server, auth_server_path).await;
+
+        let mut props = HashMap::new();
+        props.insert("credential".to_string(), "client1:secret1".to_string());
+        props.insert(
+            "rest.authorization-url".to_string(),
+            format!("{}{}", auth_server.url(), auth_server_path).to_string(),
+        );
+
+        let catalog = RestCatalog::new(
+            RestCatalogConfig::builder()
+                .uri(server.url())
+                .props(props)
+                .build(),
+        )
+        .await
+        .unwrap();
+
+        oauth_mock.assert_async().await;
+        config_mock.assert_async().await;
+        assert_eq!(
+            catalog.config.props.get("token"),
+            Some(&"ey000000000000".to_string())
+        );
+    }
+
     #[tokio::test]
     async fn test_config_override_prefix() {
         let mut server = Server::new_async().await;