public:

bool ShouldRetry(const AWSErrorDetail& error, int64_t attempted_retries) override {

if (error.message.find(kFileExistsMessage) != error.message.npos) {

// Minio returns "file exists" errors (when calling mkdir) as internal errors,

// which would trigger spurious retries.

return false;

}

return IsRetryable(error) && (attempted_retries * kRetryInterval < kMaxRetryDuration);

}

int64_t CalculateDelayBeforeNextRetry(const AWSErrorDetail& error,

int64_t attempted_retries) override {

return kRetryInterval;

}

bool IsRetryable(const AWSErrorDetail& error) const {

return error.should_retry || error.exception_name == "XMinioServerNotInitialized";

}

#ifdef _WIN32

static constexpr const char* kFileExistsMessage = "file already exists";

#else

static constexpr const char* kFileExistsMessage = "file exists";

#endif

public:

// We set this environment variable to speed up tests by ensuring

// DefaultAWSCredentialsProviderChain does not query (inaccessible)

// EC2 metadata endpoint

AwsTestMixin() : ec2_metadata_disabled_guard_("AWS_EC2_METADATA_DISABLED", "true") {}

void SetUp() override {

#ifdef AWS_CPP_SDK_S3_NOT_SHARED

auto aws_log_level = Aws::Utils::Logging::LogLevel::Fatal;

aws_options_.loggingOptions.logLevel = aws_log_level;

aws_options_.loggingOptions.logger_create_fn = [&aws_log_level] {

return std::make_shared<Aws::Utils::Logging::ConsoleLogSystem>(aws_log_level);

};

Aws::InitAPI(aws_options_);

#endif

}

void TearDown() override {

#ifdef AWS_CPP_SDK_S3_NOT_SHARED

Aws::ShutdownAPI(aws_options_);

#endif

}

private:

EnvVarGuard ec2_metadata_disabled_guard_;

#ifdef AWS_CPP_SDK_S3_NOT_SHARED

Aws::SDKOptions aws_options_;

#endif

public:

void SetUp() override {

AwsTestMixin::SetUp();

// Starting the server may fail, for example if the generated port number

// was "stolen" by another process. Run a dummy S3 operation to make sure it

// is running, otherwise retry a number of times.

Status connect_status;

int retries = kNumServerRetries;

do {

InitServerAndClient();

connect_status = OutcomeToStatus("ListBuckets", client_->ListBuckets());

} while (!connect_status.ok() && --retries > 0);

ASSERT_OK(connect_status);

}

void TearDown() override { AwsTestMixin::TearDown(); }

protected:

void InitServerAndClient() {

ASSERT_OK_AND_ASSIGN(minio_, GetMinioEnv()->GetOneServer());

client_config_.reset(new Aws::Client::ClientConfiguration());

client_config_->endpointOverride = ToAwsString(minio_->connect_string());

client_config_->scheme = Aws::Http::Scheme::HTTP;

client_config_->retryStrategy =

std::make_shared<ConnectRetryStrategy>(kRetryInterval, kMaxRetryDuration);

credentials_ = {ToAwsString(minio_->access_key()), ToAwsString(minio_->secret_key())};

bool use_virtual_addressing = false;

client_.reset(

new Aws::S3::S3Client(credentials_, *client_config_,

Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never,

use_virtual_addressing));

}

// How many times to try launching a server in a row before decreeing failure

static constexpr int kNumServerRetries = 3;

std::shared_ptr<MinioTestServer> minio_;

std::unique_ptr<Aws::Client::ClientConfiguration> client_config_;

Aws::Auth::AWSCredentials credentials_;

std::unique_ptr<Aws::S3::S3Client> client_;

