如何基于OkHttp3为Apache Drill的REST API查询添加OAuth2认证支持

核心实现思路 #

你可以通过OkHttp拦截器统一注入OAuth2访问令牌，同时搭配令牌管理器处理令牌的缓存、自动刷新逻辑，无需修改原有查询请求的业务代码。

具体实现步骤 #

• 新增OAuth2配置项
  在HTTP存储插件的配置类中新增OAuth2相关配置字段，包括：令牌获取地址tokenEndpoint、客户端IDclientId、客户端密钥clientSecret、授权范围scope，同时复用Drill现有敏感字段加密逻辑处理clientSecret，避免明文存储。
• 实现令牌管理工具类
  负责令牌的获取、缓存、过期判断、失效重置逻辑，示例代码如下：

import okhttp3.FormBody;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.time.Instant;

public class OAuth2TokenManager {
    private final OkHttpClient httpClient;
    private final String tokenEndpoint;
    private final String clientId;
    private final String clientSecret;
    private final String scope;
    private String cachedToken;
    private Instant expiresAt;
    // 提前60秒刷新令牌，避免边界过期问题
    private static final long EXPIRY_OFFSET_SECONDS = 60;

    public OAuth2TokenManager(OkHttpClient httpClient, String tokenEndpoint, String clientId, String clientSecret, String scope) {
        this.httpClient = httpClient;
        this.tokenEndpoint = tokenEndpoint;
        this.clientId = clientId;
        this.clientSecret = clientSecret;
        this.scope = scope;
    }

    public synchronized String getValidToken() throws IOException {
        if (cachedToken == null || Instant.now().isAfter(expiresAt.minusSeconds(EXPIRY_OFFSET_SECONDS))) {
            fetchNewToken();
        }
        return cachedToken;
    }

    public synchronized void invalidateToken() {
        this.cachedToken = null;
        this.expiresAt = Instant.MIN;
    }

    private void fetchNewToken() throws IOException {
        FormBody formBody = new FormBody.Builder()
                .add("grant_type", "client_credentials")
                .add("client_id", clientId)
                .add("client_secret", clientSecret)
                .add("scope", scope)
                .build();
        Request request = new Request.Builder()
                .url(tokenEndpoint)
                .post(formBody)
                .build();
        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("获取OAuth2令牌失败，响应码：" + response.code());
            }
            JsonNode jsonNode = new ObjectMapper().readTree(response.body().string());
            this.cachedToken = jsonNode.get("access_token").asText();
            long expiresIn = jsonNode.get("expires_in").asLong();
            this.expiresAt = Instant.now().plusSeconds(expiresIn);
        }
    }
}

• 实现OAuth2认证拦截器
  在OkHttp请求链路中自动注入有效令牌，同时处理401未授权的重试逻辑，示例代码如下：

import okhttp3.Interceptor;
import okhttp3.Request;
import okhttp3.Response;
import java.io.IOException;

public class OAuth2Interceptor implements Interceptor {
    private final OAuth2TokenManager tokenManager;

    public OAuth2Interceptor(OAuth2TokenManager tokenManager) {
        this.tokenManager = tokenManager;
    }

    @Override
    public Response intercept(Chain chain) throws IOException {
        Request originalRequest = chain.request();
        // 注入Bearer令牌
        Request requestWithAuth = originalRequest.newBuilder()
                .header("Authorization", "Bearer " + tokenManager.getValidToken())
                .build();
        Response response = chain.proceed(requestWithAuth);
        // 处理401情况，触发令牌刷新后重试一次
        if (response.code() == 401) {
            response.close();
            tokenManager.invalidateToken();
            Request newRequestWithAuth = originalRequest.newBuilder()
                    .header("Authorization", "Bearer " + tokenManager.getValidToken())
                    .build();
            return chain.proceed(newRequestWithAuth);
        }
        return response;
    }
}

• 集成到现有HTTP请求逻辑中
  在HttpStoragePlugin初始化OkHttpClient实例时，判断如果配置开启了OAuth2认证，就将上述拦截器添加到OkHttpClient的拦截器链中即可。

注意事项 #

• 目前上述代码适配的是OAuth2 客户端模式（Client Credentials），这是服务端调用第三方API最常用的授权模式，不建议对接需要人工交互的授权码模式接口，不适合Drill的批量查询场景。
• 令牌管理器的核心方法加了synchronized同步锁，避免多查询请求并发时重复调用令牌接口导致的资源浪费。
• 可以根据实际业务需求调整令牌提前刷新的时间偏移量，避免临界时间点使用过期令牌。

内容的提问来源于stack exchange，提问作者cgivre