Merge pull request #239 from RobertSchaffer1/update_redis_vars

Handling passwords through file, and location to new redis locations

Author: mrakitin

nslsii/__init__.py

@@ -8,6 +8,7 @@
 import appdirs
 from IPython import get_ipython
 
+from .utils import open_redis_client
 from ._version import get_versions
 
 __version__ = get_versions()["version"]
@@ -145,15 +146,9 @@ def configure_base(
         from redis import Redis
         from redis_json_dict import RedisJSONDict
 
-        md = RedisJSONDict(
-            Redis(
-                host=redis_url,
-                port=redis_port,
-                ssl=redis_ssl,
-                password=os.getenv("REDIS_PASSWORD") if redis_ssl else None
-            ),
-            prefix=redis_prefix
-        )
+        redis_client = open_redis_client(redis_ssl=redis_ssl, redis_url=redis_url, redis_prefix=redis_prefix)
+        prefix = redis_prefix if redis_prefix and not redis_ssl else ""
+        md = RedisJSONDict(redis_client=redis_client, prefix=prefix)
 
     # if RunEngine already defined grab it
     # useful when users make their own custom RunEngine

nslsii/sync_experiment/sync_experiment.py

@@ -17,6 +17,8 @@
 
 nslsii_api_client = httpx.Client(base_url="https://api.nsls2.bnl.gov")
 
+from nslsii.utils import open_redis_client
+
 
 def get_current_cycle() -> str:
     cycle_response = nslsii_api_client.get(
@@ -154,6 +156,7 @@ def switch_redis_proposal(
     beamline: str,
     username: Optional[str] = None,
     prefix: str = "",
+    redis_ssl: bool = False,
 ) -> RedisJSONDict:
     """Update information in RedisJSONDict for a specific beamline
 
@@ -167,16 +170,18 @@ def switch_redis_proposal(
         login name of the user assigned to the proposal; if None, current user will be kept
     prefix : str
         optional prefix to identify a specific endstation, e.g. `opls`
+    redis_ssl : bool
+        optional flag to enable/disable ssl connections to redis
 
     Returns
     -------
     md : RedisJSONDict
         The updated redis dictionary.
     """
-
-    redis_client = redis.Redis(host=f"info.{beamline.lower()}.nsls2.bnl.gov")
-    redis_prefix = f"{prefix}-" if prefix else ""
-    md = RedisJSONDict(redis_client=redis_client, prefix=redis_prefix)
+    location = prefix if prefix else beamline
+    redis_client = open_redis_client(redis_ssl=redis_ssl, redis_prefix=location)
+    prefix = f"{prefix}-" if prefix and not redis_ssl else ""
+    md = RedisJSONDict(redis_client=redis_client, prefix=prefix)
     username = username or md.get("username")
 
     new_data_session = f"pass-{proposal_number}"
@@ -229,7 +234,7 @@ def switch_redis_proposal(
     return md
 
 
-def sync_experiment(proposal_number, beamline, verbose=False, prefix=""):
+def sync_experiment(proposal_number, beamline, verbose=False, prefix="", redis_ssl=False):
     # Authenticate the user
     username = input("Username : ")
     authenticate(username)
@@ -241,7 +246,7 @@ def sync_experiment(proposal_number, beamline, verbose=False, prefix=""):
     redis_beamline = normalized_beamlines.get(beamline.lower(), beamline)
 
     md = switch_redis_proposal(
-        proposal_number, beamline=redis_beamline, username=username, prefix=prefix
+        proposal_number, beamline=redis_beamline, username=username, prefix=prefix, redis_ssl=redis_ssl
     )
 
     if verbose:
@@ -269,6 +274,7 @@ def main():
         "--endstation",
         dest="prefix",
         type=str,
+        default="",
         help="Prefix for redis keys (e.g. by endstation)",
         required=False,
     )
@@ -280,6 +286,13 @@ def main():
         help="Which proposal (e.g. 123456)",
         required=True,
     )
+    parser.add_argument(
+        "-s",
+        "--enable-ssl",
+        dest="redis_ssl",
+        action="store_true",
+        help="Flag to enable ssl connection with redis",
+    )
     parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction)
     args = parser.parse_args()
 
@@ -288,4 +301,5 @@ def main():
         beamline=args.beamline,
         verbose=args.verbose,
         prefix=args.prefix,
+        redis_ssl=args.redis_ssl,
     )

nslsii/utils.py

