Function test_write_dataset_s3_put_only

python/pyarrow/tests/test_dataset.py:5154–5227

(s3_server)

@pytest.mark.parquet
@pytest.mark.s3
def test_write_dataset_s3_put_only(s3_server):
    # [ARROW-15892] Testing the create_dir flag which will restrict
    # creating a new directory for writing a dataset. This is
    # required while writing a dataset in s3 where we have very
    # limited permissions and thus we can directly write the dataset
    # without creating a directory.
    from pyarrow.fs import S3FileSystem

    # write dataset with s3 filesystem
    host, port, _, _ = s3_server['connection']

    _configure_s3_limited_user(s3_server, _minio_put_only_policy,
                               'test_dataset_limited_user', 'limited123')
    fs = S3FileSystem(
        access_key='test_dataset_limited_user',
        secret_key='limited123',
        endpoint_override=f'{host}:{port}',
        scheme='http'
    )

    table = pa.table([
        pa.array(range(20)), pa.array(random.random() for _ in range(20)),
        pa.array(['a']*10 + ['b'] * 10)],
        names=["f1", "f2", "part"]
    )
    part = ds.partitioning(pa.schema([("part", pa.string())]), flavor="hive")

    # writing with filesystem object with create_dir flag set to false
    ds.write_dataset(
        table, "existing-bucket", filesystem=fs,
        format="feather", create_dir=False, partitioning=part,
        existing_data_behavior='overwrite_or_ignore'
    )
    # check roundtrip
    result = ds.dataset(
        "existing-bucket", filesystem=fs, format="ipc", partitioning="hive"
    ).to_table()
    assert result.equals(table)

    # Passing create_dir is fine if the bucket already exists
    ds.write_dataset(
        table, "existing-bucket", filesystem=fs,
        format="feather", create_dir=True, partitioning=part,
        existing_data_behavior='overwrite_or_ignore'
    )
    # check roundtrip
    result = ds.dataset(
        "existing-bucket", filesystem=fs, format="ipc", partitioning="hive"
    ).to_table()
    assert result.equals(table)

    # Error enforced by filesystem
    with pytest.raises(OSError,
                       match="Bucket 'non-existing-bucket' not found"):
        ds.write_dataset(
            table, "non-existing-bucket", filesystem=fs,
            format="feather", create_dir=True,
            existing_data_behavior='overwrite_or_ignore'