diff --git a/fragalysis_api/xcextracter/computed_set_update.py b/fragalysis_api/xcextracter/computed_set_update.py index 4601cd1a..069cfe32 100644 --- a/fragalysis_api/xcextracter/computed_set_update.py +++ b/fragalysis_api/xcextracter/computed_set_update.py @@ -1,12 +1,20 @@ +import os import sys import requests -import json import time import threading import _thread as thread REQ_URL = 'https://fragalysis.diamond.ac.uk/viewer/upload_cset/' +# The keycloak credentials (sensitive). +# These are required if you're going to use the authenticated API. +# They will be asserted when the appropriate functions are used. +KEYCLOAK_USERNAME = os.environ.get('KEYCLOAK_USERNAME') +KEYCLOAK_PASSWORD = os.environ.get('KEYCLOAK_PASSWORD') +KEYCLOAK_CLIENT_SECRET = os.environ.get('KEYCLOAK_CLIENT_SECRET') + + def get_csrf(REQ_URL): """Get a csrf token from the request url to authenticate further requests Parameters @@ -31,7 +39,111 @@ def get_csrf(REQ_URL): return csrftoken -def update_cset(REQ_URL, target_name, sdf_path, update_set='None', submit_choice=None, upload_key=None, pdb_zip_path=None, add=False): +def get_keycloak_access_token(*, + keycloak_url, + keycloak_realm, + keycloak_client_id): + """Gets an 'access token' from Keycloak. + If successful we'll get the token (a big long string). + + Each access tokens should last for a number of minutes depending on the + keycloak configuration. + + To use this function you must have set the corresponding + environment variables KEYCLOAK_USERNAME, KEYCLOAK_PASSWORD, + and KEYCLOAK_CLIENT_SECRET + + Parameters + ---------- + keycloak_url string + The keycloak authentication URL (ending in '/auth') + keycloak_realm string + The keycloak realm the stack belongs to + keycloak_client_id string + The client ID the stack is known by in keycloak + + Returns + ------- + access_token string + An API access token + """ + assert KEYCLOAK_USERNAME + assert KEYCLOAK_PASSWORD + assert KEYCLOAK_CLIENT_SECRET + + print(KEYCLOAK_USERNAME) + print(KEYCLOAK_PASSWORD) + print(KEYCLOAK_CLIENT_SECRET) + + realm_url = f'{keycloak_url}/realms/{keycloak_realm}' + url = f'{realm_url}/protocol/openid-connect/token' + data = ( + f'client_id={keycloak_client_id}' + f'&grant_type=password' + f'&username={KEYCLOAK_USERNAME}' + f'&password={KEYCLOAK_PASSWORD}' + f'&client_secret={KEYCLOAK_CLIENT_SECRET}' + ) + headers = {'Content-Type': 'application/x-www-form-urlencoded'} + resp = requests.post(url, headers=headers, data=data, timeout=4.0) + assert resp.status_code == 200 + assert 'access_token' in resp.json() + return resp.json()['access_token'] + +def upload_cset(REQ_URL, access_token, target_name, sdf_path, upload_key=None, pdb_zip_path=None): + """Upload a computed set to Fragalysis + + Parameters + ---------- + REQ_URL: str + request URL for the upload (e.g. https://fagalysis.diamond.ac.uk/viewer/upload_cset/ or viewer/update_cset) + access_token: str + a valid OIDC/Keycloak access token, inserted into the request header as a bearer token + target_name: str + the name of the target in Fragalysis that the computed set is for + sdf_path: str + path to the sdf file to upload + upload_key: str + upload key, not currently turned on, so can be any value, but not blank or null (optional) + pdb_zip_path: str + path to the zip file of pdb's to upload (optional) + """ + assert access_token + + print(f'Uploading new computed set files....') + + csrf_token = get_csrf(REQ_URL) + + payload = {'target_name': target_name, + 'submit_choice': "U", + 'upload_key': upload_key, + 'update_set': None} + + files = [ + ('sdf_file', open(sdf_path,'rb')), + + ] + + if pdb_zip_path: + files.append(('pdb_zip', open(pdb_zip_path,'rb'))) + + headers = {'X-CSRFToken': csrf_token, + 'Cookie': f'csrftoken={csrf_token}', + 'Authorization': f'Bearer {access_token}'} + + response = requests.request("POST", REQ_URL, headers=headers, data=payload, files=files) + + taskurl = get_taskidurl(response=response, REQ_URL=REQ_URL) + + if not taskurl: + raise Exception(f'Something went wrong with the upload request! \ + Please try again or email warren.thompson@diamond.ac.uk for help.\ + Response: {response.text}') + + return taskurl + + +def update_cset(REQ_URL, access_token, target_name, sdf_path, cset_name, upload_key=None, pdb_zip_path=None, add=False): """Send data to /viewer/upload_cset/ to overwrite an existing computed set, or to /viewer/update_cset/ to add new molecules without deleting the old ones. @@ -39,9 +151,11 @@ def update_cset(REQ_URL, target_name, sdf_path, update_set='None', submit_choice ---------- REQ_URL: str request URL for the upload (e.g. https://fagalysis.diamond.ac.uk/viewer/upload_cset/ or viewer/update_cset) + access_token: str + a valid OIDC/Keycloak access token, inserted into the request header as a bearer token target_name: str the name of the target in Fragalysis that the computed set is for - update_set: str + cset_name: str the name of the computed set you want to update (can be found with: "".join(submitter_name.split()) + '-' + "".join(method.split()), where submitter_name is the name in the submitter_name field in the blank mol of the uploaded sdf file, @@ -49,32 +163,32 @@ def update_cset(REQ_URL, target_name, sdf_path, update_set='None', submit_choice a set for the first time sdf_path: str path to the sdf file to upload - submit_choice: int - 0 for validate, 1 for upload (not required for update - ie. viewer/update_cset) upload_key: str upload key, not currently turned on, so can be any value, but not blank or null (optional) pdb_zip_path: str path to the zip file of pdb's to upload (optional) add: bool set to True if updating a computed set without overwriting it completely (for /viewer/update_cset/) - + Returns ------- - taskurl str - the URL to check for the status of the upload + taskurl: str + The URL of the Celery task """ - print(f'Submitting files to update {update_set}...') + assert access_token + + print(f'Submitting files to update {cset_name}...') csrf_token = get_csrf(REQ_URL) - + if not add: payload = {'target_name': target_name, - 'submit_choice': submit_choice, + 'submit_choice': "U", 'upload_key': upload_key, - 'update_set': update_set} + 'update_set': cset_name} else: payload = {'target_name': target_name, - 'update_set': update_set} + 'update_set': cset_name} files = [ ('sdf_file', open(sdf_path,'rb')), @@ -85,10 +199,98 @@ def update_cset(REQ_URL, target_name, sdf_path, update_set='None', submit_choice files.append(('pdb_zip', open(pdb_zip_path,'rb'))) headers = {'X-CSRFToken': csrf_token, - 'Cookie': f'csrftoken={csrf_token}'} + 'Cookie': f'csrftoken={csrf_token}', + 'Authorization': f'Bearer {access_token}'} response = requests.request("POST", REQ_URL, headers=headers, data=payload, files=files) + taskurl = get_taskidurl(response=response, REQ_URL=REQ_URL) + + if not taskurl: + raise Exception(f'Something went wrong with the update request! \ + Please try again or email warren.thompson@diamond.ac.uk for help.\ + Response: {response.text}') + + return taskurl + + +def delete_cset(REQ_URL, access_token, target_name, cset_name, upload_key=None): + """Delete uploaded computer set + + Parameters + ---------- + REQ_URL: str + request URL for the upload (e.g. https://fagalysis.diamond.ac.uk/viewer/upload_cset/ or viewer/update_cset) + access_token: str + a valid OIDC/Keycloak access token, inserted into the request header as a bearer token + target_name: str + the name of the target in Fragalysis that the computed set is for + cset_name: str + the name of the computed set you want to delete + (can be found with: "".join(submitter_name.split()) + '-' + "".join(method.split()), + where submitter_name is the name in the submitter_name field in the blank mol of the uploaded sdf file, + and method is the method field in the blank mol of the uploaded sdf file). Leave blank if you are adding + a set for the first time + upload_key: str + upload key, not currently turned on, so can be any value, but not blank or null (optional) + """ + assert access_token + + print(f'Deleting {cset_name}...') + + csrf_token = get_csrf(REQ_URL) + + payload = {'target_name': target_name, + 'submit_choice': "D", + 'upload_key': upload_key, + 'update_set': cset_name} + + headers = {'X-CSRFToken': csrf_token, + 'Cookie': f'csrftoken={csrf_token}', + 'Authorization': f'Bearer {access_token}'} + + response = requests.request("POST", REQ_URL, headers=headers, data=payload) + + check_url_csetdeleted(cset_name=cset_name, response=response) + +def check_url_csetdeleted(cset_name: str, response: str): + """Checks if response yields delete success + + Parameters + ---------- + cset_name: str + the name of the computed set you want to delete + (can be found with: "".join(submitter_name.split()) + '-' + "".join(method.split()), + where submitter_name is the name in the submitter_name field in the blank mol of the uploaded sdf file, + and method is the method field in the blank mol of the uploaded sdf file). Leave blank if you are adding + a set for the first time + response: str + The response from the update cset call + REQ_URL: str + request URL for the upload (e.g. https://fagalysis.diamond.ac.uk/viewer/upload_cset/ or viewer/update_cset) + """ + lines = response.text.split('\n') + delete_check = None + for l in lines: + if f'

