apoc.util.google.cloud.GCStorageURLConnection Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of apoc-common Show documentation
Data types package for Neo4j Procedures
There is a newer version: 5.25.1
/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [http://neo4j.com]
 *
 * This file is part of Neo4j.
 *
 * 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 apoc.util.google.cloud;

import com.google.cloud.storage.Blob;
import com.google.cloud.storage.BlobId;
import com.google.cloud.storage.Storage;
import com.google.cloud.storage.StorageOptions;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLConnection;
import java.nio.channels.Channels;
import java.util.Collections;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.commons.lang3.StringUtils;
import org.neo4j.util.VisibleForTesting;

public class GCStorageURLConnection extends URLConnection {

    enum AuthType {
        NONE,
        PRIVATE_KEY,
        GCP_ENVIRONMENT
    }

    private Blob blob;

    public GCStorageURLConnection(URL url) {
        super(url);
    }

    @VisibleForTesting
    public Storage getStorage(URI uri) {
        Storage storage;
        Map queryParams = getQueryParams(uri);
        AuthType authenticationType =
                AuthType.valueOf(queryParams.getOrDefault("authenticationType", AuthType.NONE.toString()));
        switch (authenticationType) {
            case PRIVATE_KEY:
                String googleAppCredentialsEnv = System.getenv("GOOGLE_APPLICATION_CREDENTIALS");
                if (StringUtils.isBlank(googleAppCredentialsEnv)) {
                    throw new RuntimeException(
                            "You must set the env variable GOOGLE_APPLICATION_CREDENTIALS as described here: https://cloud.google.com/storage/docs/reference/libraries#client-libraries-install-java");
                }
                // fall through
            case GCP_ENVIRONMENT:
                storage = StorageOptions.getDefaultInstance().getService();
                break;
            default:
                storage = StorageOptions.getUnauthenticatedInstance().getService();
        }
        return storage;
    }

    private Map getQueryParams(URI uri) {
        Map queryParams;
        if (StringUtils.isBlank(uri.getQuery())) {
            queryParams = Collections.emptyMap();
        } else {
            queryParams = Stream.of(uri.getQuery().split("&"))
                    .map(e -> e.split("="))
                    .collect(Collectors.toMap(e -> e[0], e -> e[1]));
        }
        return queryParams;
    }

    @Override
    public void connect() {
        try {
            URI uri = url.toURI();
            if (StringUtils.isBlank(uri.getPath())) {
                throw new RuntimeException("Please provide the file name");
            }
            blob = getStorage(uri)
                    .get(BlobId.of(uri.getAuthority(), uri.getPath().substring(1)));
            connected = true;
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public String getContentType() {
        return blob.getContentType();
    }

    @Override
    public InputStream getInputStream() {
        if (!connected) {
            connect();
        }
        InputStream in;
        if (blob == null) {
            in = new ByteArrayInputStream(new byte[0]);
        } else {
            in = Channels.newInputStream(blob.reader());
        }
        return in;
    }
}