import boto3
import time

WORKGROUP = "iceberg-workgroup"
DATABASE = "dnx_warehouse"

session = boto3.Session()
athena = session.client("athena", region_name="us-east-1")

# ==============================================
# QUERY
# ==============================================

QUERY = """
SELECT title,shortname from AwsDataCatalog.dnx_warehouse.bacio_transacional_loja_app_pesquisa;
"""

print("Executando query no Athena...")
response = athena.start_query_execution(
    QueryString=QUERY,
    QueryExecutionContext={"Database": DATABASE},
    WorkGroup=WORKGROUP
)

query_execution_id = response["QueryExecutionId"]
print(f"QueryExecutionId: {query_execution_id}")

# ==============================================
# AGUARDAR RESULTADO
# ==============================================

while True:
    result = athena.get_query_execution(QueryExecutionId=query_execution_id)
    state = result["QueryExecution"]["Status"]["State"]

    if state in ["SUCCEEDED", "FAILED", "CANCELLED"]:
        print("Estado final:", state)
        break

    print("Aguardando execução...")
    time.sleep(1)

# ==============================================
# RESULTADO
# ==============================================

if state == "SUCCEEDED":
    output = athena.get_query_results(QueryExecutionId=query_execution_id)
    print("\nResultados:")
    for row in output["ResultSet"]["Rows"]:
        print([col.get("VarCharValue", "") for col in row["Data"]][0])
else:
    print("Erro ao executar a query.")