Source code for operators.github_to_s3

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
import os
import shutil
import subprocess
from tempfile import mkdtemp


DEFAULT_S3_BUCKET = Config.DATA_LAKE_BUCKET




class GithubToS3Operator(BaseOperator):
    """
    this is a sin
    """

    ui_color = "#4287f5"

    @apply_defaults
    def __init__(
        self,
        repo_url: str,
        s3_key: str,
        s3_bucket_name: str = DEFAULT_S3_BUCKET,
        *args,
        **kwargs,
    ):
        super(GithubToS3Operator, self).__init__(*args, **kwargs)
        self.repo_url = repo_url
        self.s3_bucket_name = s3_bucket_name
        self.s3_key = s3_key



    def execute(self, context):
        """
        Clones a git repository to a temporary directory, uploads it to S3,
        then cleans up the temporary files.
        """
        import boto3

        temp_dir = mkdtemp(dir="/tmp")
        repo_path = os.path.join(temp_dir, self.repo_url.split("/")[-1])
        git_url = f"git@github.com:{self.repo_url}.git"

        try:
            # Clone the repository
            subprocess.run(
                ["git", "clone", git_url, repo_path],
                check=True,
                capture_output=True,
            )

            # Initialize S3 client
            s3_client = boto3.client("s3")

            # Upload all files in the cloned repo
            for root, _, files in os.walk(repo_path):
                for each_file in files:
                    file_path = os.path.join(root, each_file)
                    # Calculate relative path for S3 key
                    relative_path = os.path.relpath(file_path, repo_path)
                    s3_key = f"{self.s3_key.rstrip('/')}/{relative_path}"

                    s3_client.upload_file(file_path, self.s3_bucket_name, s3_key)

        except subprocess.CalledProcessError as e:
            raise Exception(f"Git clone failed: {e.stderr.decode()}")
        except Exception as e:
            raise Exception(f"Error during process: {str(e)}")
        finally:
            # Clean up temp directory
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)