io.trino.gateway.proxyserver.ProxyResponseHandler Maven / Gradle / Ivy

Go to download
/*
 * Licensed 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 io.trino.gateway.proxyserver;

import com.google.common.collect.ListMultimap;
import io.airlift.http.client.HeaderName;
import io.airlift.http.client.Request;
import io.airlift.http.client.Response;
import io.airlift.http.client.ResponseHandler;
import io.trino.gateway.proxyserver.ProxyResponseHandler.ProxyResponse;

import java.io.IOException;

import static com.google.common.io.ByteStreams.toByteArray;
import static java.util.Objects.requireNonNull;

public class ProxyResponseHandler
        implements ResponseHandler
{
    @Override
    public ProxyResponse handleException(Request request, Exception exception)
    {
        throw new ProxyException("Request to remote Trino server failed", exception);
    }

    @Override
    public ProxyResponse handle(Request request, Response response)
    {
        try {
            return new ProxyResponse(response.getStatusCode(), response.getHeaders(), toByteArray(response.getInputStream()));
        }
        catch (IOException e) {
            throw new ProxyException("Failed reading response from remote Trino server", e);
        }
    }

    public record ProxyResponse(
            int statusCode,
            ListMultimap headers,
            byte[] body)
    {
        public ProxyResponse
        {
            requireNonNull(headers, "headers is null");
        }
    }
}