Skip to content

Commit

Permalink
OAK-10670: add support for service principal in oak-upgrade (apache#1329
Browse files Browse the repository at this point in the history
)

* OAK-10670: add service principal support in oak upgrade

* OAK-10670: add service principal support in oak upgrade

* OAK-10670: add service principal support in oak upgrade

* OAK-10670: add service principal support in oak upgrade

* OAK-10670: minor change

* OAK-10670: delete tempDir after closing filestore

* OAK-10670: add license to new files
  • Loading branch information
t-rana authored Mar 4, 2024
1 parent 00255b1 commit 8f97e2a
Show file tree
Hide file tree
Showing 8 changed files with 449 additions and 37 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -16,17 +16,6 @@
*/
package org.apache.jackrabbit.oak.segment.azure;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.Paths;
import java.security.InvalidKeyException;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.List;

import com.azure.core.credential.TokenRequestContext;
import com.azure.identity.ClientSecretCredential;
import com.azure.identity.ClientSecretCredentialBuilder;
Expand All @@ -41,13 +30,25 @@
import com.microsoft.azure.storage.blob.CloudBlob;
import com.microsoft.azure.storage.blob.CloudBlobContainer;
import com.microsoft.azure.storage.blob.CloudBlobDirectory;
import com.microsoft.azure.storage.blob.LeaseStatus;
import com.microsoft.azure.storage.blob.ListBlobItem;
import org.apache.jackrabbit.oak.commons.Buffer;
import org.apache.jackrabbit.oak.segment.spi.RepositoryNotReachableException;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.Paths;
import java.security.InvalidKeyException;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.List;

public final class AzureUtilities {
public static final String AZURE_ACCOUNT_NAME = "AZURE_ACCOUNT_NAME";
public static final String AZURE_SECRET_KEY = "AZURE_SECRET_KEY";
Expand Down Expand Up @@ -110,7 +111,7 @@ public static void deleteAllEntries(CloudBlobDirectory directory) throws IOExcep
}

