"""
README Generator Operator
This module provides an operator for generating README files from a template
and uploading them to Google Drive or S3. The README serves as a warning to
users that exported data should not be modified and re-imported.
Usage (Google Drive):
from operators.cya.generate_readme.export import GenerateReadmeOperator
upload_readme = GenerateReadmeOperator(
task_id="upload_readme",
google_drive_id="1ABC123xyz",
client_file_name="export_data.xlsx",
overwrite=True,
)
Usage (S3):
upload_readme = GenerateReadmeOperator(
task_id="upload_readme",
s3_path="s3://my-bucket/exports/",
client_file_name="export_data.xlsx",
overwrite=True,
)
"""
from __future__ import annotations
import os
import re
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Sequence
from airflow.exceptions import AirflowException
from airflow.utils.context import Context
from googleapiclient.errors import HttpError
from googleapiclient.http import MediaFileUpload
from jinja2 import Template, UndefinedError
from operators.gdrive.gdrive_base import GDriveBaseOperator
################################################################################
## Constants
################################################################################
TEMPLATE_PATH = Path(__file__).parent / "template" / "readme.txt"
################################################################################
## Operator
################################################################################
[docs]
class GenerateReadmeOperator(GDriveBaseOperator):
"""
Generate a README file from template and upload to Google Drive or S3.
The Jinja2 template placeholders {{ timestamp }} and {{ client_file }} are
replaced with the logical execution date and the provided client file name.
Exactly one of google_drive_id or s3_path must be provided.
:param google_drive_id: The Google Drive folder ID to upload to.
:param s3_path: The S3 path to upload to (e.g., "s3://bucket/prefix/").
:param client_file_name: The name of the client file this README accompanies.
:param overwrite: If True, overwrite existing file with same name. If False, create new file.
"""
ui_color = "#A6E6A6"
template_fields: Sequence[str] = GDriveBaseOperator.template_fields + (
"client_file_name",
"s3_path",
)
def __init__(
self,
*,
google_drive_id: str | None = None,
s3_path: str | None = None,
client_file_name: str,
overwrite: bool = False,
**kwargs,
):
super().__init__(google_drive_id=google_drive_id, **kwargs)
self.s3_path = s3_path
self.client_file_name = client_file_name
self.overwrite = overwrite
############################################################################
## Template Rendering
############################################################################
def _render_template_content(self, logical_date: datetime) -> str:
"""Load the README template and render with Jinja2."""
try:
template_content = TEMPLATE_PATH.read_text(encoding="utf-8")
except FileNotFoundError as e:
raise AirflowException(
f"README template not found at {TEMPLATE_PATH}"
) from e
try:
template = Template(template_content)
return template.render(
timestamp=logical_date.strftime("%Y-%m-%d %H:%M:%S UTC"),
client_file=self.client_file_name,
)
except UndefinedError as e:
raise AirflowException(f"Template rendering failed: {e}") from e
def _get_readme_filename(self) -> str:
"""Derive the README filename from client_file_name, sanitizing special chars."""
stem = Path(self.client_file_name).stem
stem = re.sub(r"[^\w.-]+", "_", stem).strip("_") or "file"
return f"README_{stem}.txt"
############################################################################
## S3 Helpers
############################################################################
def _parse_s3_path(self, s3_path: str) -> tuple[str, str]:
"""Parse 's3://bucket/prefix/' into (bucket, prefix)."""
path = s3_path.replace("s3://", "")
parts = path.split("/", 1)
bucket = parts[0]
prefix = parts[1] if len(parts) > 1 else ""
if prefix and not prefix.endswith("/"):
prefix += "/"
return bucket, prefix
def _upload_to_s3(self, local_path: str, filename: str) -> dict:
"""Upload README to S3 using instance IAM credentials."""
import boto3
bucket, prefix = self._parse_s3_path(self.s3_path)
key = f"{prefix}{filename}"
s3 = boto3.client("s3")
if not self.overwrite:
try:
s3.head_object(Bucket=bucket, Key=key)
self.log.info(
"File %s already exists and overwrite=False; creating with unique name",
key,
)
timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
name_parts = filename.rsplit(".", 1)
if len(name_parts) == 2:
filename = f"{name_parts[0]}_{timestamp}.{name_parts[1]}"
else:
filename = f"{filename}_{timestamp}"
key = f"{prefix}{filename}"
except s3.exceptions.ClientError:
pass
s3.upload_file(local_path, bucket, key)
full_path = f"s3://{bucket}/{key}"
self.log.info("Uploaded %s to %s", filename, full_path)
return {"s3_path": full_path, "bucket": bucket, "key": key}
############################################################################
## Drive File Lookup
############################################################################
def _find_existing_file(self, filename: str) -> str | None:
"""Search for existing file by name in the target folder. Returns first match ID."""
safe_name = filename.replace("\\", "\\\\").replace("'", "\\'")
safe_folder_id = self.google_drive_id.replace("\\", "\\\\").replace("'", "\\'")
query = (
f"'{safe_folder_id}' in parents and "
f"name = '{safe_name}' and "
"trashed = false"
)
result = (
self.drive.files()
.list(
q=query,
spaces="drive",
fields="files(id)",
pageSize=1,
supportsAllDrives=True,
includeItemsFromAllDrives=True,
)
.execute()
)
files = result.get("files", [])
return files[0]["id"] if files else None
############################################################################
## GDrive Upload
############################################################################
def _upload_to_gdrive(self, local_path: str, filename: str) -> dict:
"""Upload README to Google Drive folder."""
media = MediaFileUpload(local_path, mimetype="text/plain")
existing_file_id = None
if self.overwrite:
existing_file_id = self._find_existing_file(filename)
if existing_file_id:
self.log.info("Overwriting existing README (id=%s)", existing_file_id)
try:
uploaded = (
self.drive.files()
.update(
fileId=existing_file_id,
media_body=media,
fields="id, name, webViewLink",
supportsAllDrives=True,
)
.execute()
)
except HttpError as e:
if e.resp is not None and e.resp.status == 404:
self.log.warning(
"Existing README (id=%s) not found during update; creating new",
existing_file_id,
)
existing_file_id = None
else:
raise
if not existing_file_id:
if self.overwrite:
self.log.info(
"overwrite=True but no existing README '%s' found; creating new",
filename,
)
file_metadata = {
"name": filename,
"parents": [self.google_drive_id],
}
uploaded = (
self.drive.files()
.create(
body=file_metadata,
media_body=media,
fields="id, name, webViewLink",
supportsAllDrives=True,
)
.execute()
)
self.log.info("Uploaded %s (id=%s)", uploaded.get("name"), uploaded.get("id"))
return uploaded
############################################################################
## Execute
############################################################################
[docs]
def execute(self, context: Context) -> dict:
has_gdrive = bool(self.google_drive_id)
has_s3 = bool(self.s3_path)
if has_gdrive == has_s3:
raise AirflowException(
"Exactly one of google_drive_id or s3_path must be provided"
)
if not self.client_file_name:
raise AirflowException("client_file_name must be provided and non-empty")
logical_date = context.get("logical_date") or datetime.now(timezone.utc)
content = self._render_template_content(logical_date)
readme_filename = self._get_readme_filename()
destination = self.google_drive_id if has_gdrive else self.s3_path
self.log.info(
"Uploading README for '%s' to %s",
self.client_file_name,
destination,
)
with tempfile.TemporaryDirectory() as tmpdir:
local_path = os.path.join(tmpdir, readme_filename)
with open(local_path, "w", encoding="utf-8") as f:
f.write(content)
if has_gdrive:
return self._upload_to_gdrive(local_path, readme_filename)
else:
return self._upload_to_s3(local_path, readme_filename)