eXParser-PLD/src/APIHandler.py

import os, requests
from dotenv import load_dotenv
from getpass import getpass
import elabapi_python as elabapi


class APIHandler:
    """
    Class which handles all interactions with the eLabFTW API.
    It provides methods to retrieve data from the API and download attachments.
    It relies minimally on the elabapi-python library, which is used only for downloading attachments
    (since the API doesn't support downloading attachments AFAIK).

    Args:
        api_key: str:           A valid API key for the eLabFTW instance where the data is stored, with permissions to access the relevant entries.
                                eLabFTW's API keys are well documented here: https://doc.elabftw.net/docs/usage/api/.
                                If you don't have an API key and are uncapable of creating one, contact your eLabFTW administrator.
                                Or RTFM and create one yourself, it's not that hard.
        ELABFTW_API_URL: str:   Complete URL of the eLabFTW instance's root for the API endpoints.
                                In full caps because it won't (shouldn't) be changed much.
    """

    # TO-DO: remove static url.
    def __init__(self, api_key="", ELABFTW_API_URL=None):
        """Init method, api_key suggested but not required (empty by default)."""
        # if not ELABFTW_API_URL:
        #    load_dotenv()
        #    ELABFTW_API_URL = os.getenv("ELABFTW_API_URL") or input(
        #        "Enter a valid eLabFTW API URL (ends with '/api/v2)': "
        #    )
        self.api_key = api_key
        self.auth = {"Authorization": api_key}
        self.content = {"Content-Type": "application/json"}
        self.header = {**self.auth, **self.content}
        self.elaburl = ELABFTW_API_URL

    def get_entry_from_elabid(self, elabid, entryType="items"):
        """
        Returns raw data (as dictionary) from its elabid and entry type.

        Args:
            elabid: int:     elabftw internal id of the selected resource.
            entryType: str:  Resource type. Anything other than "experiments" or "items" WILL raise an error.
        """
        if entryType not in ["experiments", "items"]:
            raise Exception(
                "You can only download attachments from experiments or items."
            )

        header = self.header
        response = requests.get(
            headers=header, url=f"{self.elaburl}/{entryType}/{elabid}", verify=True
        )

        # Response is 5xx = server error:
        if response.status_code // 100 == 5:
            raise ConnectionError(
                f"There's a problem on the server. Status code: {response.status_code}."
            )

        # Response is 4xx = client error:
        if response.status_code // 100 == 4:
            match response.status_code:
                case 401 | 403:
                    # Forbidden or unauthorized:
                    raise ConnectionError(
                        f"Invalid API key, authentication method or elabid. Check if an item with ID = {elabid} actually exists."
                    )
                case 404:
                    # Lapalissian:
                    raise ConnectionError(
                        "404: Not Found. This means there's no resource with this elabid (wrong elabid?) on your eLabFTW (wrong endpoint?)."
                    )
                case 400:
                    # I genuinely have no idea:
                    raise ConnectionError(
                        "400: Bad Request. This means the API endpoint you tried to reach is invalid. Did you tamper with the source code? If not, contact the developer."
                    )
                case _:
                    # For some fucking reason, this is the only error I actually get from the API...
                    raise ConnectionError(
                        f"HTTP request failed with status code: {response.status_code} (NOTE: 4xx means user's fault)."
                    )

        entry_data = response.json()
        return entry_data

    def download_attachment_data(self, elabid, upload_id, entryType="experiments"):
        """
        Downloads a specific attachment of a certain eLabFTW experiment (default) or item.
        Only returns its binary data. Use method download_attachment_to_disk to save to file.

        NOTE: Output is a dictionary where:
            * The key is the attachment's filename;
            * The value is the attachment's binary data.

        Args:
            elabid: int:     eLabFTW internal ID of the selected resource.
            upload_id: int:  eLabFTW internal ID of the selected upload.
            entryType: str:  Resource type. Anything other than "experiments" or "items" WILL raise an error.
        """
        if entryType not in ["experiments", "items"]:
            raise Exception(
                "You can only download attachments from experiments or items."
            )

        config = elabapi.Configuration()
        config.api_key["api_key"] = self.api_key
        config.api_key_prefix["api_key"] = "Authorization"
        config.host = self.elaburl
        config.debug = False
        api_client = elabapi.ApiClient(config)
        api_client.set_default_header(
            header_name="Authorization", header_value=self.api_key
        )
        uploads_api = elabapi.UploadsApi(api_client)

        # Scans through the attachments and selects the one with corresponing ID.
        attachment = {
            upload.real_name: uploads_api.read_upload(
                entryType, elabid, upload_id, format="binary", _preload_content=False
            ).data
            for upload in uploads_api.read_uploads(entryType, elabid)
            if upload.id == upload_id
        }

        return attachment

    def download_attachment_to_disk(
        self,
        elabid,
        upload_id,
        entryType="experiments",
        dump_dir="output/attachments",
        # persistent=True,
    ):
        """
        Downloads a specific attachment of a certain eLabFTW experiment (default) or item.
        Downloads their binary data through method download_attachments_data and dumps it to dump_dir.
        Returns full path of the output file.

        Args:
            elabid: int:        eLabFTW internal ID of the selected resource.
            upload_id: int:     eLabFTW internal ID of the selected upload.
            entryType: str:     Resource type. Anything other than "experiments" or "items" WILL raise an error.
            dump_dir: str:      Directory to which to save the attachments. Default is "output/attachments".
            persistent: bool:   [Unused] Decides if the files will stay on disk after all operations are completed.
                                If set to False, deletes the file upon exiting. Default = True.
        """

        if entryType not in ["experiments", "items"]:
            raise Exception(
                "You can only download attachments from experiments or items."
            )

        uploads = self.download_attachment_data(elabid, upload_id, entryType=entryType)
        for file in uploads:
            raw_data = uploads[file]
            full_path = os.path.join(dump_dir, f"exp{elabid}-{file}")
            with open(full_path, "wb") as f:
                f.write(raw_data)
        return full_path


# Testing methods
if __name__ == "__main__":
    api_key = getpass("Paste API key here [no echo]: ")
    handler = APIHandler(api_key=api_key)
    handler.download_attachment_to_disk(elabid=58, upload_id=81)