File size: 3,831 Bytes
e9711df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import pandas as pd
import psycopg2
from psycopg2.extras import RealDictCursor
import os
from io import BytesIO
app = FastAPI()
# Database connection parameters
DB_CONFIG = {
"host": "13.126.242.31",
"database": "aml",
"user": "dev_cbs_admin",
"password": "Finovate@2023"
}
def get_db_connection():
"""Create database connection"""
return psycopg2.connect(**DB_CONFIG, cursor_factory=RealDictCursor)
def read_columns_from_file(template_id: str):
"""Read column names from output file"""
filename = f'output_{template_id}.txt'
try:
with open(filename, 'r') as f:
content = f.read()
# Safely evaluate the string representation of the list
columns = eval(content)
return columns
except Exception as e:
raise HTTPException(status_code=404, detail=f"Error reading columns from file: {str(e)}")
@app.get("/download/{template_id}")
async def download_excel(template_id: str):
"""
Download data as Excel file based on template ID
"""
# Read columns from selected file
columns = read_columns_from_file(template_id)
if not columns:
raise HTTPException(status_code=404, detail="No columns found in output file")
try:
# Connect to database
conn = get_db_connection()
cur = conn.cursor()
# Get existing columns from the table
cur.execute("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'large_transactions'
""")
existing_columns = {row['column_name'].lower(): row['column_name'] for row in cur.fetchall()}
# Filter out non-existent columns and use actual column names from DB
valid_columns = []
for col in columns:
if col.lower() in existing_columns:
valid_columns.append(existing_columns[col.lower()])
if not valid_columns:
raise HTTPException(status_code=404, detail="None of the specified columns exist in the table")
# Create SQL query with proper column quoting
quoted_columns = [f'"{col}"' for col in valid_columns]
query = f"""
SELECT {', '.join(quoted_columns)}
FROM large_transactions
"""
# Execute query
cur.execute(query)
rows = cur.fetchall()
# Convert to DataFrame
df = pd.DataFrame(rows)
# Create Excel file in memory
output = BytesIO()
with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
df.to_excel(writer, index=False, sheet_name='Data')
output.seek(0)
# Close database connection
cur.close()
conn.close()
# Return the Excel file as a downloadable response
return StreamingResponse(
output,
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
headers={
"Content-Disposition": f"attachment; filename=large_transactions_{template_id}.xlsx"
}
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error fetching data: {str(e)}")
@app.get("/templates")
async def list_templates():
"""List all available template IDs"""
try:
output_files = [f.replace('output_', '').replace('.txt', '')
for f in os.listdir('.')
if f.startswith('output_') and f.endswith('.txt')]
return {"templates": output_files}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error listing templates: {str(e)}")
|