public static CloudBlobDirectory cloudBlobDirectoryFrom(StorageCredentials credentials,
String uri, String dir) throws URISyntaxException, StorageException {
String uri, String dir) throws URISyntaxException, StorageException {
StorageUri storageUri = new StorageUri(new URI(uri));
CloudBlobContainer container = new CloudBlobContainer(storageUri, credentials);

Expand All @@ -120,7 +121,7 @@ public static CloudBlobDirectory cloudBlobDirectoryFrom(StorageCredentials crede
}

public static CloudBlobDirectory cloudBlobDirectoryFrom(String connection, String containerName,
String dir) throws InvalidKeyException, URISyntaxException, StorageException {
String dir) throws InvalidKeyException, URISyntaxException, StorageException {
CloudStorageAccount cloud = CloudStorageAccount.parse(connection);
CloudBlobContainer container = cloud.createCloudBlobClient().getContainerReference(containerName);
container.createIfNotExists();
Expand All @@ -140,7 +141,7 @@ public static StorageCredentialsToken storageCredentialAccessTokenFrom(String ac
}

private static ResultSegment<ListBlobItem> listBlobsInSegments(CloudBlobDirectory directory,
ResultContinuation token) throws IOException {
ResultContinuation token) throws IOException {
ResultSegment<ListBlobItem> result = null;
IOException lastException = null;
for (int sleep = 10; sleep <= 10000; sleep *= 10) { //increment the sleep time in steps.
Expand Down Expand Up @@ -171,6 +172,21 @@ private static ResultSegment<ListBlobItem> listBlobsInSegments(CloudBlobDirector
}
}

public static void deleteAllBlobs(@NotNull CloudBlobDirectory directory) throws URISyntaxException, StorageException, InterruptedException {
for (ListBlobItem blobItem : directory.listBlobs()) {
if (blobItem instanceof CloudBlob) {
CloudBlob cloudBlob = (CloudBlob) blobItem;
if (cloudBlob.getProperties().getLeaseStatus() == LeaseStatus.LOCKED) {
cloudBlob.breakLease(0);
}
cloudBlob.deleteIfExists();
} else if (blobItem instanceof CloudBlobDirectory) {
CloudBlobDirectory cloudBlobDirectory = (CloudBlobDirectory) blobItem;
deleteAllBlobs(cloudBlobDirectory);
}
}
}

private static class ByteBufferOutputStream extends OutputStream {

@NotNull
Expand All @@ -182,7 +198,7 @@ public ByteBufferOutputStream(@NotNull Buffer buffer) {

@Override
public void write(int b) {
buffer.put((byte)b);
buffer.put((byte) b);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* limitations under the License.
*/
@Internal(since = "1.0.0")
@Version("2.2.0")
@Version("2.3.0")
package org.apache.jackrabbit.oak.segment.azure;

import org.apache.jackrabbit.oak.commons.annotations.Internal;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -172,12 +172,12 @@ public static SegmentArchiveManager createArchiveManager(SegmentNodeStorePersist
public static CloudBlobDirectory createCloudBlobDirectory(String path) {
return createCloudBlobDirectory(path, ENVIRONMENT);
}

public static CloudBlobDirectory createCloudBlobDirectory(String path, Environment environment) {
Map<String, String> config = parseAzureConfigurationFromUri(path);

String accountName = config.get(KEY_ACCOUNT_NAME);

StorageCredentials credentials;
if (config.containsKey(KEY_SHARED_ACCESS_SIGNATURE)) {
credentials = new StorageCredentialsSharedAccessSignature(config.get(KEY_SHARED_ACCESS_SIGNATURE));
Expand All @@ -196,6 +196,11 @@ public static CloudBlobDirectory createCloudBlobDirectory(String path, Environme
}
}

@NotNull
public static StorageCredentials getStorageCredentialsFromAccountAndEnv(@NotNull String accountName) {
return getStorageCredentialsFromAccountAndEnv(accountName, ENVIRONMENT);
}

@NotNull
private static StorageCredentials getStorageCredentialsFromAccountAndEnv(String accountName, Environment environment) {
String clientId = environment.getVariable(AZURE_CLIENT_ID);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,17 +16,17 @@
*/
package org.apache.jackrabbit.oak.upgrade.cli.node;

import static org.apache.jackrabbit.oak.segment.SegmentCache.DEFAULT_SEGMENT_CACHE_MB;
import static org.apache.jackrabbit.oak.upgrade.cli.node.FileStoreUtils.asCloseable;

import java.io.File;
import java.io.IOException;
import java.net.URISyntaxException;
import java.security.InvalidKeyException;

import com.microsoft.azure.storage.StorageCredentials;
import com.microsoft.azure.storage.StorageCredentialsSharedAccessSignature;
import com.microsoft.azure.storage.StorageException;
import com.microsoft.azure.storage.blob.CloudBlobDirectory;
import org.apache.commons.lang3.StringUtils;
import org.apache.jackrabbit.guava.common.io.Closer;
import org.apache.jackrabbit.guava.common.io.Files;
import org.apache.jackrabbit.oak.segment.SegmentNodeStoreBuilders;
import org.apache.jackrabbit.oak.segment.azure.AzurePersistence;
import org.apache.jackrabbit.oak.segment.azure.AzureUtilities;
import org.apache.jackrabbit.oak.segment.azure.tool.ToolUtils;
import org.apache.jackrabbit.oak.segment.file.FileStore;
import org.apache.jackrabbit.oak.segment.file.FileStoreBuilder;
import org.apache.jackrabbit.oak.segment.file.InvalidFileStoreVersionException;
Expand All @@ -35,15 +35,17 @@
import org.apache.jackrabbit.oak.spi.state.NodeStore;
import org.apache.jackrabbit.oak.upgrade.cli.node.FileStoreUtils.NodeStoreWithFileStore;

import org.apache.jackrabbit.guava.common.io.Closer;
import org.apache.jackrabbit.guava.common.io.Files;
import com.microsoft.azure.storage.StorageCredentials;
import com.microsoft.azure.storage.StorageCredentialsAccountAndKey;
import com.microsoft.azure.storage.StorageException;
import com.microsoft.azure.storage.blob.CloudBlobDirectory;
import java.io.File;
import java.io.IOException;
import java.net.URISyntaxException;
import java.security.InvalidKeyException;

import static org.apache.jackrabbit.oak.segment.SegmentCache.DEFAULT_SEGMENT_CACHE_MB;
import static org.apache.jackrabbit.oak.upgrade.cli.node.FileStoreUtils.asCloseable;

public class SegmentAzureFactory implements NodeStoreFactory {
private final String accountName;
private final String sasToken;
private final String uri;
private final String connectionString;
private final String containerName;
Expand All @@ -58,6 +60,7 @@ public static class Builder {
private final boolean readOnly;

private String accountName;
private String sasToken;
private String uri;
private String connectionString;
private String containerName;
Expand All @@ -73,6 +76,11 @@ public Builder accountName(String accountName) {
return this;
}

public Builder sasToken(String sasToken) {
this.sasToken = sasToken;
return this;
}

public Builder uri(String uri) {
this.uri = uri;
return this;
Expand All @@ -95,6 +103,7 @@ public SegmentAzureFactory build() {

public SegmentAzureFactory(Builder builder) {
this.accountName = builder.accountName;
this.sasToken = builder.sasToken;
this.uri = builder.uri;
this.connectionString = builder.connectionString;
this.containerName = builder.containerName;
Expand Down Expand Up @@ -142,12 +151,17 @@ public NodeStore create(BlobStore blobStore, Closer closer) throws IOException {
private AzurePersistence createAzurePersistence() throws StorageException, URISyntaxException, InvalidKeyException {
CloudBlobDirectory cloudBlobDirectory = null;

if (accountName != null && uri != null) {
String key = System.getenv("AZURE_SECRET_KEY");
StorageCredentials credentials = new StorageCredentialsAccountAndKey(accountName, key);
cloudBlobDirectory = AzureUtilities.cloudBlobDirectoryFrom(credentials, uri, dir);
} else if (connectionString != null && containerName != null) {
// connection string will take precedence over accountkey / sas / service principal
if (StringUtils.isNoneBlank(connectionString, containerName)) {
cloudBlobDirectory = AzureUtilities.cloudBlobDirectoryFrom(connectionString, containerName, dir);
} else if (StringUtils.isNoneBlank(accountName, uri)) {
StorageCredentials credentials = null;
if (StringUtils.isNotBlank(sasToken)) {
credentials = new StorageCredentialsSharedAccessSignature(sasToken);
} else {
credentials = ToolUtils.getStorageCredentialsFromAccountAndEnv(accountName);
}
cloudBlobDirectory = AzureUtilities.cloudBlobDirectoryFrom(credentials, uri, dir);
}

if (cloudBlobDirectory == null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.KEY_CONNECTION_STRING;
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.KEY_CONTAINER_NAME;
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.KEY_DIR;
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.KEY_SHARED_ACCESS_SIGNATURE;
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.KEY_STORAGE_URI;
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.isCustomAzureConnectionString;
import static org.apache.jackrabbit.oak.segment.azure.util.AzureConfigurationParserUtils.parseAzureConfigurationFromCustomConnection;
Expand Down Expand Up @@ -175,6 +176,7 @@ public StoreFactory createFactory(String[] paths, MigrationDirection direction,
return new StoreFactory(new SegmentAzureFactory.Builder(config.get(KEY_DIR),
migrationOptions.getCacheSizeInMB(), direction == MigrationDirection.SRC)
.accountName(config.get(KEY_ACCOUNT_NAME))
.sasToken(config.get(KEY_SHARED_ACCESS_SIGNATURE))
.uri(config.get(KEY_STORAGE_URI))
.build()
);
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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.apache.jackrabbit.oak.upgrade.cli;

import org.apache.jackrabbit.oak.segment.azure.AzureUtilities;
import org.apache.jackrabbit.oak.segment.azure.util.Environment;
import org.apache.jackrabbit.oak.upgrade.cli.container.NodeStoreContainer;
import org.apache.jackrabbit.oak.upgrade.cli.container.SegmentAzureServicePrincipalNodeStoreContainer;
import org.apache.jackrabbit.oak.upgrade.cli.container.SegmentTarNodeStoreContainer;

import java.io.IOException;

import static org.junit.Assume.assumeNotNull;

public class SegmentTarToSegmentAzureServicePrincipalTest extends AbstractOak2OakTest {
private static boolean skipTest = true;
private static final Environment ENVIRONMENT = new Environment();
private final NodeStoreContainer source;
private final NodeStoreContainer destination;

@Override
public void prepare() throws Exception {
assumeNotNull(ENVIRONMENT.getVariable(AzureUtilities.AZURE_ACCOUNT_NAME), ENVIRONMENT.getVariable(AzureUtilities.AZURE_TENANT_ID),
ENVIRONMENT.getVariable(AzureUtilities.AZURE_CLIENT_ID), ENVIRONMENT.getVariable(AzureUtilities.AZURE_CLIENT_SECRET));
skipTest = false;
super.prepare();
}

@Override
public void clean() throws IOException {
if (!skipTest) {
super.clean();
}
}

public SegmentTarToSegmentAzureServicePrincipalTest() throws IOException {
source = new SegmentTarNodeStoreContainer();
destination = new SegmentAzureServicePrincipalNodeStoreContainer();
}

@Override
protected NodeStoreContainer getSourceContainer() {
return source;
}

@Override
protected NodeStoreContainer getDestinationContainer() {
return destination;
}

@Override
protected String[] getArgs() {
return new String[]{source.getDescription(), destination.getDescription()};
}
}
Loading

0 comments on commit 8f97e2a

Please sign in to comment.