Dp 2534 aws s3

bin/init-s3.sh

@@ -1,15 +1,39 @@
 #!/bin/bash
 
-if [ -n "$BUCKET_NAME" ]; then
-    echo "Creating bucket: s3://$BUCKET_NAME"
-    awslocal s3 mb s3://"$BUCKET_NAME"
+echo "=================================================="
+echo "    STARTING S3 SEEDING PROCESS           "
+echo "=================================================="
 
-    if [ -n "$AWS_SEED_DATA" ]; then
-        echo "AWS_SEED_DATA is set ($AWS_SEED_DATA). Syncing files..."
-        awslocal s3 sync /tmp/aws_seed_data/ s3://"$BUCKET_NAME"/
+if [ -n "$AWS_CITY_ARTS_BUCKET" ]; then
+    echo "Creating City Arts bucket: s3://$AWS_CITY_ARTS_BUCKET"
+    awslocal s3 mb s3://"$AWS_CITY_ARTS_BUCKET"
+
+    if [ -d "/tmp/aws_seed_city_data" ]; then
+        echo "Syncing City Arts data from /tmp/aws_seed_city_data..."
+        awslocal s3 sync /tmp/aws_seed_city_data/ s3://"$AWS_CITY_ARTS_BUCKET"/
     else
-        echo "AWS_SEED_DATA is not set. Skipping file seeding."
+        echo "Warning: /tmp/aws_seed_city_data directory not found. Skipping City Arts seeding."
     fi
 else
-    echo "No BUCKET_NAME environment variable specified. Skipping initialization."
+    echo "AWS_CITY_ARTS_BUCKET is not specified. Skipping City Arts initialization."
 fi
+
+echo "--------------------------------------------------"
+
+if [ -n "$AWS_TEST_BUCKET" ]; then
+    echo "Creating Test bucket: s3://$AWS_TEST_BUCKET"
+    awslocal s3 mb s3://"$AWS_TEST_BUCKET"
+
+    if [ -d "/tmp/aws_seed_test_data" ]; then
+        echo "Syncing Test data from /tmp/aws_seed_test_data..."
+        awslocal s3 sync /tmp/aws_seed_test_data/ s3://"$AWS_TEST_BUCKET"/
+    else
+        echo "Warning: /tmp/aws_seed_test_data directory not found. Skipping Test seeding."
+    fi
+else
+    echo "AWS_TEST_BUCKET is not specified. Skipping Test initialization."
+fi
+
+echo "=================================================="
+echo "          SEEDING PROCESS COMPLETE                "
+echo "=================================================="

docker-compose.yml

@@ -65,6 +65,8 @@ x-airflow-common: &airflow-common
     AIRFLOW__LOGGING__LOGGING_LEVEL: DEBUG
     AIRFLOW__SCHEDULER__ENABLE_HEALTH_CHECK: 'true'
     AIRFLOW_CONN_SMTP_DEFAULT: 'smtp://mailpit:1025?disable_ssl=true&disable_tls=true'
+    AIRFLOW_CONN_AWS_S3_CITY_ARTS: ${AIRFLOW_CONN_AWS_S3_CITY_ARTS:-} 
+    AIRFLOW_CONN_AWS_S3_TEST: ${AIRFLOW_CONN_AWS_S3_TEST:-} 
     CI: ${CI:-}
   volumes:
     - ./:/opt/airflow
@@ -307,15 +309,21 @@ services:
       SERVICES: s3
       BUCKET_NAME: "${AWS_BUCKET_NAME:-city-arts}"
       AWS_SEED_DATA: "${AWS_SEED_DATA:-/tmp/seed_data}"
-      AWS_ACCESS_KEY_ID: "${AWS_ACCESS_KEY_ID:-test}"
-      AWS_SECRET_ACCESS_KEY: "${AWS_SECRET_ACCESS_KEY:-test}"
-      AWS_DEFAULT_REGION: "${AWS_DEFAULT_REGION:-us-west-2}"
-      AWS_ENDPOINT_URL: "${AWS_ENDPOINT_URL:-http://localstack:4566}"
+      AWS_CITY_ARTS_BUCKET: "${AWS_CITY_ARTS_BUCKET:-cityartsmedia}"
+      AWS_TEST_BUCKET: "${AWS_TEST_BUCKET:-testbucket}"
+
       DEBUG: ${DEBUG:-0}
     volumes:
       - "./bin/init-s3.sh:/etc/localstack/init/ready.d/init-s3.sh"
