Source code for operators.http_to_s3

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from io import BytesIO




class HttpToS3Operator(BaseOperator):
    """
    Operator to upload a file to s3 from an http link
    :param url: path to HTTP file.
    :type url: str
    :param bucket: name of the s3 bucket
    :type bucket: str
    :param key: path to the s3 file within bucket
    :type key: str
    :param if_exists: action to take if file already in s3. 'noop', 'overwrite'
    :type if_exists: str
    """

    ui_color = "#0380fc"

    @apply_defaults
    def __init__(
        self, url, bucket, key, if_exists, headers=None, verify=None, *args, **kwargs
    ):
        super(HttpToS3Operator, self).__init__(*args, **kwargs)
        self.url = url
        self.bucket = bucket
        self.key = key
        self.if_exists = if_exists
        self.headers = headers
        self.verify = verify

    def _http_to_s3(self):
        import boto3
        import requests

        r = requests.get(
            self.url,
            headers=self.headers,
            verify=self.verify if self.verify is not None else True,
        )
        content = BytesIO(r.content)
        client = boto3.client("s3")
        client.put_object(Body=content, Bucket=self.bucket, Key=self.key)



    def execute(self, context):
        import boto3

        if self.if_exists == "noop":
            s3 = boto3.resource("s3")
            s3_bucket = s3.Bucket(self.bucket)
            objs = list(s3_bucket.objects.filter(Prefix=self.key))
            if not (len(objs) > 0 and objs[0].key == self.key):  # if file not in s3
                self._http_to_s3()
        else:
            self._http_to_s3()