Compound set "{cset_name}" deleted

' in l: + delete_check = True + print(f'Computed set: {cset_name} deleted') + break + if not delete_check: + raise Exception(f'Something went wrong with the delete request! \ + Please try again or email warren.thompson@diamond.ac.uk for help.\ + Response: {response.text}') + +def get_taskidurl(response: str, REQ_URL: str): + """Checks if response yields a Celery task ID and returns the URL + + Parameters + ---------- + response: str + The response from the update cset call + REQ_URL: str + request URL for the upload (e.g. https://fagalysis.diamond.ac.uk/viewer/upload_cset/ or viewer/update_cset) + """ lines = response.text.split('\n') taskurl = None for l in lines: @@ -100,14 +302,8 @@ def update_cset(REQ_URL, target_name, sdf_path, update_set='None', submit_choice taskid = l.split('/')[-2] print(f'upload task id: {taskid}') taskurl = f'{REQ_URL.replace("/viewer/update_cset/","/viewer/update_task/")}{taskid}' - if taskurl: break - if not taskurl: - raise Exception(f'Something went wrong with the upload/update request! \ - Please try again or email rachael.skyner@diamond.ac.uk for help.\ - Response: {response.text}') - return taskurl @@ -184,13 +380,48 @@ def get_task_response(taskurl): # ---- NB: The major difference here is the REQ_URL. For new data, or to overwrite data, use https://fragalysis.diamond.ac.uk/viewer/upload_cset/. # To add new molecules to an existing set, use https://fragalysis.diamond.ac.uk/viewer/update_cset/. ---- # +# You must set the following environment variables +# to avoid assertions when you run the code: +# +# - KEYCLOAK_USERNAME (a valid user for the chosen stack) +# - KEYCLOAK_PASSWORD (the user's password) +# - KEYCLOAK_CLIENT_SECRET (this is a stack-specific uuid4-like value) +# +# i.e.: +# +# export KEYCLOAK_USERNAME=someone +# export KEYCLOAK_PASSWORD=someone1234 +# export KEYCLOAK_CLIENT_SECRET=00000000-0000-0000-0000-000000000000 +# +# to get an OIDC/Keycloak access token: +# ------------------------------------- +# Here we get a token for the staging stack... +# +# access_token = get_keycloak_access_token( +# keycloak_url = 'https://keycloak.xchem.diamond.ac.uk/auth', +# keycloak_realm = 'xchem', +# keycloak_client_id = 'fragalysis-xchem') +# print(f'access_token="{access_token}"') +# +# to upload a new cset: +# ------------------------------ +# taskurl = upload_cset(REQ_URL='https://fragalysis.xchem.diamond.ac.uk/viewer/upload_cset/', +# access_token=access_token, +# target_name='Mpro', +# upload_key='1', +# sdf_path='/Users/uzw12877/Downloads/Test_upload/Top_100_three_hop_XCOS_1.4_2020-07-28.sdf', +# pdb_zip_path='/Users/uzw12877/Downloads/Test_upload/receptor_pdbs.zip') +# task_response, json_results = get_task_response(taskurl) +# print(task_response) +# print(json_results) +# # to overwrite an existing cset: # ------------------------------ -# taskurl = update_cset(REQ_URL='https://fragalysis.diamond.ac.uk/viewer/upload_cset/', +# taskurl = update_cset(REQ_URL='https://fragalysis.xchem.diamond.ac.uk/viewer/upload_cset/', +# access_token=access_token, # target_name='Mpro', -# submit_choice='1', # upload_key='1', -# update_set='WT-xCOS3-ThreeHop', +# cset_name='WT-xCOS3-ThreeHop', # sdf_path='/Users/uzw12877/Downloads/Test_upload/Top_100_three_hop_XCOS_1.4_2020-07-28.sdf', # pdb_zip_path='/Users/uzw12877/Downloads/Test_upload/receptor_pdbs.zip') # task_response, json_results = get_task_response(taskurl) @@ -199,13 +430,24 @@ def get_task_response(taskurl): # # to update an existing cset: # --------------------------- -# taskurl = update_cset(REQ_URL='https://fragalysis.diamond.ac.uk/viewer/update_cset/', +# taskurl = update_cset(REQ_URL='https://fragalysis.xchem.diamond.ac.uk/viewer/update_cset/', +# access_token=access_token, # target_name='Mpro', -# update_set='WT-xCOS3-ThreeHop', +# cset_name='WT-xCOS3-ThreeHop', # sdf_path='/Users/uzw12877/Downloads/Test_upload/Top_100_three_hop_XCOS_1.4_2020-07-28 copy.sdf', # pdb_zip_path='/Users/uzw12877/Downloads/Test_upload/receptor_pdbs copy.zip', # add=True) # task_response, json_results = get_task_response(taskurl) # print(task_response) # print(json_results) - +# +# to delete an existing cset: +# --------------------------- +# delete_cset(REQ_URL='https://fragalysis.xchem.diamond.ac.uk/viewer/update_cset/', +# access_token=access_token, +# target_name='Mpro', +# cset_name='WT-xCOS3-ThreeHop', +# upload_key='1') +# task_response, json_results = get_task_response(taskurl) +# print(task_response) +# print(json_results) \ No newline at end of file