# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
#
# DISCLAIMER: This code is provided as a sample/educational resource only.
# It is NOT intended for use in production environments without thorough review,
# testing, and adaptation to your specific requirements. AWS makes no warranties,
# express or implied, about the completeness, reliability, or suitability of this
# code for any particular purpose. Use at your own risk.
#
# Intended to run as a Glue 5.1 Spark job, as the Data-Analyst definer IAM role,
# in producer account that owns the definer role.
# Glue database, base tables and the view are present in a central catalog account.
#
# This script UPDATES the existing Glue multi-dialect/protected view
# using glueclient.update_table() with
# ViewUpdateAction='REPLACE', adding two additional columns to the view.


import logging
import sys
import os
import boto3
import time
from datetime import datetime
from pyspark.sql import SparkSession

# Logging configuration
formatter = logging.Formatter('[%(asctime)s] %(levelname)s @ line %(lineno)d: %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)
handler.setFormatter(formatter)
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(handler)

catalog_name = "spark_catalog"

# Initialize Spark session
logger.info("Initializing Spark Session")
try:
    spark = (SparkSession
                .builder
                .appName("GlueView-BothDialects")
                .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
                .config(f"spark.sql.catalog.{catalog_name}", "org.apache.iceberg.spark.SparkSessionCatalog")
                .config(f"spark.sql.catalog.{catalog_name}.io-impl", "org.apache.iceberg.aws.s3.S3FileIO")
                .config(f"spark.sql.catalog.{catalog_name}.catalog", "org.apache.iceberg.aws.glue.GlueCatalog")
                .config(f"spark.sql.catalog.{catalog_name}.warehouse", "s3://<your-bucket-name>/<prefix>/transaction-table1")
                .config(f"spark.sql.catalog.{catalog_name}.glue.account-id", "<your-producer-account-id>")
                .config(f"spark.sql.catalog.{catalog_name}.client.region", "<your-region>")
                .config("spark.sql.defaultCatalog", catalog_name)
                .config(f"spark.sql.catalog.{catalog_name}.glue.skip-name-validation", "true")
                .getOrCreate())
except Exception as e:
    logger.error(f"Failed to initialize Spark session: {e}")
    raise

spark.sparkContext.setLogLevel("INFO")
logger.info("Spark session initialized successfully - running as Data-Analyst")

print('boto3 version = ', boto3.__version__)

try:
    glueclient = boto3.client('glue')
    stsclient = boto3.client("sts")
    catalog_id = stsclient.get_caller_identity()["Account"]
except Exception as e:
    logger.error(f"Failed to initialize AWS clients or retrieve account ID: {e}")
    raise

database_rl_name  = 'rl_bankdata_icebergdb'
original_db_name  = 'bankdata_icebergdb'
basetable1_name   = 'transaction_table1'
basetable2_name   = 'transaction_table2'
view_name         = 'view_fromglue'

## sql_query should refer to the original database name and table name, not resource links
## added 'transaction_amount' (from transaction_table1) and 'transaction_date'
## (from transaction_table2) as the additional columns for this update
sql_query = (
    "SELECT a.transaction_id, a.transaction_type, b.transaction_location, "
    "a.transaction_amount, b.transaction_date "
    "FROM bankdata_icebergdb.transaction_table1 a "
    "RIGHT JOIN bankdata_icebergdb.transaction_table2 b "
    "ON a.transaction_id = b.transaction_id"
)

## you need backticks in SQL query if you have special characters or unusual notations
## in database or table names. Underscore is a valid character; no backticks needed.

try:
    logger.info("Showing available databases")
    spark.sql("show databases").show()

    spark.sql(f"show tables from {database_rl_name}").show()

    logger.info("Select from base table using resourcelink of db")
    spark.sql(f"select * from {database_rl_name}.{basetable1_name} limit 7").show()
    spark.sql(f"select * from {database_rl_name}.{basetable2_name} limit 7").show()
except Exception as e:
    logger.error(f"Failed during Spark SQL exploration queries: {e}")
    raise


## Spark SQL cannot be used to update a Glue multi-dialect/protected view.
## You need to use boto3 or the CLI.

def ensure_view_exists(glue_client, catalog_id, database_name, table_name):
    """
    Fail fast if the view does not already exist, instead of relying on
    update_table()/the readiness poller to eventually surface an
    EntityNotFoundException. This addresses the case where someone runs the
    update script before the create script has ever run.
    """
    try:
        glue_client.get_table(
            CatalogId=catalog_id,
            DatabaseName=database_name,
            Name=table_name
        )
        logger.info(f"Confirmed view '{table_name}' exists in database '{database_name}'; proceeding with update.")
    except glue_client.exceptions.EntityNotFoundException:
        raise RuntimeError(
            f"View '{table_name}' does not exist in database '{database_name}'. "
            "Run the create-view script first before running this update script."
        )


try:
    ensure_view_exists(
        glue_client=glueclient,
        catalog_id=catalog_id,
        database_name=database_rl_name,
        table_name=view_name
    )
except RuntimeError as e:
    logger.error(f"Pre-update existence check failed: {e}")
    raise
except Exception as e:
    logger.error(f"Unexpected error while checking view existence: {e}")
    raise


try:
    logger.info(f"Updating Glue view '{view_name}' in database '{database_rl_name}'")
    response_updatetable = glueclient.update_table(
        CatalogId=catalog_id,
        DatabaseName=database_rl_name,  ## use resource link name here, not original database name
        ViewUpdateAction='REPLACE',
        Force=False,
        TableInput={
            'Name': view_name,
            'Description': 'ViewFromGlueBothDialects',
            'StorageDescriptor': {
                'Columns': [
                    {'Name': 'transaction_id',       'Type': 'string'},
                    {'Name': 'transaction_type',     'Type': 'string'},
                    {'Name': 'transaction_location', 'Type': 'string'},
                    {'Name': 'transaction_amount',   'Type': 'double'},  ## newly added column
                    {'Name': 'transaction_date',     'Type': 'date'},    ## newly added column
                ]
            },
            'ViewDefinition': {
                'IsProtected': True,
                'Definer': 'arn:aws:iam::<your-producer-account-id>:role/Data-Analyst',
                'Representations': [
                    {
                        'Dialect': 'SPARK',
                        'DialectVersion': '1.0',
                        'ViewOriginalText': sql_query,
                        'ViewExpandedText': sql_query
                    },
                    {
                        'Dialect': 'ATHENA',
                        'DialectVersion': '3',
                        'ViewOriginalText': sql_query,
                        'ValidationConnection': 'glue-view-validation-connection',
                    }
                ],
                'SubObjects': [
                    'arn:aws:glue:<your-region>:<your-central-account-id>:table/bankdata_icebergdb/transaction_table1',  ## use original database name and table names of source account
                    'arn:aws:glue:<your-region>:<your-central-account-id>:table/bankdata_icebergdb/transaction_table2'
                ]
            }
        }
    )
    logger.info("update_table API call succeeded, polling for view readiness...")
except Exception as e:
    logger.error(f"Failed to update Glue view: {e}")
    raise


def wait_for_view_ready(glue_client, catalog_id, database_name, table_name,
                        poll_interval_seconds=15, max_attempts=8):
    """
    Poll Glue get_table() and check Table.Status.State until it reaches SUCCESS.
    Possible State values observed from the API: SUCCESS | FAILED | IN_PROGRESS.
    Returns True on SUCCESS, raises RuntimeError on FAILED or if max_attempts exceeded.

    Note: existence of the view is validated up front by ensure_view_exists()
    before update_table() is ever called, so EntityNotFoundException here is
    not expected during normal operation and is allowed to propagate
    immediately (fail fast) rather than being treated as a transient/retryable
    condition.
    """
    for attempt in range(1, max_attempts + 1):
        response = glue_client.get_table(
            CatalogId=catalog_id,
            DatabaseName=database_name,
            Name=table_name,
            IncludeStatusDetails=True
        )
        table = response.get('Table', {})
        status = table.get('Status', {})
        state = status.get('State')
        action = status.get('Action', 'UNKNOWN')
        logger.info(f"[Attempt {attempt}/{max_attempts}] View status — Action: {action}, State: {state}")

        if state == 'SUCCESS':
            logger.info(f"View '{table_name}' updated successfully (Action: {action}, State: {state}).")
            return True

        if state == 'FAILED':
            details = status.get('Details', {})
            raise RuntimeError(
                f"View '{table_name}' update failed. Action: {action}, Details: {details}"
            )

        # state is IN_PROGRESS or absent — keep polling
        logger.info(f"View not ready yet (State: {state}), waiting {poll_interval_seconds}s...")
        time.sleep(poll_interval_seconds)

    raise RuntimeError(
        f"View '{table_name}' did not reach SUCCESS after {max_attempts} attempts "
        f"({max_attempts * poll_interval_seconds}s total)."
    )


try:
    wait_for_view_ready(
        glue_client=glueclient,
        catalog_id=catalog_id,
        database_name=database_rl_name,
        table_name=view_name
    )
except RuntimeError as e:
    logger.error(f"View readiness check failed: {e}")
    raise
except Exception as e:
    logger.error(f"Unexpected error while polling for view readiness: {e}")
    raise


try:
    logger.info("-----Select from the updated view-------")
    spark.sql(f"select * from {database_rl_name}.{view_name} limit 7").show()
except Exception as e:
    logger.error(f"Failed to query the updated view: {e}")
    raise

logger.info("Glue view update and query completed successfully.")