You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
80 lines
2.6 KiB
80 lines
2.6 KiB
from arango import ArangoClient, exceptions |
|
import pandas as pd |
|
import yaml |
|
|
|
def get_documents(query=False, collection=False, fields=[], filter = '', df=False, index=False, field_names=False, limit=False): |
|
""" |
|
This function retrieves documents from a specified collection or based on a query in ArangoDB. |
|
|
|
Parameters: |
|
query (str): If specified this will be the query. Defaults to False. |
|
collection (str): The name of the collection from which to retrieve documents. Defaults to False. |
|
fields (list): The fields of the documents to retrieve. If empty, all fields are retrieved. |
|
filter (str): AQL filter to apply to the retrieval. Defaults to no filter. |
|
df (bool): If True, the result is returned as a pandas DataFrame. Defaults to False. |
|
index (bool): If True and df is True, the DataFrame is indexed. Defaults to False. |
|
field_names (dict): If provided, these field names will replace the original field names in the result. |
|
|
|
Returns: |
|
list or DataFrame: The retrieved documents as a list of dictionaries or a DataFrame. |
|
""" |
|
if query: |
|
pass |
|
|
|
else: |
|
if fields == []: |
|
return_fields = 'doc' |
|
else: |
|
fields_dict = {} |
|
for field in fields: |
|
fields_dict[field] = field |
|
|
|
if field_names: |
|
for k, v in field_names.items(): |
|
fields_dict[k] = v |
|
|
|
fields_list = [f'{v}: doc.{k}' for k, v in fields_dict.items()] |
|
fields_string = ', '.join(fields_list) |
|
return_fields = f"{{{fields_string}}}" |
|
|
|
if filter != None and filter != '': |
|
filter = f'filter {filter}' |
|
else: |
|
filter = '' |
|
|
|
if limit: |
|
limit = f'limit {limit}' |
|
else: |
|
limit = '' |
|
|
|
query = f''' |
|
for doc |
|
in {collection} |
|
{filter} |
|
{limit} |
|
return {return_fields} |
|
''' |
|
try: |
|
cursor = arango_db.aql.execute(query) |
|
except exceptions.AQLQueryExecuteError: |
|
print('ERROR:\n', query) |
|
exit() |
|
result = [i for i in cursor] |
|
if df: |
|
result = pd.DataFrame(result) |
|
if index: |
|
result.set_index(index, inplace=True) |
|
return result |
|
|
|
|
|
with open('config.yml', 'r') as f: |
|
config = yaml.safe_load(f) |
|
|
|
db = config['arango']['db'] |
|
username = config['arango']['username'] |
|
pwd = config['arango']['pwd_lasse'] |
|
|
|
# Initialize the database for ArangoDB. |
|
client = ArangoClient(hosts=config['arango']['hosts']) |
|
arango_db = client.db(db, username=username, password=pwd) |
|
|
|
|