AWSTemplateFormatVersion: '2010-09-09'
Description: >
  Athena Workgroup and Query Cost Review Remediator.
  Review-only module that inventories Athena workgroups, workgroup cost controls,
  query result settings, encryption, named queries, data catalogs, and recent
  query execution samples across enabled Regions. This module intentionally does
  not stop, delete, or modify Athena resources.

Parameters:
  ProjectName:
    Type: String
    Default: athena-cost-review
    Description: Short base name used for created resources.

  IncludeDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: Include detailed Athena inventory in Lambda response.

  IncludeTags:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: Include tags where supported.

  MaxQueryExecutionsPerWorkgroup:
    Type: Number
    Default: 50
    MinValue: 0
    MaxValue: 50
    Description: Maximum recent query executions to sample per workgroup. Set 0 to skip query sampling.

  LargeQueryBytesWarningThreshold:
    Type: Number
    Default: 10737418240
    MinValue: 1
    Description: Per-query data scanned threshold in bytes for review flag. Default is 10 GiB.

  HighSampleBytesWarningThreshold:
    Type: Number
    Default: 107374182400
    MinValue: 1
    Description: Per-workgroup sampled data scanned threshold in bytes for review flag. Default is 100 GiB.

  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
    Description: Comma-separated Regions to skip.

  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
    Description: Whether to create an enabled EventBridge schedule.

  ScheduleExpression:
    Type: String
    Default: rate(30 days)
    Description: EventBridge schedule expression used when ScheduleEnabled is true.

Conditions:
  IsScheduleEnabled: !Equals [!Ref ScheduleEnabled, 'true']

