From 9638896d2e1c744374f3fe51697462a65dfde1f5 Mon Sep 17 00:00:00 2001
From: Lanqing Yuan <yuanlq@uchicago.edu>
Date: Sat, 10 Jun 2023 14:25:32 -0500
Subject: [PATCH] automatically appending local rucio path (#1182)

* automatically appending midway rucio

* annotations

* do nothing if on dali compute node

* better coding style

* better coding style

* better coding style

* better coding style

* Fix some style

* Little bug fix

---------

Co-authored-by: dachengxu <dx2227@columbia.edu>
---
 straxen/contexts.py | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/straxen/contexts.py b/straxen/contexts.py
index 8ce4962ce..684e1defd 100644
--- a/straxen/contexts.py
+++ b/straxen/contexts.py
@@ -7,6 +7,8 @@
 import warnings
 import typing as ty
 from pandas.util._decorators import deprecate_kwarg
+import socket
+
 
 common_opts = dict(
     register_all=[],
@@ -109,6 +111,34 @@ def xenonnt(cmt_version='global_ONLINE', xedocs_version=None,
     return st
 
 
+def find_rucio_local_path(include_rucio_local, _rucio_local_path):
+    """
+    Check the hostname to determine which rucio local path to use. Note that access to
+    /dali/lgrandi/rucio/ is possible only if you are on dali compute node or login node.
+
+    :param include_rucio_local: add the rucio local storage frontend.
+        This is only needed if one wants to do a fuzzy search in the
+        data the runs database is out of sync with rucio
+    :param _rucio_local_path: str, path of local RSE of rucio. Only use
+        for testing!
+    """
+    hostname = socket.gethostname()
+    # if you are on dali compute node, do nothing
+    if ('dali' in hostname) and ('login' not in hostname):
+        _include_rucio_local = include_rucio_local
+        __rucio_local_path = _rucio_local_path
+    # Assumed the only other option is 'midway' or login nodes, 
+    # where we have full access to dali and project space. 
+    # This doesn't make sense outside XENON but we don't care.
+    else:
+        _include_rucio_local = True
+        __rucio_local_path = '/project/lgrandi/rucio/'
+        print('You specified _auto_append_rucio_local=True and you are not on dali compute nodes,'
+              'so we will add the following rucio local path: ', __rucio_local_path)
+
+    return _include_rucio_local, __rucio_local_path
+
+
 @deprecate_kwarg('_minimum_run_number', 'minimum_run_number')
 @deprecate_kwarg('_maximum_run_number', 'maximum_run_number')
 @deprecate_kwarg('_include_rucio_remote', 'include_rucio_remote')
@@ -125,6 +155,7 @@ def xenonnt_online(output_folder: str = './strax_data',
 
                    # Frontend options
                    download_heavy: bool = False,
+                   _auto_append_rucio_local: bool = True,
                    _rucio_path: str = '/dali/lgrandi/rucio/',
                    _rucio_local_path: ty.Optional[str] = None,
                    _raw_paths: ty.Optional[str] = ['/dali/lgrandi/xenonnt/raw'],
@@ -156,6 +187,8 @@ def xenonnt_online(output_folder: str = './strax_data',
     :param download_heavy: bool, whether or not to allow downloads of
         heavy data (raw_records*, less the aqmon)
 
+    :param _auto_append_rucio_local: bool, whether or not to automatically append the 
+        rucio local path
     :param _rucio_path: str, path of rucio
     :param _rucio_local_path: str, path of local RSE of rucio. Only use
         for testing!
@@ -179,6 +212,10 @@ def xenonnt_online(output_folder: str = './strax_data',
                  straxen.LEDCalibration,
                  straxen.LEDAfterpulseProcessing])
 
+    if _auto_append_rucio_local:
+        include_rucio_local, _rucio_local_path = find_rucio_local_path(
+            include_rucio_local, _rucio_local_path)
+
     st.storage = [
         straxen.RunDB(
             readonly=not we_are_the_daq,