-      - "${AWS_SEED_DATA:-./tmp/seed_data}:/tmp/aws_seed_data"
+      - "${AWS_SEED_CITY_DATA:-./tmp/aws_seed_city_data}:/tmp/aws_seed_city_data"
+      - "${AWS_SEED_TEST_DATA:-./tmp/aws_seed_test_data}:/tmp/aws_seed_test_data"
       - "./.data/localstack:/var/lib/localstack"
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:4566/_localstack/health"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+      start_period: 15s
 
 volumes:
   postgres-db-volume:

mokelumne/dags/fetch_from_s3.py

@@ -0,0 +1,102 @@
+import io
+import os
+import logging
+from pathlib import Path
+from datetime import datetime
+import pandas as pd
+
+from airflow.sdk import Param, dag, task
+from airflow.exceptions import AirflowFailException
+from mokelumne.util.s3_utils import list_bucket_files, download_single_s3_file
+
+logger = logging.getLogger(__name__)
+
+@dag(
+    dag_id="s3_download",
+    description="Download files from AWS S3 bucket to srv directory",
+    schedule=None,
+    catchup=False,
+    params={
+        "s3_bucket": Param(
+            default="cityartsmedia", 
+            type="string", 
+            description="The name of your S3 bucket."
+        ),
+        "s3_conn": Param(
+            default="AWS_S3_CITY_ARTS",
+            type="string",
+            description="The S3 connection."
+        ),
+        "destination_directory": Param(
+            default="/srv/pa/city_arts/incoming", 
+            type="string", 
+            description="The absolute directory path where files will be saved."
+        ),
+        "file_prefix": Param(
+            default="", 
+            type=["null", "string"], 
+            description="File prefix. Default is any prefix"
+        ),
+        "file_extension": Param(
+            default="", 
+            type=["null", "string"], 
+            description="File extension. Default is any extension"
+        ),
+    },
+)
+def retrieve_S3_bucket():
+
+    @task
+    def validate_destination(params: dict):
+        """Checks that the destination path exists."""
+        dest_path = params["destination_directory"]
+        destination_path = Path(dest_path)
+        if not destination_path.exists():
+            raise AirflowFailException(
+                f"Destination directory does not exist: {destination_path}"
+            )
+
+    @task
+    def get_bucket_file_names(params: dict) -> list:
+        """Get a list of files for a given bucket. Can be filtered by prefix and extension"""
+        bucket = params["s3_bucket"]
+        s3_conn = params["s3_conn"]
+        file_prefix = params["file_prefix"]
+        file_extension = params["file_extension"]
+
+        prefix = file_prefix.strip() if file_prefix else None
+        extension = file_extension.strip() if file_extension else None
+
+        file_names = list_bucket_files(bucket_name=bucket, conn_id = s3_conn, file_prefix=prefix, file_extension=extension)
+        if file_names:
+            for name in file_names:
+                logger.info("Found file: %s", name)
+        else:
+            logger.info("The bucket is empty!")
+
+        return file_names
+
+    @task(max_active_tis_per_dag=4)
+    def retrieve_file_from_bucket(file_key: str, params: dict):
+        """Download from S3"""
+        bucket = params["s3_bucket"]
+        s3_conn = params["s3_conn"]
+        dest_dir = params["destination_directory"]
+
+        try:
+            download_single_s3_file(
+                file_key=file_key,
+                conn_id=s3_conn,
+                bucket_name=bucket,
+                dest_dir=dest_dir
+            )
+        except Exception as ex:
+            raise AirflowFailException(f"Task aborted. Details: {str(ex)}")
+
+    validated_destination = validate_destination()
+    bucket_filenames = get_bucket_file_names()
+    retrieved_files = retrieve_file_from_bucket.expand(file_key=bucket_filenames)
+
+    validated_destination >> bucket_filenames >> retrieved_files
+
+retrieve_S3_bucket()

