Python Guides

Upload Files

Set up environment

!pip install -U athenaintel
import os
import io
import asyncio
from athena.client import Athena
ATHENA_API_KEY = os.environ["ATHENA_API_KEY"]
athena = Athena(
    api_key=ATHENA_API_KEY,
)

Upload a data frame or a plot

Athena SDK can save pandas DataFrames as assets out of the box:

data_frame = pd.DataFrame({"column": [1, 2, 3]})
result = athena.tools.save_asset(data_frame, name="My data frame")
print(result)

The result will contain the asset_id.

It can also save most plots objects which support IPython display protocol, for example:

fig = plt.figure(figsize=(10, 6))
plt.plot([0, 1, 2, 3, 4, 5], [1, 2, 4, 8, 16, 32])
plt.show()
client.tools.save_asset(fig, name="Powers of two")

You can set the parent_folder_id to upload the file (or plot/data frame) to a specific folder.

Upload a file

To upload a file using the Athena SDK, you can use the athena.tools.save_asset() method. This method accepts a file tuple containing the filename, file content as a BytesIO object, and the MIME type.

Here’s an example of how to upload an Excel file:

async def upload_file():
    # Prepare the file for upload
    file_bytes = io.BytesIO()
    with open("your_file.xlsx", "rb") as f:
        file_bytes.write(f.read())
        file_bytes.seek(0)  # Reset the cursor of the BytesIO object
    # Create the file tuple
    file_tuple = (
        "your_file.xlsx",
        file_bytes,
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    )
    # Upload the file
    result = await athena.tools.save_asset(file_tuple)
    print(result)
# Run the async function
await upload_file()

In this example:

  1. We open the file and read its contents into a BytesIO object.
  2. We create a tuple containing the filename, the BytesIO object, and the MIME type.
  3. We call athena.tools.save_asset() with our file tuple.
  4. The function returns the result of the upload operation.

Using with FastAPI

If you’re using FastAPI and want to upload files received from a client, you can use the UploadFile object:

from fastapi import FastAPI, UploadFile
app = FastAPI()
@app.post("/upload/")
async def upload_file(file: UploadFile):
    file_tuple = (file.filename, file.file, file.content_type)
    result = await athena.tools.save_asset(file_tuple)
    return {"message": "File uploaded successfully", "result": result}

This endpoint will accept file uploads and forward them to the Athena API using the SDK.

Remember to handle exceptions and implement proper error checking in your production code.