adding delete and put interfaces for the S3 plugin
This commit is contained in:
parent
cd29b2b870
commit
bfe89e131e
|
@ -9,9 +9,12 @@
|
|||
from flask import current_app
|
||||
from .sts import sts_client
|
||||
|
||||
from botocore.exceptions import ClientError
|
||||
from lemur.extensions import sentry
|
||||
|
||||
|
||||
@sts_client("s3", service_type="resource")
|
||||
def put(bucket_name, region, prefix, data, encrypt, **kwargs):
|
||||
def put(bucket_name, prefix, data, encrypt, **kwargs):
|
||||
"""
|
||||
Use STS to write to an S3 bucket
|
||||
"""
|
||||
|
@ -32,4 +35,37 @@ def put(bucket_name, region, prefix, data, encrypt, **kwargs):
|
|||
ServerSideEncryption="AES256",
|
||||
)
|
||||
else:
|
||||
bucket.put_object(Key=prefix, Body=data, ACL="bucket-owner-full-control")
|
||||
try:
|
||||
bucket.put_object(Key=prefix, Body=data, ACL="bucket-owner-full-control")
|
||||
except ClientError:
|
||||
sentry.captureException()
|
||||
|
||||
|
||||
@sts_client("s3", service_type="client")
|
||||
def delete(bucket_name, prefix, **kwargs):
|
||||
"""
|
||||
Use STS to delete an object
|
||||
"""
|
||||
try:
|
||||
response = kwargs["client"].delete_object(Bucket=bucket_name, Key=prefix)
|
||||
current_app.logger.debug(f"Delete data from S3."
|
||||
f"Bucket: {bucket_name},"
|
||||
f"Prefix: {prefix},"
|
||||
f"Status_code: {response}")
|
||||
return response['ResponseMetadata']['HTTPStatusCode'] < 300
|
||||
except ClientError:
|
||||
sentry.captureException()
|
||||
|
||||
|
||||
@sts_client("s3", service_type="client")
|
||||
def get(bucket_name, prefix, **kwargs):
|
||||
"""
|
||||
Use STS to get an object
|
||||
"""
|
||||
try:
|
||||
response = kwargs["client"].get_object(Bucket=bucket_name, Key=prefix)
|
||||
current_app.logger.debug(f"Get data from S3. Bucket: {bucket_name},"
|
||||
f"Prefix: {prefix}")
|
||||
return response['Body'].read().decode("utf-8")
|
||||
except ClientError:
|
||||
sentry.captureException()
|
||||
|
|
|
@ -0,0 +1,38 @@
|
|||
import boto3
|
||||
from moto import mock_sts, mock_s3
|
||||
|
||||
|
||||
@mock_sts()
|
||||
@mock_s3()
|
||||
def test_put_delete_s3_object(app):
|
||||
from lemur.plugins.lemur_aws.s3 import put, delete, get
|
||||
|
||||
bucket = "public-bucket"
|
||||
account = "123456789012"
|
||||
path = "some_path/foo"
|
||||
|
||||
s3_client = boto3.client('s3')
|
||||
s3_client.create_bucket(Bucket=bucket)
|
||||
|
||||
data = "dummy data"
|
||||
put(bucket_name=bucket,
|
||||
prefix=path,
|
||||
data=data,
|
||||
encrypt=None,
|
||||
account_number=account)
|
||||
|
||||
response = get(bucket_name=bucket, prefix=path, account_number=account)
|
||||
|
||||
# put data, and getting the same data
|
||||
assert (response == data)
|
||||
|
||||
response = get(bucket_name="wrong-bucket", prefix=path, account_number=account)
|
||||
|
||||
# attempting to get thccle wrong data
|
||||
assert (response is None)
|
||||
|
||||
delete(bucket_name=bucket, prefix=path, account_number=account)
|
||||
response = get(bucket_name=bucket, prefix=path, account_number=account)
|
||||
|
||||
# delete data, and getting the same data
|
||||
assert (response is None)
|
Loading…
Reference in New Issue