mokelumne/util/s3_utils.py

@@ -0,0 +1,62 @@
+import logging
+import os
+
+from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.sdk.exceptions import AirflowFailException
+
+logger = logging.getLogger(__name__)
+
+def list_bucket_files(bucket_name: str, conn_id: str, file_prefix: str | None = None,
+                      file_extension: str | None = None) -> list:
+    """Get a list of files for a given bucket"""
+
+    bucket_name = bucket_name.strip()
+
+    hook = S3Hook(aws_conn_id=conn_id)
+
+    if not hook.check_for_bucket(bucket_name=bucket_name):
+        raise AirflowFailException(f"Bucket {bucket_name} was not found.")
+
+    prefix = file_prefix.strip() if file_prefix else ""
+    extension = file_extension.strip() if file_extension else ""
+
+    keys = hook.list_keys(bucket_name=bucket_name, prefix=prefix)
+
+    if keys is None:
+        return []
+
+    return [key for key in keys if key.lower().endswith(extension.lower())]
+
+
+def download_single_s3_file(file_key: str, bucket_name: str, conn_id: str, dest_dir: str) -> str:
+    """
+    Downloads a single file from S3 to a local directory.
+    Returns the absolute path to the downloaded local file.
+    """
+
+    timeout_settings = {
+        "connect_timeout": 60,
+        "read_timeout": 300
+    }
+
+    hook = S3Hook(aws_conn_id=conn_id, config=timeout_settings)
+    logger.info("Going to download file: %s", file_key)
+
+    local_file_name = os.path.basename(file_key)
+    local_output_path = os.path.join(dest_dir, local_file_name)
+
+    try:
+        hook.download_file(
+            key=file_key,
+            bucket_name=bucket_name,
+            local_path=dest_dir,
+            preserve_file_name=True,
+            use_autogenerated_subdir=False  
+        )
+        logger.warning("Successfully saved file to: %s", local_output_path)
+        return local_output_path
+
+    except Exception as ex:
+        if os.path.exists(local_output_path):
+            logger.error("File exists. Will not download: %s", local_output_path)
+        raise RuntimeError(f"Failed to download {file_key}: {str(ex)}") from ex

test/dags/test_fetch_from_s3.py