@@ -0,0 +1,99 @@
+import os
+from redis import Redis
+import socket
+
+redis_hosts = [
+    "xf02id1-six-redis1.nsls2.bnl.gov",
+    "xf03id1-hxn-redis1.nsls2.bnl.gov",
+    "xf04bm-maia-redis1.nsls2.bnl.gov",
+    "xf04bm-xfm-redis1.nsls2.bnl.gov",
+    "xf04id1-isr-redis1.nsls2.bnl.gov",
+    "xf05id2-srx-redis1.nsls2.bnl.gov",
+    "xf06bm-bmm-redis1.nsls2.bnl.gov",
+    "xf07bm-qas-redis1.nsls2.bnl.gov",
+    "xf07id1-haxpes-redis1.nsls2.bnl.gov",
+    "xf07id1-nexafs-redis1.nsls2.bnl.gov",
+    "xf07id1-rsoxs-redis1.nsls2.bnl.gov",
+    "xf07id1-ucal-redis1.nsls2.bnl.gov",
+    "xf08bm-tes-redis1.nsls2.bnl.gov",
+    "xf08id1-iss-redis1.nsls2.bnl.gov",
+    "xf09id1-cdi-redis1.nsls2.bnl.gov",
+    "xf10id1-ixs-redis1.nsls2.bnl.gov",
+    "xf11id1-chx-redis1.nsls2.bnl.gov",
+    "xf11bm-cms-redis1.nsls2.bnl.gov",
+    "xf12id1-opls-redis1.nsls2.bnl.gov",
+    "xf12id2-smi-redis1.nsls2.bnl.gov",
+    "xf16id1-lix-redis1.nsls2.bnl.gov",
+    "xf17bm-xfp-redis1.nsls2.bnl.gov",
+    "xf17id1-amx-redis1.nsls2.bnl.gov",
+    "xf17id2-fmx-redis1.nsls2.bnl.gov",
+    "xf18id1-fxi-redis1.nsls2.bnl.gov",
+    "xf19id2-nyx-redis1.nsls2.bnl.gov",
+    "xf21id1-arpes-redis1.nsls2.bnl.gov",
+    "xf21id1-xpeem-redis1.nsls2.bnl.gov",
+    "xf23id1-csx-redis1.nsls2.bnl.gov",
+    "xf23id2-ios-redis1.nsls2.bnl.gov",
+    "xf27id1-hex-redis1.nsls2.bnl.gov",
+    "xf28id1-pdf-redis1.nsls2.bnl.gov",
+    "xf28id2-xpd-redis1.nsls2.bnl.gov",
+    "xf28id2-xpdd-redis1.nsls2.bnl.gov",
+    "xf31id1-tst-redis1.nsls2.bnl.gov",
+]
+
+
+def open_redis_client(
+    redis_url=None,
+    redis_port=None,
+    redis_ssl=False,
+    redis_prefix: str = "",
+    redis_db: int = 0,
+) -> Redis:
+    """
+    Helper function to get the redis client connection.
+    """
+    redis_url = os.getenv("REDIS_HOST", redis_url)
+    if not redis_url:
+        if redis_ssl:
+            client_loc_id = (
+                redis_prefix if redis_prefix else socket.gethostname().split("-")[0]
+            )
+            client_locations = [
+                location for location in redis_hosts if client_loc_id in location
+            ]
+            if len(client_locations) != 1:
+                raise RuntimeError(
+                    "Failed to derive redis server url, please specify using the "
+                    "redis_url argument."
+                )
+            else:
+                redis_url = client_locations[0]
+        else:
+            tla = os.getenv("BEAMLINE_ACRONYM").lower()
+            redis_url = f"info.{tla}.nsls2.bnl.gov"
+
+    if redis_ssl:
+        redis_pw = os.getenv("REDIS_PASSWORD")
+        if not redis_pw:
+            redis_secret_file = os.getenv(
+                "REDIS_SECRET_FILE", "/etc/bluesky/redis.secret"
+            )
+            with open(redis_secret_file, "r", encoding="utf-8") as password_file:
+                redis_pw = password_file.read().strip()
+    else:
+        redis_pw = None
+
+    redis_port = os.getenv("REDIS_PORT", redis_port)
+    if not redis_port:
+        if not redis_ssl:
+            redis_port = 6379
+        else:
+            redis_port = 6380
+
+    conn = Redis(
+        host=redis_url,
+        port=redis_port,
+        ssl=redis_ssl,
+        password=redis_pw,
+        db=redis_db,
+    )
+    return conn