org.projectnessie.versioned.transfer.ExportCommon Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nessie-versioned-transfer Show documentation
The newest version!
/*
 * Copyright (C) 2022 Dremio
 *
 * 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 org.projectnessie.versioned.transfer;

import static org.projectnessie.versioned.storage.common.objtypes.StandardObjType.UNIQUE;
import static org.projectnessie.versioned.transfer.ExportImportConstants.EXPORT_METADATA;
import static org.projectnessie.versioned.transfer.ExportImportConstants.HEADS_AND_FORKS;
import static org.projectnessie.versioned.transfer.ExportImportConstants.REPOSITORY_DESCRIPTION;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import org.projectnessie.api.NessieVersion;
import org.projectnessie.nessie.relocated.protobuf.UnsafeByteOperations;
import org.projectnessie.versioned.storage.common.logic.RepositoryDescription;
import org.projectnessie.versioned.storage.common.objtypes.CustomObjType;
import org.projectnessie.versioned.storage.common.objtypes.UniqueIdObj;
import org.projectnessie.versioned.storage.common.objtypes.UpdateableObj;
import org.projectnessie.versioned.storage.common.persist.Obj;
import org.projectnessie.versioned.storage.common.persist.ObjId;
import org.projectnessie.versioned.storage.common.persist.ObjType;
import org.projectnessie.versioned.storage.serialize.SmileSerialization;
import org.projectnessie.versioned.transfer.files.ExportFileSupplier;
import org.projectnessie.versioned.transfer.related.TransferRelatedObjects;
import org.projectnessie.versioned.transfer.serialize.TransferTypes;
import org.projectnessie.versioned.transfer.serialize.TransferTypes.ExportMeta;
import org.projectnessie.versioned.transfer.serialize.TransferTypes.ExportVersion;
import org.projectnessie.versioned.transfer.serialize.TransferTypes.HeadsAndForks;
import org.projectnessie.versioned.transfer.serialize.TransferTypes.RelatedObj;
import org.projectnessie.versioned.transfer.serialize.TransferTypes.RepositoryDescriptionProto;

abstract class ExportCommon {
  final ExportFileSupplier exportFiles;
  final NessieExporter exporter;
  final TransferRelatedObjects transferRelatedObjects;
  final ExportVersion exportVersion;

  private Batcher genericObjBatcher;

  ExportCommon(
      ExportFileSupplier exportFiles, NessieExporter exporter, ExportVersion exportVersion) {
    this.exportFiles = exportFiles;
    this.exporter = exporter;

    this.exportVersion = verifyExportVersion(exportVersion);

    this.transferRelatedObjects =
        new CompositeTransferRelatedObjects(exporter.genericObjectResolvers());
  }

  private static ExportVersion verifyExportVersion(ExportVersion exportVersion) {
    if (exportVersion == null) {
      exportVersion = ExportVersion.UNRECOGNIZED;
    }
    switch (exportVersion) {
      case V1:
        throw new IllegalArgumentException(
            "Cannot export using export-version " + exportVersion.getNumber());
      case V2:
      case V3:
        break;
      default:
        throw new IllegalArgumentException(
            "Unimplemented export-version " + exportVersion.getNumber());
    }

    return exportVersion;
  }

  ExportMeta exportRepo() throws IOException {
    ExportContext exportContext = createExportContext(getExportVersion());
    try {
      // Generic objects are written by multiple export tasks, so we need to open (and close) its
      // `Batcher` instance here.
      try (Batcher genericObjBatcher =
          new Batcher<>(exporter.commitBatchSize(), objs -> mapGenericObjs(objs, exportContext))) {
        this.genericObjBatcher = genericObjBatcher;

        exporter.progressListener().progress(ProgressEvent.STARTED);
        prepare(exportContext);

        exporter.progressListener().progress(ProgressEvent.START_COMMITS);
        HeadsAndForks headsAndForks = exportCommits(exportContext);
        exportContext.commitOutput.finishCurrentFile();
        writeHeadsAndForks(headsAndForks);
        exporter.progressListener().progress(ProgressEvent.END_COMMITS);

        exporter.progressListener().progress(ProgressEvent.START_NAMED_REFERENCES);
        exportReferences(exportContext);
        exportContext.namedReferenceOutput.finishCurrentFile();
        exporter.progressListener().progress(ProgressEvent.END_NAMED_REFERENCES);

        exporter.progressListener().progress(ProgressEvent.START_META);
        writeRepositoryInformation();

      } finally {
        this.genericObjBatcher = null;
      }

      // Must close the `Batcher` before `ExportContext.finish()` to have the right values in
      // `ExportMeta`.

      ExportMeta meta = exportContext.finish();
      writeExportMeta(meta);
      exporter.progressListener().progress(ProgressEvent.END_META, meta);

      exporter.progressListener().progress(ProgressEvent.FINISHED);

      return meta;
    } finally {
      // Ensure that all output streams are closed.
      exportContext.closeSilently();
    }
  }

  final long currentTimestampMillis() {
    return exporter.clock().millis();
  }

  ExportContext createExportContext(ExportVersion exportVersion) {
    return new ExportContext(
        exportFiles,
        exporter,
        ExportMeta.newBuilder()
            .setNessieVersion(NessieVersion.NESSIE_VERSION)
            .setCreatedMillisEpoch(currentTimestampMillis())
            .setVersion(exportVersion));
  }

  ExportVersion getExportVersion() {
    return exportVersion;
  }

  abstract void prepare(ExportContext exportContext);

  abstract void exportReferences(ExportContext exportContext);

  abstract HeadsAndForks exportCommits(ExportContext exportContext);

  void writeRepositoryInformation() throws IOException {
    handleGenericObjs(transferRelatedObjects.repositoryRelatedObjects());

    RepositoryDescription repositoryDescription =
        exporter.repositoryLogic().fetchRepositoryDescription();
    if (repositoryDescription != null) {
      writeRepositoryDescription(
          TransferTypes.RepositoryDescriptionProto.newBuilder()
              .putAllProperties(repositoryDescription.properties())
              .setRepositoryId(exporter.persist().config().repositoryId())
              .setRepositoryCreatedTimestampMillis(
                  repositoryDescription.repositoryCreatedTime().toEpochMilli())
              .setOldestCommitTimestampMillis(
                  repositoryDescription.oldestPossibleCommitTime().toEpochMilli())
              .setDefaultBranchName(repositoryDescription.defaultBranchName())
              .build());
    }
  }

  void writeExportMeta(ExportMeta meta) throws IOException {
    try (OutputStream output = exportFiles.newFileOutput(EXPORT_METADATA)) {
      meta.writeTo(output);
    }
  }

  void writeHeadsAndForks(HeadsAndForks hf) throws IOException {
    try (OutputStream output = exportFiles.newFileOutput(HEADS_AND_FORKS)) {
      hf.writeTo(output);
    }
  }

  void writeRepositoryDescription(RepositoryDescriptionProto repositoryDescription)
      throws IOException {
    try (OutputStream output = exportFiles.newFileOutput(REPOSITORY_DESCRIPTION)) {
      repositoryDescription.writeTo(output);
    }
  }

  void handleGenericObjs(Set ids) {
    if (ids.isEmpty()) {
      return;
    }

    ObjId[] idArray = ids.toArray(ObjId[]::new);
    Obj[] objs = exporter.persist().fetchObjsIfExist(idArray);
    Arrays.stream(objs).filter(Objects::nonNull).forEach(genericObjBatcher::add);
  }

  private void mapGenericObjs(List objs, ExportContext exportContext) {
    for (Obj o : objs) {
      RelatedObj custom = mapGenericObj(o);
      exportContext.writeGeneric(custom);
      // use the same progress value than commits
      exporter.progressListener().progress(ProgressEvent.GENERIC_WRITTEN);
    }
  }

  private RelatedObj mapGenericObj(Obj obj) {
    RelatedObj.Builder genericObj = RelatedObj.newBuilder().setId(obj.id().asBytes());
    ObjType objType = obj.type();
    if (objType.equals(UNIQUE)) {
      UniqueIdObj uniqueIdObj = (UniqueIdObj) obj;
      genericObj
          .setTypeName(UNIQUE.shortName())
          .setSpace(uniqueIdObj.space())
          .setData(uniqueIdObj.value());
    } else if (objType instanceof CustomObjType) {
      // 'obj' can be a 'GenericObj' or, if the actual `ObjType` implementations are on the
      // classpath, the "real" instance, so we need to handle both cases.
      byte[] serialized =
          SmileSerialization.serializeObj(
              obj, compression -> genericObj.setCompression(compression.valueString()));
      genericObj
          .setTypeName(objType.shortName())
          .setEncoding(TransferTypes.Encoding.Smile)
          .setData(UnsafeByteOperations.unsafeWrap(serialized));
      if (obj instanceof UpdateableObj) {
        UpdateableObj updateableObj = (UpdateableObj) obj;
        String versionToken = updateableObj.versionToken();
        if (versionToken != null) {
          genericObj.setVersionToken(versionToken);
        }
      }
    } else {
      throw new IllegalArgumentException("Unhandled object type " + objType.name());
    }
    return genericObj.build();
  }
}