from config import Config
from .base import BaseUtil
from typing import List
# Secrets Manager id holding the long lived IAM user credentials used to sign S3 URLs.
# Presigned URLs signed with temporary (task role / STS) credentials die with the
# session token, which caps them far below the 7 day maximum, so signing is done with
# the automation user instead.
SIGNING_SECRET_ID = "airflow/connections/airflow_automation_user"
# field names inside that secret's json payload
ACCESS_KEY_FIELD = "ACCESS_KEY"
SECRET_KEY_FIELD = "SECRET_KEY"
SESSION_TOKEN_FIELD = "SESSION_TOKEN" # optional; only set for temporary credentials
# secrets manager lookups are cached per process; get_presigned_url is called in loops
_signing_credentials_cache = {}
[docs]
class S3Util:
"""
S3 utilities for airflow
"""
[docs]
@staticmethod
def get_signing_credentials(
secret_id: str = SIGNING_SECRET_ID, refresh: bool = False
) -> dict:
"""
Read the automation user credentials from Secrets Manager as boto3 client kwargs
:param secret_id: the secrets manager secret id holding the credentials json
:param refresh: skip the in process cache and re-read the secret
:return: boto3 credential kwargs, or an empty dict when the secret is missing or
incomplete, in which case the caller falls back to the default credential
chain (task role) rather than failing outright
"""
import json
if not refresh and secret_id in _signing_credentials_cache:
return _signing_credentials_cache[secret_id]
try:
secret = json.loads(BaseUtil.get_aws_secret(secret_id))
except Exception as e:
print(f"Couldn't read signing credentials from '{secret_id}': {e}")
return {}
access_key = secret.get(ACCESS_KEY_FIELD)
secret_key = secret.get(SECRET_KEY_FIELD)
if not access_key or not secret_key:
# never log the values, only which fields were present
print(
f"Secret '{secret_id}' has no "
f"{ACCESS_KEY_FIELD}/{SECRET_KEY_FIELD} pair. "
f"Fields found: {sorted(secret)}"
)
return {}
credentials = {
"aws_access_key_id": access_key,
"aws_secret_access_key": secret_key,
}
session_token = secret.get(SESSION_TOKEN_FIELD)
if session_token:
# temporary credentials: the url will expire with the token, not ExpiresIn
print(f"Secret '{secret_id}' carries a session token, expect a short ttl")
credentials["aws_session_token"] = session_token
_signing_credentials_cache[secret_id] = credentials
return credentials
[docs]
@staticmethod
def store_s3_file(bucket: str, key: str, file_path: str, overwrite=True) -> dict:
"""
Store a file on s3
:param bucket: the bucket
:param key: the key in s3
:param file_path: the file path to store from
:param overwrite: call the ifNoneMatch value (to prevent overwriting)
:return: the response from boto
"""
import boto3
client = boto3.client("s3")
if not overwrite:
return client.put_object(
Body=open(file_path, "rb").read(),
Bucket=bucket,
Key=key,
ifNoneMatch="*", # Prevent overwriting
)
else:
print("Careful. Allowing file overwriting now")
with open(file_path, "rb") as rfh:
return client.put_object(
Body=rfh.read(),
Bucket=bucket,
Key=key,
)
[docs]
@staticmethod
def get_presigned_url(
bucket: str,
key: str,
expiration: int = 604800,
response_type=None,
*,
region_name: str = None,
signature_version: str = "s3v4",
) -> str:
"""
Get a presigned URL for an S3 object
Signed with the automation user credentials from Secrets Manager so the url
actually lives for the full expiration, see get_signing_credentials
:param bucket: the bucket name
:param key: the key of the object
:param expiration: the expiration time in seconds (default: 604800 seconds = 7 days, Max)
:param response_type: optional ResponseContentType to force on the download
The two below are keyword only, so the first four positions stay exactly as
callers already pass them (e.g. bucket, key, expiration, response_type=...)
:param region_name: the bucket's region. sigv4 urls must be signed for the
region the bucket lives in, so pass it when the bucket is not in the
session default region. None leaves it to the boto session
:param signature_version: s3v4 (default) or s3 for legacy sigv2 query auth.
boto3 presigns s3 with sigv2 unless told otherwise, and buckets created
after June 2020 reject sigv2 outright, so v4 is the safer default
"""
import boto3
from botocore.client import Config as BotoConfig
client = boto3.client(
"s3",
region_name=region_name,
config=BotoConfig(signature_version=signature_version),
**S3Util.get_signing_credentials(),
)
params={"Bucket": bucket, "Key": key}
if response_type:
params['ResponseContentType']= response_type
try:
url = client.generate_presigned_url(
ClientMethod="get_object",
Params=params,
ExpiresIn=expiration,
)
except Exception as e:
print(f"Couldn't get a presigned URL for client method '{str(e)}'.")
raise
return url
[docs]
@staticmethod
def delete_s3_folder(
bucket: str,
prefix: str,
override_panic: bool = False,
max_delete: int = 1_000_000,
) -> list[dict]:
"""
Delete a folder from s3
Ensure prefix ends with a slash so it's a folder due to key prefix matching behavior of s3
Panics if the prefix ends with any of the panic prefixes in the config
:param bucket: the bucket to delete from
:param prefix: the prefix match to delete, be careful with partial matches
:param override_panic: whether or not to override the panic prefixes in the config
:param max_delete: the threshold/max authorized amount of items to delete,
:return: the response from boto
"""
import boto3
if not prefix.endswith("/"):
prefix += "/"
if (
any(
# e.g. transformed/cai/sbo/temp_enrich_raw/ => will not panic on `raw`, and delete folder `temp_enrich_raw`
prefix[:-1].split("/")[-1] == each_prefix
for each_prefix in Config.PANIC_PREFIXES
)
and not override_panic
):
raise ValueError(f"Cannot delete {prefix}")
client = boto3.client("s3")
all_keys = S3Util.get_existing_s3_files(
bucket=bucket, prefix=prefix, file_type="", full_key=True
)
print(
f"""
About to delete {len(all_keys)} files from {bucket} with prefix {prefix}
First five keys (head) samples: {all_keys[:5]}
"""
)
if len(all_keys) > max_delete:
raise ValueError(
f"Cannot delete a total of {len(all_keys)} files; need to set higher {max_delete=}"
)
return [client.delete_object(Bucket=bucket, Key=key) for key in all_keys]
[docs]
@staticmethod
def get_existing_s3_files(
bucket: str,
prefix: str,
file_type: str = "csv",
full_key: bool = False,
return_raw_objects=False,
) -> list[str]:
"""
Get files from s3 that match the given file type. Deals with pagination.
:param str bucket: the bucket to get from
:param str prefix: the prefix to match
:param str, optional file_type: the file type to match, without leading dot. Default: csv
:param bool, optional full_key: whether or not to return the full key. Default: False
:return: the list of files
"""
import boto3
client = boto3.client("s3")
s3_resp = client.list_objects_v2(Bucket=bucket, Prefix=prefix)
all_objects = []
while s3_resp.get("IsTruncated", False):
all_objects.extend(s3_resp["Contents"])
s3_resp = client.list_objects_v2(
Bucket=bucket,
Prefix=prefix,
ContinuationToken=s3_resp["NextContinuationToken"],
)
all_objects.extend(s3_resp.get("Contents", []))
if return_raw_objects:
return [
each_file
for each_file in all_objects
if each_file["Key"].endswith(file_type)
]
else:
return [
each_file["Key"].split("/")[-1] if not full_key else each_file["Key"]
for each_file in all_objects
if each_file["Key"].endswith(file_type)
]
[docs]
@staticmethod
def s3_url_to_table(url: str) -> str:
"""
Convert an s3 url to a table name
:param url: the url to convert
:return: the table name
"""
return str(
BaseUtil.normalize_string(url)
.replace(f"s3__{Config.DEFAULT_BUCKET}", "")
.replace("_manual_ingress_", "")
.replace("_automated_ingress_", "")
.replace("_xlsx", "")
)[:53]
[docs]
@staticmethod
def get_raw_csv_columns(
bucket: str, prefix: str, header_row: int = 0
) -> dict[str, str]:
"""
Get arbitrary file from s3, returns a dict of column names as strings
:param bucket: the bucket to get from
:param prefix: the prefix to match
:param header_row: row to use as header (0-based index). Default: 0 (first row)
:return: the dict of column names
"""
import boto3
import pandas as pd
if not prefix.endswith("/"):
prefix += "/"
client = boto3.client("s3")
# Get S3 response and handle empty results
response = client.list_objects_v2(Bucket=bucket, Prefix=prefix)
if "Contents" not in response:
raise ValueError(f"No files found in s3://{bucket}/{prefix}")
# Get the first non-empty file
valid_files = [
each_file["Key"] for each_file in response["Contents"] if each_file["Size"]
]
if not valid_files:
raise ValueError(f"No non-empty files found in s3://{bucket}/{prefix}")
tgt_file = valid_files[0]
print(f"Reading columns from: s3://{bucket}/{tgt_file}")
df = None
all_exceptions = []
for encoding in ["utf-8", "cp1252", "latin1", "iso-8859-1"]:
try:
# If header_row is 0, this is the default behavior (first row is header)
# If header_row > 0, we'll skip that many rows and use the next row as header
df = pd.read_csv(
f"s3://{bucket}/{tgt_file}",
encoding=encoding,
header=header_row, # Use specified header row
nrows=1, # Only need one row of data
)
print(
f"Successfully read file with {encoding} encoding, header at row {header_row}"
)
break
except UnicodeDecodeError as ude:
print(f"Failed to read with {encoding} encoding, trying next...")
all_exceptions.append(f"{encoding}: {str(ude)}")
except Exception as e:
print(f"Error reading file with {encoding}: {str(e)}")
all_exceptions.append(f"{encoding}: {str(e)}")
# If all encoding attempts failed
if df is None:
error_msg = "\n".join(all_exceptions)
raise ValueError(
f"Failed to read file with any encoding: s3://{bucket}/{tgt_file}\nErrors:\n{error_msg}"
)
# Verify we got valid column names
if df.empty or len(df.columns) == 0:
raise ValueError(
f"No valid columns found in the file: s3://{bucket}/{tgt_file}"
)
# Check if any column names are unnamed/numeric indices (which happens when header=None)
unnamed_columns = [col for col in df.columns if str(col).startswith("Unnamed:")]
if unnamed_columns and len(unnamed_columns) / len(df.columns) > 0.5:
print(
f"Warning: Most columns appear unnamed. Header row {header_row} may not contain column names."
)
return {col: "string" for col in df.columns}
[docs]
@staticmethod
def copy_between_buckets(
source_path: str,
target_path: str,
pattern: str = "",
) -> List[str]:
"""Copies a set of s3 objects between two buckets. it will replace data in the destination if they already exist
:param str source_path: the source path to copy data from. make sure to include s3://
:param str target_path: the target path to copy data to. make sure to include s3://
:param str | None pattern: file name regex pattern to match. Default: None (matches all files)
:return: List of new objects paths.
"""
import awswrangler as wr
import re
all_source_keys = wr.s3.list_objects(source_path)
if pattern:
all_source_keys = [
key for key in all_source_keys if re.match(pattern, key.split("/")[-1])
]
print(f"{len(all_source_keys)} files will be copied between two buckets!")
copied_files = wr.s3.copy_objects(
paths=all_source_keys,
source_path=source_path,
target_path=target_path,
)
return copied_files
[docs]
@staticmethod
def folder_exists_and_not_empty(bucket: str, path: str) -> bool:
"""
Check if a folder exists and is not empty.
:param bucket: the bucket name
:type bucket: str
:param path: prefix
:type path: str
:return: True if folder exists and not empty
:rtype: bool
"""
import awswrangler as wr
# Ensure path ends with '/'
path = path.rstrip("/") + "/"
# Remove leading '/' if present
path = path.lstrip("/")
s3_full_directory = f"s3://{bucket}/{path}"
# Use list_objects with chunked=True to get an iterator
iter_objects = wr.s3.list_objects(s3_full_directory, chunked=True)
try:
next(iter_objects)
return True
except StopIteration:
# If StopIteration is raised, the iterator is empty
return False
[docs]
@staticmethod
def download_folder(
s3_path: str, local_path: str, recursive=True, verbose=False
) -> None:
"""
Download all files in the directory from s3 to local airflow server.
Recommend passing a local_path that is created from TemporaryDirectory() to ensure cleanup.
:param s3_path: the s3 path of the folder to download
:param local_path: the local directory to download the contents into
:param recursive: whether to download recursively, or ignore subdirectories
:param verbose: whether to print out the files being downloaded
"""
import awswrangler as wr
import os
clean_s3_path = s3_path.rstrip("/") + "/"
clean_local_path = local_path.rstrip("/") + "/"
files = wr.s3.list_objects(s3_path)
for file in files:
relative_file_path = file[len(clean_s3_path) :].lstrip("/")
local_file_path = clean_local_path + relative_file_path
if not recursive and "/" in relative_file_path:
continue
if verbose:
print(f"Downloading {file} to {local_file_path}")
os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
wr.s3.download(file, local_file_path)