From 3c498317f3968af83606e748bd740ecde79dbd1f Mon Sep 17 00:00:00 2001 From: Eugene M Date: Tue, 28 May 2024 12:27:41 -0400 Subject: [PATCH 1/2] ENH: add utitlity to start/switch beamline experiment --- nslsii/sync_redis/__init__.py | 1 + nslsii/sync_redis/sync_redis.py | 139 ++++++++++++++++++++++++++++++++ requirements.txt | 6 +- setup.py | 6 ++ 4 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 nslsii/sync_redis/__init__.py create mode 100644 nslsii/sync_redis/sync_redis.py diff --git a/nslsii/sync_redis/__init__.py b/nslsii/sync_redis/__init__.py new file mode 100644 index 00000000..d09eb803 --- /dev/null +++ b/nslsii/sync_redis/__init__.py @@ -0,0 +1 @@ +from .sync_redis import main, switch_experiment, validate_proposal \ No newline at end of file diff --git a/nslsii/sync_redis/sync_redis.py b/nslsii/sync_redis/sync_redis.py new file mode 100644 index 00000000..0b93d492 --- /dev/null +++ b/nslsii/sync_redis/sync_redis.py @@ -0,0 +1,139 @@ +from ldap3 import Server, Connection, NTLM +from ldap3.core.exceptions import LDAPInvalidCredentialsResult + +import json +import re +import redis +import httpx +import warnings +from datetime import datetime +from getpass import getpass +from redis_json_dict import RedisJSONDict +from typing import Dict, Any +import argparse + +data_session_re = re.compile(r"^pass-(?P\d+)$") + +nslsii_api_client = httpx.Client(base_url="https://api.nsls2.bnl.gov") + + +def get_current_cycle() -> str: + cycle_response = nslsii_api_client.get(f"/v1/facility/nsls2/cycles/current").raise_for_status() + return cycle_response.json()["cycle"] + + +def validate_proposal(data_session_value, beamline) -> Dict[str, Any]: + + data_session_match = data_session_re.match(data_session_value) + + if data_session_match is None: + raise ValueError(f"RE.md['data_session']='{data_session_value}' " f"is not matched by regular expression '{data_session_re.pattern}'") + + try: + current_cycle = get_current_cycle() + proposal_number = data_session_match.group("proposal_number") + proposal_response = nslsii_api_client.get(f"/v1/proposal/{proposal_number}").raise_for_status() + proposal_data = proposal_response.json()["proposal"] + if "error_message" in proposal_data: + raise ValueError( + f"while verifying data_session '{data_session_value}' " f"an error was returned by {proposal_response.url}: " f"{proposal_data}" + ) + else: + if current_cycle not in proposal_data["cycles"]: + raise ValueError(f"Proposal {data_session_value} is not valid in the current NSLS2 cycle ({current_cycle}).") + if beamline.upper() not in proposal_data["instruments"]: + raise ValueError( + f"Wrong beamline ({beamline.upper()}) for proposal {data_session_value} ({', '.join(proposal_data['instruments'])})." + ) + # data_session is valid! + + except httpx.RequestError as rerr: + # give the user a warning but allow the run to start + proposal_data = {} + warnings.warn(f"while verifying data_session '{data_session_value}' " f"the request {rerr.request.url!r} failed with " f"'{rerr}'") + + finally: + return proposal_data + + +def authenticate(username): + + auth_server = Server("dc2.bnl.gov", use_ssl=True) + + try: + connection = Connection( + auth_server, user=f"BNL\\{username}", password=getpass("Password : "), authentication=NTLM, auto_bind=True, raise_exceptions=True + ) + print(f"\nAuthenticated as : {connection.extend.standard.who_am_i()}") + + except LDAPInvalidCredentialsResult: + raise RuntimeError(f"Invalid credentials for user '{username}'.") from None + + +def should_they_be_here(username, new_data_session, beamline): + + user_access_json = nslsii_api_client.get(f"/v1/data-session/{username}").json() + + if "nsls2" in user_access_json["facility_all_access"]: + return True + + elif beamline.lower() in user_access_json["beamline_all_access"]: + return True + + elif new_data_session in user_access_json["data_sessions"]: + return True + + return False + + +class AuthorizationError(Exception): ... + + +def switch_experiment(proposal_number, beamline, verbose=False, prefix=""): + + redis_client = redis.Redis(host=f"info.{beamline.lower()}.nsls2.bnl.gov") + + md = RedisJSONDict(redis_client=redis_client, prefix=prefix) + + new_data_session = f"pass-{proposal_number}" + username = input("Username : ") + + if (new_data_session == md.get("data_session")) and (username == md.get("username")): + + warnings.warn(f"Experiment {new_data_session} was already started by the same user.") + + else: + + proposal_data = validate_proposal(new_data_session, beamline) + users = proposal_data.pop("users") + + authenticate(username) + + if not should_they_be_here(username, new_data_session, beamline): + raise AuthorizationError(f"User '{username}' is not allowed to take data on proposal {new_data_session}") + + md["data_session"] = new_data_session + md["username"] = username + md["start_datetime"] = datetime.now().isoformat() + md["cycle"] = get_current_cycle() + + print(f"Started experiment {new_data_session}.") + if verbose: + print(json.dumps(proposal_data, indent=2)) + print("Users:") + print(users) + + +def main(): + # Used by the `sync-redis` command + + parser = argparse.ArgumentParser(description="Start or switch beamline experiment and record it in Redis") + parser.add_argument("-b", "--beamline", dest="beamline", type=str, help="Which beamline (e.g. CHX)", required=True) + parser.add_argument("-p", "--proposal", dest="proposal", type=int, help="Which proposal (e.g. 123456)", required=True) + parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction) + + args = parser.parse_args() + switch_experiment(proposal_number=args.proposal, beamline=args.beamline, verbose=args.verbose) + + +# e.g. start_experiment(proposal_number=314062, beamline="chx") diff --git a/requirements.txt b/requirements.txt index 0df30124..7cdb2611 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,13 +4,17 @@ bluesky-kafka>=0.8.0 caproto databroker h5py +httpx ipython ipywidgets +ldap3 matplotlib msgpack >=1.0.0 msgpack-numpy numpy ophyd psutil +pycryptodome pyolog -redis \ No newline at end of file +redis +redis-json-dict \ No newline at end of file diff --git a/setup.py b/setup.py index 31de79ae..8e8563fb 100644 --- a/setup.py +++ b/setup.py @@ -27,4 +27,10 @@ description='Tools for data collection and analysis at NSLS-II', author='Brookhaven National Laboratory', install_requires=requirements, + entry_points={ + "console_scripts": [ + # 'command = some.module:some_function', + "sync-redis = nslsii.sync_redis:main", + ], + }, ) From 14f27ff5f96bb22d95a8790f2b810adb94a4e5ef Mon Sep 17 00:00:00 2001 From: Eugene M Date: Tue, 28 May 2024 13:44:16 -0400 Subject: [PATCH 2/2] cleanup --- nslsii/sync_redis/__init__.py | 2 +- nslsii/sync_redis/sync_redis.py | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/nslsii/sync_redis/__init__.py b/nslsii/sync_redis/__init__.py index d09eb803..c56bc8d5 100644 --- a/nslsii/sync_redis/__init__.py +++ b/nslsii/sync_redis/__init__.py @@ -1 +1 @@ -from .sync_redis import main, switch_experiment, validate_proposal \ No newline at end of file +from .sync_redis import main, switch_experiment, validate_proposal diff --git a/nslsii/sync_redis/sync_redis.py b/nslsii/sync_redis/sync_redis.py index 0b93d492..6982dccd 100644 --- a/nslsii/sync_redis/sync_redis.py +++ b/nslsii/sync_redis/sync_redis.py @@ -134,6 +134,3 @@ def main(): args = parser.parse_args() switch_experiment(proposal_number=args.proposal, beamline=args.beamline, verbose=args.verbose) - - -# e.g. start_experiment(proposal_number=314062, beamline="chx")