@@ -0,0 +1,134 @@
+"""Test the s3_download DAG."""
+import pytest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+from airflow.exceptions import AirflowFailException
+from test.util.dag_helper import get_dag
+
+with patch("mokelumne.util.s3_utils.list_bucket_files"), \
+     patch("mokelumne.util.s3_utils.download_single_s3_file"):
+    DAG = get_dag("s3_download")
+
+
+class TestS3DownloadDAGStructure:
+    """Test basic DAG structure, configurations, and upstream dependencies."""
+
+    def test_dag_loads(self):
+        """Test that the DAG loads without errors."""
+        assert DAG is not None
+        assert DAG.dag_id == "s3_download"
+
+    def test_all_required_tasks_exist(self):
+        """Test that all required tasks are present in the DAG."""
+        required_tasks = [
+            "validate_destination",
+            "get_bucket_file_names",
+            "retrieve_file_from_bucket",
+        ]
+        task_ids = [t.task_id for t in DAG.tasks]
+        for task_id in required_tasks:
+            assert task_id in task_ids, f"Task '{task_id}' not found in DAG"
+
+    def test_dag_task_order(self):
+        """Test that tasks execute in the correct linear order."""
+        validate_dest = DAG.get_task("validate_destination")
+        get_filenames = DAG.get_task("get_bucket_file_names")
+        retrieve_file = DAG.get_task("retrieve_file_from_bucket")
+
+        # Verify: validate_destination >> get_bucket_file_names >> retrieve_file_from_bucket
+        assert validate_dest in get_filenames.upstream_list
+        assert get_filenames in retrieve_file.upstream_list
+
+    def test_dag_parameters_exist(self):
+        """Verify that default DAG parameters are defined correctly."""
+        dag_params = DAG.params
+        assert "s3_bucket" in dag_params
+        assert "destination_directory" in dag_params
+        assert "file_prefix" in dag_params
+        assert "file_extension" in dag_params
+        assert "s3_conn" in dag_params
+
+
+class TestValidateDestinationTask:
+    """Test validate_destination task logic"""
+
+    def test_validate_destination_success(self, tmp_path):
+        """Should pass silently if the destination directory exists."""
+        validate_task = DAG.get_task("validate_destination").python_callable
+        params = {"destination_directory": str(tmp_path)}
+        validate_task(params=params)
+
+    def test_validate_destination_raises_error(self):
+        """Should throw AirflowFailException if the directory is missing."""
+        validate_task = DAG.get_task("validate_destination").python_callable
+        params = {"destination_directory": "/nonexistent/absolute/path/to/folder"}
+        with pytest.raises(AirflowFailException) as exc_info:
+            validate_task(params=params)
+        assert "Destination directory does not exist:" in str(exc_info.value)
+
+
+class TestGetBucketFileNamesTask:
+    """Test get_bucket_file_names task logic and utility filtering calls."""
+
+    def test_get_bucket_file_names_success(self):
+        """Should clean and pass configurations to list_bucket_files utility."""
+        get_filenames_task = DAG.get_task("get_bucket_file_names").python_callable
+        params = {
+            "s3_bucket": "my-bucket",
+            "file_prefix": " raw_data/ ",
+            "file_extension": " .wav",
+            "s3_conn": "AWS_S3_CITY_ARTS"
+        }
+        mock_list_files = MagicMock(return_value=["file1.wav", "file2.wav"])
+
+        with patch.dict(get_filenames_task.__globals__, {"list_bucket_files": mock_list_files}):
+            result = get_filenames_task(params=params)
+
+        assert result == ["file1.wav", "file2.wav"]
+        mock_list_files.assert_called_once_with(
+            bucket_name="my-bucket",
+            file_prefix="raw_data/",
+            file_extension=".wav",
+            conn_id="AWS_S3_CITY_ARTS"
+        )
+
+
+class TestRetrieveFileFromBucketTask:
+    """Test retrieve_file_from_bucket task and download."""
+
+    def test_retrieve_file_from_bucket_success(self):
+        """Should test download utility with mappings."""
+        retrieve_task = DAG.get_task("retrieve_file_from_bucket").python_callable
+        params = {
+            "s3_bucket": "my-bucket",
+            "destination_directory": "/srv/pa/incoming",
+            "s3_conn": "AWS_S3_CITY_ARTS"
+        }
+        mock_download_file = MagicMock()
+
+        with patch.dict(retrieve_task.__globals__, {"download_single_s3_file": mock_download_file}):
+            retrieve_task(file_key="data/audio_file.wav", params=params)
+
+        mock_download_file.assert_called_once_with(
+            file_key="data/audio_file.wav",
+            bucket_name="my-bucket",
+            dest_dir="/srv/pa/incoming",
+            conn_id="AWS_S3_CITY_ARTS"
+        )
+
+    def test_retrieve_file_from_bucket_failure(self):
+        """Should catch utility exceptions and re-raise them as AirflowFailException."""
+        retrieve_task = DAG.get_task("retrieve_file_from_bucket").python_callable
+        params = {
+            "s3_bucket": "my-bucket",
+            "destination_directory": "/srv/pa/incoming",
+            "s3_conn": "AWS_S3_CITY_ARTS"
+        }
+        mock_download_file = MagicMock(side_effect=RuntimeError("Disk full error"))
+
+        with patch.dict(retrieve_task.__globals__, {"download_single_s3_file": mock_download_file}):
+            with pytest.raises(AirflowFailException) as exc_info:
+                retrieve_task(file_key="data/audio_file.wav", params=params)
+
+        assert "Task aborted. Details: Disk full error" in str(exc_info.value)
+