const std::string& expected) {

auto length = static_cast<int64_t>(expected.length());

ASSERT_EQ(result.GetContentLength(), length);

auto& stream = result.GetBody();

std::string actual;

actual.resize(length + 1);

stream.read(&actual[0], length + 1);

ASSERT_EQ(stream.gcount(), length); // EOF was reached before length + 1

actual.resize(length);

ASSERT_EQ(actual.size(), expected.size());

ASSERT_TRUE(actual == expected); // Avoid ASSERT_EQ on large data

const std::string& key, const std::string& expected) {

Aws::S3::Model::GetObjectRequest req;

req.SetBucket(ToAwsString(bucket));

req.SetKey(ToAwsString(key));

ARROW_AWS_ASSIGN_OR_FAIL(auto result, client->GetObject(req));

AssertGetObject(result, expected);

std::string path;

S3Options options;

ASSERT_OK_AND_ASSIGN(options, S3Options::FromUri("s3://", &path));

ASSERT_EQ(options.region, "");

ASSERT_EQ(options.scheme, "https");

ASSERT_EQ(options.endpoint_override, "");

ASSERT_EQ(path, "");

ASSERT_OK_AND_ASSIGN(options, S3Options::FromUri("s3:", &path));

ASSERT_EQ(path, "");

ASSERT_OK_AND_ASSIGN(options, S3Options::FromUri("s3://access:secret@mybucket", &path));

ASSERT_EQ(path, "mybucket");

const auto creds = options.credentials_provider->GetAWSCredentials();

ASSERT_EQ(creds.GetAWSAccessKeyId(), "access");

ASSERT_EQ(creds.GetAWSSecretKey(), "secret");

ASSERT_OK_AND_ASSIGN(options, S3Options::FromUri("s3://mybucket/", &path));

ASSERT_NE(options.region, ""); // Some region was chosen

ASSERT_EQ(options.scheme, "https");

ASSERT_EQ(options.endpoint_override, "");

ASSERT_EQ(path, "mybucket");

ASSERT_OK_AND_ASSIGN(options, S3Options::FromUri("s3://mybucket/foo/bar/", &path));

ASSERT_NE(options.region, "");

ASSERT_EQ(options.scheme, "https");

ASSERT_EQ(options.endpoint_override, "");

ASSERT_EQ(path, "mybucket/foo/bar");

// Region resolution with a well-known bucket

ASSERT_OK_AND_ASSIGN(

options, S3Options::FromUri("s3://aws-earth-mo-atmospheric-ukv-prd/", &path));

ASSERT_EQ(options.region, "eu-west-2");

// Explicit region override

ASSERT_OK_AND_ASSIGN(

options,

S3Options::FromUri(

"s3://mybucket/foo/bar/?region=utopia&endpoint_override=localhost&scheme=http",

&path));

ASSERT_EQ(options.region, "utopia");

ASSERT_EQ(options.scheme, "http");

ASSERT_EQ(options.endpoint_override, "localhost");

ASSERT_EQ(path, "mybucket/foo/bar");

// Missing bucket name

ASSERT_RAISES(Invalid, S3Options::FromUri("s3:///foo/bar/", &path));

// Invalid option

ASSERT_RAISES(Invalid, S3Options::FromUri("s3://mybucket/?xxx=zzz", &path));

S3Options options;

// session token is optional and should default to empty string

options = S3Options::FromAccessKey("access", "secret");

ASSERT_EQ(options.GetAccessKey(), "access");

ASSERT_EQ(options.GetSecretKey(), "secret");

ASSERT_EQ(options.GetSessionToken(), "");

options = S3Options::FromAccessKey("access", "secret", "token");

ASSERT_EQ(options.GetAccessKey(), "access");

ASSERT_EQ(options.GetSecretKey(), "secret");

ASSERT_EQ(options.GetSessionToken(), "token");

S3Options options;

// arn should be only required argument

options = S3Options::FromAssumeRole("my_role_arn");

options = S3Options::FromAssumeRole("my_role_arn", "session");

options = S3Options::FromAssumeRole("my_role_arn", "session", "id");

options = S3Options::FromAssumeRole("my_role_arn", "session", "id", 42);

// test w/ custom STSClient (will not use DefaultAWSCredentialsProviderChain)

Aws::Auth::AWSCredentials test_creds = Aws::Auth::AWSCredentials("access", "secret");

std::shared_ptr<Aws::STS::STSClient> sts_client =

std::make_shared<Aws::STS::STSClient>(Aws::STS::STSClient(test_creds));

options = S3Options::FromAssumeRole("my_role_arn", "session", "id", 42, sts_client);

ASSERT_OK_AND_EQ("us-east-2", ResolveS3BucketRegion("voltrondata-labs-datasets"));

// Taken from a registry of open S3-hosted datasets

// at https://github.com/awslabs/open-data-registry

ASSERT_OK_AND_EQ("eu-west-2",

ResolveS3BucketRegion("aws-earth-mo-atmospheric-ukv-prd"));

// Same again, cached

ASSERT_OK_AND_EQ("eu-west-2",

ResolveS3BucketRegion("aws-earth-mo-atmospheric-ukv-prd"));

ASSERT_OK_AND_EQ("us-west-2", ResolveS3BucketRegion("ursa-labs-r-test"));

// Same again, cached

ASSERT_OK_AND_EQ("us-west-2", ResolveS3BucketRegion("ursa-labs-r-test"));

auto maybe_region = ResolveS3BucketRegion("ursa-labs-non-existent-bucket");

ASSERT_RAISES(IOError, maybe_region);

ASSERT_THAT(maybe_region.status().message(),

::testing::HasSubstr("Bucket 'ursa-labs-non-existent-bucket' not found"));

ASSERT_OK_AND_ASSIGN(auto fs, FileSystemFromUri("s3://"));

auto s3fs = checked_pointer_cast<S3FileSystem>(fs);

ASSERT_EQ(s3fs->region(), "us-east-1");

public:

void SetUp() override { S3TestMixin::SetUp(); }

protected:

public:

void SetUp() override {

S3TestMixin::SetUp();

// Most tests will create buckets

options_.allow_bucket_creation = true;

options_.allow_bucket_deletion = true;

MakeFileSystem();

// Set up test bucket

{

Aws::S3::Model::CreateBucketRequest req;

req.SetBucket(ToAwsString("bucket"));

ASSERT_OK(OutcomeToStatus("CreateBucket", client_->CreateBucket(req)));

req.SetBucket(ToAwsString("empty-bucket"));

ASSERT_OK(OutcomeToStatus("CreateBucket", client_->CreateBucket(req)));

}

{

Aws::S3::Model::PutObjectRequest req;

req.SetBucket(ToAwsString("bucket"));

req.SetKey(ToAwsString("emptydir/"));

req.SetBody(std::make_shared<std::stringstream>(""));

ASSERT_OK(OutcomeToStatus("PutObject", client_->PutObject(req)));

// NOTE: no need to create intermediate "directories" somedir/ and

// somedir/subdir/

req.SetKey(ToAwsString("somedir/subdir/subfile"));

req.SetBody(std::make_shared<std::stringstream>("sub data"));

ASSERT_OK(OutcomeToStatus("PutObject", client_->PutObject(req)));

req.SetKey(ToAwsString("somefile"));

req.SetBody(std::make_shared<std::stringstream>("some data"));

req.SetContentType("x-arrow/test");

ASSERT_OK(OutcomeToStatus("PutObject", client_->PutObject(req)));

req.SetKey(ToAwsString("otherdir/1/2/3/otherfile"));

req.SetBody(std::make_shared<std::stringstream>("other data"));

ASSERT_OK(OutcomeToStatus("PutObject", client_->PutObject(req)));

}

}

void MakeFileSystem() {

options_.ConfigureAccessKey(minio_->access_key(), minio_->secret_key());

options_.scheme = "http";

options_.endpoint_override = minio_->connect_string();

if (!options_.retry_strategy) {

options_.retry_strategy = std::make_shared<ShortRetryStrategy>();

}

ASSERT_OK_AND_ASSIGN(fs_, S3FileSystem::Make(options_));

}

template <typename Matcher>

void AssertMetadataRoundtrip(const std::string& path,

const std::shared_ptr<const KeyValueMetadata>& metadata,

Matcher&& matcher) {

ASSERT_OK_AND_ASSIGN(auto output, fs_->OpenOutputStream(path, metadata));

ASSERT_OK(output->Close());

ASSERT_OK_AND_ASSIGN(auto input, fs_->OpenInputStream(path));

ASSERT_OK_AND_ASSIGN(auto got_metadata, input->ReadMetadata());

ASSERT_NE(got_metadata, nullptr);

ASSERT_THAT(got_metadata->sorted_pairs(), matcher);

}

void AssertInfoAllBucketsRecursive(const std::vector<FileInfo>& infos) {

AssertFileInfo(infos[0], "bucket", FileType::Directory);

AssertFileInfo(infos[1], "bucket/emptydir", FileType::Directory);

AssertFileInfo(infos[2], "bucket/otherdir", FileType::Directory);

AssertFileInfo(infos[3], "bucket/otherdir/1", FileType::Directory);

AssertFileInfo(infos[4], "bucket/otherdir/1/2", FileType::Directory);

AssertFileInfo(infos[5], "bucket/otherdir/1/2/3", FileType::Directory);

AssertFileInfo(infos[6], "bucket/otherdir/1/2/3/otherfile", FileType::File, 10);

AssertFileInfo(infos[7], "bucket/somedir", FileType::Directory);

AssertFileInfo(infos[8], "bucket/somedir/subdir", FileType::Directory);

AssertFileInfo(infos[9], "bucket/somedir/subdir/subfile", FileType::File, 8);

AssertFileInfo(infos[10], "bucket/somefile", FileType::File, 9);

AssertFileInfo(infos[11], "empty-bucket", FileType::Directory);

}

void TestOpenOutputStream() {

std::shared_ptr<io::OutputStream> stream;

// Nonexistent

ASSERT_RAISES(IOError, fs_->OpenOutputStream("nonexistent-bucket/somefile"));

// URI

ASSERT_RAISES(Invalid, fs_->OpenOutputStream("s3:bucket/newfile1"));

// Create new empty file

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile1"));

ASSERT_OK(stream->Close());

AssertObjectContents(client_.get(), "bucket", "newfile1", "");

// Create new file with 1 small write

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile2"));

ASSERT_OK(stream->Write("some data"));

ASSERT_OK(stream->Close());

AssertObjectContents(client_.get(), "bucket", "newfile2", "some data");

// Create new file with 3 small writes

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile3"));

ASSERT_OK(stream->Write("some "));

ASSERT_OK(stream->Write(""));

ASSERT_OK(stream->Write("new data"));

ASSERT_OK(stream->Close());

AssertObjectContents(client_.get(), "bucket", "newfile3", "some new data");

// Create new file with some large writes

std::string s1, s2, s3, s4, s5, expected;

s1 = random_string(6000000, /*seed =*/42); // More than the 5 MB minimum part upload

s2 = "xxx";

s3 = random_string(6000000, 43);

s4 = "zzz";

s5 = random_string(600000, 44);

expected = s1 + s2 + s3 + s4 + s5;

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile4"));

for (auto input : {s1, s2, s3, s4, s5}) {

ASSERT_OK(stream->Write(input));

// Clobber source contents. This shouldn't reflect in the data written.

input.front() = 'x';

input.back() = 'x';

}

ASSERT_OK(stream->Close());

AssertObjectContents(client_.get(), "bucket", "newfile4", expected);

// Overwrite

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile1"));

ASSERT_OK(stream->Write("overwritten data"));

ASSERT_OK(stream->Close());

AssertObjectContents(client_.get(), "bucket", "newfile1", "overwritten data");

// Overwrite and make empty

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile1"));

ASSERT_OK(stream->Close());

AssertObjectContents(client_.get(), "bucket", "newfile1", "");

// Open file and then lose filesystem reference

ASSERT_EQ(fs_.use_count(), 1); // needed for test to work

std::weak_ptr<S3FileSystem> weak_fs(fs_);

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/newfile99"));

fs_.reset();

ASSERT_OK(stream->Write("some other data"));

ASSERT_OK(stream->Close());

ASSERT_TRUE(weak_fs.expired());

AssertObjectContents(client_.get(), "bucket", "newfile99", "some other data");

}

void TestOpenOutputStreamAbort() {

std::shared_ptr<io::OutputStream> stream;

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/somefile"));

ASSERT_OK(stream->Write("new data"));

// Abort() cancels the multipart upload.

ASSERT_OK(stream->Abort());

ASSERT_EQ(stream->closed(), true);

AssertObjectContents(client_.get(), "bucket", "somefile", "some data");

}

void TestOpenOutputStreamDestructor() {

std::shared_ptr<io::OutputStream> stream;

ASSERT_OK_AND_ASSIGN(stream, fs_->OpenOutputStream("bucket/somefile"));

ASSERT_OK(stream->Write("new data"));

// Destructor implicitly closes stream and completes the multipart upload.

stream.reset();

AssertObjectContents(client_.get(), "bucket", "somefile", "new data");

}

protected:

S3Options options_;

std::shared_ptr<S3FileSystem> fs_;