Resources:
  ReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: AthenaCostReviewPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: RegionDiscovery
                Effect: Allow
                Action:
                  - ec2:DescribeRegions
                Resource: '*'
              - Sid: AthenaRead
                Effect: Allow
                Action:
                  - athena:ListWorkGroups
                  - athena:GetWorkGroup
                  - athena:ListTagsForResource
                  - athena:ListQueryExecutions
                  - athena:BatchGetQueryExecution
                  - athena:ListNamedQueries
                  - athena:ListDataCatalogs
                  - athena:GetDataCatalog
                Resource: '*'

  ReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Ref ProjectName
      Description: Review-only Lambda for Athena workgroup and query cost visibility.
      Runtime: python3.12
      Handler: index.handler
      Timeout: 900
      MemorySize: 512
      Role: !GetAtt ReviewRole.Arn
      Environment:
        Variables:
          INCLUDE_DETAILS: !Ref IncludeDetails
          INCLUDE_TAGS: !Ref IncludeTags
          MAX_QUERY_EXECUTIONS_PER_WORKGROUP: !Ref MaxQueryExecutionsPerWorkgroup
          LARGE_QUERY_BYTES_WARNING_THRESHOLD: !Ref LargeQueryBytesWarningThreshold
          HIGH_SAMPLE_BYTES_WARNING_THRESHOLD: !Ref HighSampleBytesWarningThreshold
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
      Code:
        ZipFile: |
          import json
          import os
          from datetime import datetime, timezone

          import boto3
          from botocore.exceptions import ClientError, EndpointConnectionError

          def utc_now():
              return datetime.now(timezone.utc)

          def env_bool(name, default=False):
              return os.environ.get(name, str(default)).strip().lower() in ("1", "true", "yes", "y")

          def parse_bool(value, default=False):
              if value is None:
                  return default
              return str(value).strip().lower() in ("1", "true", "yes", "y")

          def csv_list(value):
              raw = value or ""
              return [x.strip() for x in raw.split(",") if x.strip()]

          def error_record(service, operation, err):
              if isinstance(err, ClientError):
                  return {
                      "service": service,
                      "operation": operation,
                      "error_code": err.response.get("Error", {}).get("Code", "Unknown"),
                      "error_message": err.response.get("Error", {}).get("Message", str(err))
                  }
              return {
                  "service": service,
                  "operation": operation,
                  "error_code": type(err).__name__,
                  "error_message": str(err)
              }

          def is_access_denied(err):
              text = str(err).lower()
              if isinstance(err, ClientError):
                  code = err.response.get("Error", {}).get("Code", "").lower()
                  msg = err.response.get("Error", {}).get("Message", "").lower()
                  text = f"{code} {msg}"
              return "accessdenied" in text or "not authorized" in text or "explicit deny" in text or "service control policy" in text

          def is_not_found_or_unsupported(err):
              if isinstance(err, ClientError):
                  code = err.response.get("Error", {}).get("Code", "").lower()
                  text = (code + " " + err.response.get("Error", {}).get("Message", "")).lower()
                  return (
                      "invalidrequestexception" in text and "not found" in text or
                      "resourcenotfound" in text or
                      "unsupportedoperation" in text
                  )
              return False

          def safe_call(region_result, service_name, operation_name, func, default):
              try:
                  return func()
              except EndpointConnectionError as e:
                  region_result["failed"] += 1
                  region_result["failures"].append(error_record(service_name, operation_name, e))
                  return default
              except ClientError as e:
                  if is_not_found_or_unsupported(e):
                      return default
                  if is_access_denied(e):
                      region_result["access_denied"] += 1
                  else:
                      region_result["failed"] += 1
                  region_result["failures"].append(error_record(service_name, operation_name, e))
                  return default
              except Exception as e:
                  region_result["failed"] += 1
                  region_result["failures"].append(error_record(service_name, operation_name, e))
                  return default

          def get_enabled_regions(excluded):
              ec2 = boto3.client("ec2", region_name="us-east-1")
              response = ec2.describe_regions(AllRegions=True)
              regions = []
              for r in response.get("Regions", []):
                  status = r.get("OptInStatus")
                  name = r.get("RegionName")
                  if status in (None, "opt-in-not-required", "opted-in") and name not in excluded:
                      regions.append(name)
              return sorted(regions)

          def tags_to_dict(tags):
              out = {}
              for tag in tags or []:
                  key = tag.get("Key") or tag.get("key")
                  if key:
                      out[key] = tag.get("Value") or tag.get("value")
              return out

          def list_workgroups(client):
              items = []
              token = None
              while True:
                  kwargs = {}
                  if token:
                      kwargs["NextToken"] = token
                  response = client.list_work_groups(**kwargs)
                  items.extend(response.get("WorkGroups", []) or [])
                  token = response.get("NextToken")
                  if not token:
                      return items

          def list_data_catalogs(client):
              items = []
              token = None
              while True:
                  kwargs = {"MaxResults": 50}
                  if token:
                      kwargs["NextToken"] = token
                  response = client.list_data_catalogs(**kwargs)
                  items.extend(response.get("DataCatalogsSummary", []) or [])
                  token = response.get("NextToken")
                  if not token:
                      return items

          def count_named_queries(client, workgroup):
              total = 0
              token = None
              while True:
                  kwargs = {"WorkGroup": workgroup}
                  if token:
                      kwargs["NextToken"] = token
                  response = client.list_named_queries(**kwargs)
                  total += len(response.get("NamedQueryIds", []) or [])
                  token = response.get("NextToken")
                  if not token:
                      return total

          def list_query_execution_ids(client, workgroup, max_items):
              if max_items <= 0:
                  return []
              ids = []
              token = None
              while True:
                  kwargs = {"WorkGroup": workgroup, "MaxResults": min(50, max_items)}
                  if token:
                      kwargs["NextToken"] = token
                  response = client.list_query_executions(**kwargs)
                  ids.extend(response.get("QueryExecutionIds", []) or [])
                  if len(ids) >= max_items:
                      return ids[:max_items]
                  token = response.get("NextToken")
                  if not token:
                      return ids

          def batch_get_query_executions(client, ids):
              out = []
              for i in range(0, len(ids), 50):
                  chunk = ids[i:i+50]
                  if not chunk:
                      continue
                  response = client.batch_get_query_execution(QueryExecutionIds=chunk)
                  out.extend(response.get("QueryExecutions", []) or [])
              return out

          def workgroup_item(wg, detail, tags, query_stats, include_tags):
              cfg = detail.get("Configuration", {}) or {}
              result_cfg = cfg.get("ResultConfiguration", {}) or {}
              enc_cfg = result_cfg.get("EncryptionConfiguration", {}) or {}
              engine = cfg.get("EngineVersion", {}) or {}
              item = {
                  "name": detail.get("Name") or wg.get("Name"),
                  "state": detail.get("State") or wg.get("State"),
                  "creation_time": detail.get("CreationTime") or wg.get("CreationTime"),
                  "description": detail.get("Description"),
                  "enforce_workgroup_configuration": cfg.get("EnforceWorkGroupConfiguration"),
                  "publish_cloudwatch_metrics_enabled": cfg.get("PublishCloudWatchMetricsEnabled"),
                  "bytes_scanned_cutoff_per_query": cfg.get("BytesScannedCutoffPerQuery"),
                  "requester_pays_enabled": cfg.get("RequesterPaysEnabled"),
                  "result_output_location_present": bool(result_cfg.get("OutputLocation")),
                  "expected_bucket_owner_present": bool(result_cfg.get("ExpectedBucketOwner")),
                  "acl_configuration_present": bool(result_cfg.get("AclConfiguration")),
                  "encryption_option": enc_cfg.get("EncryptionOption"),
                  "kms_key_present": bool(enc_cfg.get("KmsKey")),
                  "selected_engine_version": engine.get("SelectedEngineVersion"),
                  "effective_engine_version": engine.get("EffectiveEngineVersion"),
                  **query_stats
              }
              if include_tags:
                  item["tags"] = tags or {}
              return item

          def query_stats(executions, thresholds):
              stats = {
                  "sampled_queries": len(executions),
                  "sampled_succeeded_queries": 0,
                  "sampled_failed_queries": 0,
                  "sampled_cancelled_queries": 0,
                  "sampled_running_queries": 0,
                  "sampled_queued_queries": 0,
                  "sampled_total_bytes_scanned": 0,
                  "sampled_largest_query_bytes_scanned": 0,
                  "sampled_large_queries": 0,
                  "latest_sampled_query_submission_time": None
              }
              latest = None
              for q in executions or []:
                  status = (q.get("Status") or {}).get("State")
                  if status == "SUCCEEDED":
                      stats["sampled_succeeded_queries"] += 1
                  elif status == "FAILED":
                      stats["sampled_failed_queries"] += 1
                  elif status == "CANCELLED":
                      stats["sampled_cancelled_queries"] += 1
                  elif status == "RUNNING":
                      stats["sampled_running_queries"] += 1
                  elif status == "QUEUED":
                      stats["sampled_queued_queries"] += 1

                  data = ((q.get("Statistics") or {}).get("DataScannedInBytes") or 0)
                  stats["sampled_total_bytes_scanned"] += int(data)
                  stats["sampled_largest_query_bytes_scanned"] = max(stats["sampled_largest_query_bytes_scanned"], int(data))
                  if int(data) >= thresholds["large_query_bytes_warning_threshold"]:
                      stats["sampled_large_queries"] += 1

                  submitted = (q.get("Status") or {}).get("SubmissionDateTime")
                  if submitted and (latest is None or submitted > latest):
                      latest = submitted

              stats["latest_sampled_query_submission_time"] = latest
              return stats

          def scan_region(region, settings):
              rr = {
                  "region": region,
                  "workgroups": 0,
                  "enabled_workgroups": 0,
                  "disabled_workgroups": 0,
                  "workgroups_enforcing_configuration": 0,
                  "workgroups_with_cloudwatch_metrics": 0,
                  "workgroups_with_bytes_scanned_cutoff": 0,
                  "workgroups_with_result_location": 0,
                  "workgroups_with_encryption": 0,
                  "workgroups_with_kms_encryption": 0,
                  "named_queries": 0,
                  "data_catalogs": 0,
                  "sampled_queries": 0,
                  "sampled_succeeded_queries": 0,
                  "sampled_failed_queries": 0,
                  "sampled_cancelled_queries": 0,
                  "sampled_running_queries": 0,
                  "sampled_queued_queries": 0,
                  "sampled_total_bytes_scanned": 0,
                  "sampled_largest_query_bytes_scanned": 0,
                  "sampled_large_queries": 0,
                  "workgroups_over_sample_bytes_threshold": 0,
                  "review_flags": [],
                  "workgroup_inventory": [],
                  "data_catalog_inventory": [],
                  "access_denied": 0,
                  "failed": 0,
                  "failures": []
              }

              ath = boto3.client("athena", region_name=region)

              workgroups = safe_call(rr, "Athena", "ListWorkGroups", lambda: list_workgroups(ath), [])
              rr["workgroups"] = len(workgroups)

              catalogs = safe_call(rr, "Athena", "ListDataCatalogs", lambda: list_data_catalogs(ath), [])
              rr["data_catalogs"] = len(catalogs)
              if settings["include_details"]:
                  rr["data_catalog_inventory"] = [
                      {"name": c.get("CatalogName"), "type": c.get("Type")} for c in catalogs
                  ]

              for wg in workgroups:
                  name = wg.get("Name")
                  detail_resp = safe_call(rr, "Athena", "GetWorkGroup", lambda n=name: ath.get_work_group(WorkGroup=n), {})
                  detail = detail_resp.get("WorkGroup", {}) or {}
                  state = detail.get("State") or wg.get("State")
                  if state == "ENABLED":
                      rr["enabled_workgroups"] += 1
                  elif state == "DISABLED":
                      rr["disabled_workgroups"] += 1

                  cfg = detail.get("Configuration", {}) or {}
                  result_cfg = cfg.get("ResultConfiguration", {}) or {}
                  enc_cfg = result_cfg.get("EncryptionConfiguration", {}) or {}

                  if cfg.get("EnforceWorkGroupConfiguration"):
                      rr["workgroups_enforcing_configuration"] += 1
                  if cfg.get("PublishCloudWatchMetricsEnabled"):
                      rr["workgroups_with_cloudwatch_metrics"] += 1
                  if cfg.get("BytesScannedCutoffPerQuery"):
                      rr["workgroups_with_bytes_scanned_cutoff"] += 1
                  if result_cfg.get("OutputLocation"):
                      rr["workgroups_with_result_location"] += 1
                  if enc_cfg.get("EncryptionOption"):
                      rr["workgroups_with_encryption"] += 1
                  if enc_cfg.get("KmsKey"):
                      rr["workgroups_with_kms_encryption"] += 1

                  nq = safe_call(rr, "Athena", "ListNamedQueries", lambda n=name: count_named_queries(ath, n), 0)
                  rr["named_queries"] += nq

                  ids = safe_call(
                      rr,
                      "Athena",
                      "ListQueryExecutions",
                      lambda n=name: list_query_execution_ids(ath, n, settings["max_query_executions_per_workgroup"]),
                      []
                  )
                  executions = safe_call(
                      rr,
                      "Athena",
                      "BatchGetQueryExecution",
                      lambda qids=ids: batch_get_query_executions(ath, qids),
                      []
                  ) if ids else []

                  qstats = query_stats(executions, settings)
                  for key in [
                      "sampled_queries", "sampled_succeeded_queries", "sampled_failed_queries",
                      "sampled_cancelled_queries", "sampled_running_queries", "sampled_queued_queries",
                      "sampled_total_bytes_scanned", "sampled_largest_query_bytes_scanned",
                      "sampled_large_queries"
                  ]:
                      rr[key] += int(qstats.get(key, 0) or 0)

                  if qstats["sampled_total_bytes_scanned"] >= settings["high_sample_bytes_warning_threshold"]:
                      rr["workgroups_over_sample_bytes_threshold"] += 1

                  tags = {}
                  if settings["include_tags"] and detail.get("Arn"):
                      tag_resp = safe_call(
                          rr,
                          "Athena",
                          "ListTagsForResource",
                          lambda arn=detail.get("Arn"): ath.list_tags_for_resource(ResourceARN=arn).get("Tags", []),
                          []
                      )
                      tags = tags_to_dict(tag_resp)

                  if settings["include_details"]:
                      rr["workgroup_inventory"].append(workgroup_item(wg, detail, tags, qstats, settings["include_tags"]))

              if rr["workgroups"] > 0:
                  rr["review_flags"].append("athena_workgroups_present")
              if rr["workgroups_with_bytes_scanned_cutoff"] < rr["enabled_workgroups"] and rr["enabled_workgroups"] > 0:
                  rr["review_flags"].append("workgroups_without_bytes_scanned_cutoff")
              if rr["workgroups_enforcing_configuration"] < rr["enabled_workgroups"] and rr["enabled_workgroups"] > 0:
                  rr["review_flags"].append("workgroups_not_enforcing_configuration")
              if rr["sampled_large_queries"] > 0:
                  rr["review_flags"].append("large_sampled_queries_present")
              if rr["workgroups_over_sample_bytes_threshold"] > 0:
                  rr["review_flags"].append("workgroups_over_sample_bytes_threshold")
              if rr["sampled_failed_queries"] > 0 or rr["sampled_cancelled_queries"] > 0:
                  rr["review_flags"].append("failed_or_cancelled_queries_sampled")

              return rr

          def handler(event, context):
              event = event or {}
              excluded = set(csv_list(event.get("excluded_regions_csv") or os.environ.get("EXCLUDED_REGIONS_CSV", "ap-northeast-3")))
              settings = {
                  "include_details": parse_bool(event.get("include_details"), env_bool("INCLUDE_DETAILS", True)),
                  "include_tags": parse_bool(event.get("include_tags"), env_bool("INCLUDE_TAGS", True)),
                  "max_query_executions_per_workgroup": int(event.get("max_query_executions_per_workgroup") or os.environ.get("MAX_QUERY_EXECUTIONS_PER_WORKGROUP", "50")),
                  "large_query_bytes_warning_threshold": int(event.get("large_query_bytes_warning_threshold") or os.environ.get("LARGE_QUERY_BYTES_WARNING_THRESHOLD", "10737418240")),
                  "high_sample_bytes_warning_threshold": int(event.get("high_sample_bytes_warning_threshold") or os.environ.get("HIGH_SAMPLE_BYTES_WARNING_THRESHOLD", "107374182400"))
              }

              summary = {
                  "started_at": utc_now().isoformat(),
                  "review_only": True,
                  "include_details": settings["include_details"],
                  "include_tags": settings["include_tags"],
                  "excluded_regions": sorted(list(excluded)),
                  "query_sampling": {
                      "max_query_executions_per_workgroup": settings["max_query_executions_per_workgroup"],
                      "large_query_bytes_warning_threshold": settings["large_query_bytes_warning_threshold"],
                      "high_sample_bytes_warning_threshold": settings["high_sample_bytes_warning_threshold"]
                  },
                  "totals": {
                      "regions_scanned": 0,
                      "workgroups": 0,
                      "enabled_workgroups": 0,
                      "disabled_workgroups": 0,
                      "workgroups_enforcing_configuration": 0,
                      "workgroups_with_cloudwatch_metrics": 0,
                      "workgroups_with_bytes_scanned_cutoff": 0,
                      "workgroups_with_result_location": 0,
                      "workgroups_with_encryption": 0,
                      "workgroups_with_kms_encryption": 0,
                      "named_queries": 0,
                      "data_catalogs": 0,
                      "sampled_queries": 0,
                      "sampled_succeeded_queries": 0,
                      "sampled_failed_queries": 0,
                      "sampled_cancelled_queries": 0,
                      "sampled_running_queries": 0,
                      "sampled_queued_queries": 0,
                      "sampled_total_bytes_scanned": 0,
                      "sampled_largest_query_bytes_scanned": 0,
                      "sampled_large_queries": 0,
                      "workgroups_over_sample_bytes_threshold": 0,
                      "regions_with_workgroups": 0,
                      "regions_with_review_flags": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "region_results": [],
                  "failures": []
              }

              try:
                  regions = get_enabled_regions(excluded)
              except Exception as e:
                  summary["totals"]["failed"] += 1
                  summary["failures"].append(error_record("EC2", "DescribeRegions", e))
                  summary["ended_at"] = utc_now().isoformat()
                  print(json.dumps(summary, indent=2, default=str))
                  return json.loads(json.dumps(summary, default=str))

              total_keys = [
                  k for k in summary["totals"].keys()
                  if k not in ("regions_scanned", "regions_with_workgroups", "regions_with_review_flags")
              ]

              for region in regions:
                  rr = scan_region(region, settings)
                  for key in total_keys:
                      summary["totals"][key] += int(rr.get(key, 0) or 0)

                  if rr.get("workgroups", 0) > 0:
                      summary["totals"]["regions_with_workgroups"] += 1
                  if rr.get("review_flags"):
                      summary["totals"]["regions_with_review_flags"] += 1

                  summary["totals"]["regions_scanned"] += 1
                  summary["failures"].extend([dict(f, region=region) for f in rr.get("failures", [])])
                  summary["region_results"].append(rr)

              summary["ended_at"] = utc_now().isoformat()
              print(json.dumps(summary, indent=2, default=str))
              return json.loads(json.dumps(summary, default=str))

  ReviewScheduleRule:
    Type: AWS::Events::Rule
    Condition: IsScheduleEnabled
    Properties:
      Name: !Sub '${ProjectName}-schedule'
      Description: Scheduled run for Athena Cost Review Remediator.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt ReviewFunction.Arn
          Id: AthenaCostReviewTarget
          Input: '{"source":"eventbridge-schedule"}'

  ReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: IsScheduleEnabled
    Properties:
      FunctionName: !Ref ReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt ReviewScheduleRule.Arn

Outputs:
  ReviewFunctionName:
    Description: Lambda function name.
    Value: !Ref ReviewFunction

  RemediatorFunctionName:
    Description: Lambda function name alias.
    Value: !Ref ReviewFunction

  ReviewFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt ReviewFunction.Arn

  ReviewOnly:
    Description: This module is intentionally review-only.
    Value: 'true'
