mala/_arango.py

import os
import re
#import pandas as pd
from arango import ArangoClient, exceptions
from arango.database import StandardDatabase
from dotenv import load_dotenv

class ArangoDB:
    def __init__(self, database=None):
            """
            Initializes a connection to an ArangoDB database using the configuration
            """

            load_dotenv(".env")

            host = os.environ['ARANGO_HOSTS']
            if database:
                db = database
            else:
                db = os.environ['ARANGO_DB']
            username = os.environ['ARANGO_USERNAME']
            pwd = os.environ['ARANGO_PWD_LASSE']

            # Initialize the database for ArangoDB.
            self.client: ArangoClient = ArangoClient(hosts=host)
            self.db: StandardDatabase = self.client.db(db, username=username, password=pwd)

    def fix_key_name(self, string):
        """
        Makes a string a valid ArangoDB key name.

        Args:
        string (str): The string to fix.

        Returns:
        str: The fixed string.
        """
        string = string.replace("å", "a").replace('ä', 'a').replace('ö', 'o').replace('Å', 'A').replace('Ä', 'A').replace('Ö', 'O')
        string = re.sub(r"[^a-zA-Z0-9_\_\-\.\@\(\)\+,=\;\$\!\*\'\%]", "_", string)
        encoded_string = string.encode('utf-8')
        if len(encoded_string) > 254:
            string = encoded_string[:254].decode('utf-8', 'ignore')
        return string
    
    def get_persons(self, confirmed=True):
        """
        Retrieves a list of persons from the database.

        Args:
            confirmed (bool, optional): If True, only retrieves confirmed persons. Defaults to True.

        Returns:
            dict: A dictionary containing two keys:
                - 'names': A list of person names.
                - 'dict': A dictionary mapping person names to their corresponding keys.
        """
        confirmed_string = ''
        if confirmed:
            confirmed_string = 'filter person.confirmed == true'
                 
        query = f"""
        FOR person IN persons
        {confirmed_string}
        RETURN {{'name': person.name, '_key': person._key}}
        """
        persons = [i for i in self.db.aql.execute(query)]
        names  = [document['name'] for document in persons]
        dict_persons = {document['name']: document['_key'] for document in persons}
        return {'names': names, 'dict':dict_persons}

arango = ArangoDB()
db = arango.db

if __name__ == '__main__':
    arango = ArangoDB()
    print(len(arango.get_persons(confirmed=False)['names']))