google-nomulus/java/google/registry/storage/drive/DriveConnection.java

// Copyright 2016 The Domain Registry Authors. All Rights Reserved.
//
// 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 google.registry.storage.drive;

import com.google.api.client.http.ByteArrayContent;
import com.google.api.services.drive.Drive;
import com.google.api.services.drive.Drive.Children;
import com.google.api.services.drive.model.ChildList;
import com.google.api.services.drive.model.ChildReference;
import com.google.api.services.drive.model.File;
import com.google.api.services.drive.model.ParentReference;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.net.MediaType;

import java.io.IOException;
import java.util.List;

import javax.annotation.Nullable;
import javax.inject.Inject;

/** Class encapsulating parameters and state for accessing the Drive API. */
public class DriveConnection {

  private static final MediaType GOOGLE_FOLDER =
      MediaType.create("application", "vnd.google-apps.folder");

  /** Drive client instance wrapped by this class. */
  @Inject Drive drive;
  @Inject public DriveConnection() {}

  /**
   * Creates a folder with the given parent.
   *
   * @returns the folder id.
   */
  public String createFolder(String title, String parentFolderId) throws IOException {
    return drive.files()
        .insert(createFileReference(title, GOOGLE_FOLDER, parentFolderId))
        .execute()
        .getId();
  }

  /**
   * Creates a file with the given parent.
   *
   * @returns the file id.
   */
  public String createFile(String title, MediaType mimeType, String parentFolderId, byte[] bytes)
      throws IOException {
    return drive.files()
        .insert(
            createFileReference(title, mimeType, parentFolderId),
            new ByteArrayContent(mimeType.toString(), bytes))
        .execute()
        .getId();
  }

  /**
   * Creates a file with the given parent or updates the existing one if a file already exists with
   * that same title and parent.
   *
   * @throws IllegalStateException if multiple files with that name exist in the given folder.
   * @throws IOException if communication with Google Drive fails for any reason.
   * @returns the file id.
   */
  public String createOrUpdateFile(
      String title,
      MediaType mimeType,
      String parentFolderId,
      byte[] bytes) throws IOException {
    List<String> existingFiles = listFiles(parentFolderId, String.format("title = '%s'", title));
    if (existingFiles.size() > 1) {
      throw new IllegalStateException(String.format(
          "Could not update file '%s' in Drive folder id '%s' because multiple files with that "
              + "name already exist.",
          title,
          parentFolderId));
    }
    return existingFiles.isEmpty()
        ? createFile(title, mimeType, parentFolderId, bytes)
        : updateFile(existingFiles.get(0), title, mimeType, bytes);
  }

  /**
   * Updates the file with the given id in place, setting the title, content, and mime type to
   * the newly specified values.
   *
   * @returns the file id.
   */
  public String updateFile(String fileId, String title, MediaType mimeType, byte[] bytes)
      throws IOException {
    File file = new File().setTitle(title);
    return drive.files()
        .update(fileId, file, new ByteArrayContent(mimeType.toString(), bytes))
        .execute()
        .getId();
  }
  /**
   * Returns a list of Drive file ids for all files in Google Drive in the folder with the
   * specified id.
   */
  public List<String> listFiles(String parentFolderId) throws IOException {
    return listFiles(parentFolderId, null);
  }

  /**
   * Returns a list of Drive file ids for all files in Google Drive in the folder with the
   * specified id and matching the given Drive query.
   *
   * @see <a href="https://developers.google.com/drive/web/search-parameters">The query format</a>
   */
  public List<String> listFiles(String parentFolderId, String query) throws IOException {
    ImmutableList.Builder<String> result = new ImmutableList.Builder<>();
    Children.List req = drive.children().list(parentFolderId);
    if (!Strings.isNullOrEmpty(query)) {
      req.setQ(query);
    }
    do {
      ChildList files = req.execute();
      for (ChildReference child : files.getItems()) {
        result.add(child.getId());
      }
      req.setPageToken(files.getNextPageToken());
    } while (!Strings.isNullOrEmpty(req.getPageToken()));
    return result.build();
  }

  /** Constructs an object representing a file (or folder) with a given title and parent. */
  private File createFileReference(
      String title, MediaType mimeType, @Nullable String parentFolderId) {
    return new File()
        .setTitle(title)
        .setMimeType(mimeType.toString())
        .setParents(parentFolderId == null
            ? null
            : ImmutableList.of(new ParentReference().setId(parentFolderId)));
  }
}