From 2ab6457f0b966283d0280d23f8bfe0d66dc3a0d1 Mon Sep 17 00:00:00 2001
From: kcho <kevincho@bwh.harvard.edu>
Date: Mon, 1 Apr 2024 18:13:33 -0400
Subject: [PATCH] feat: reuse xnat connection

---
 lochness/xnat/__init__.py | 17 ++++++++++++-----
 scripts/sync.py           |  7 ++++++-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/lochness/xnat/__init__.py b/lochness/xnat/__init__.py
index d8caac41..8a7c41bc 100644
--- a/lochness/xnat/__init__.py
+++ b/lochness/xnat/__init__.py
@@ -218,8 +218,17 @@ def clear_stored_credentials(dataorc_bindary_path: Path):
     child.expect(pexpect.EOF, timeout=None)
 
 
+def get_xnat_session(Lochness, subject):
+    alias = subject.xnat.keys()[0]
+    keyring = Lochness['keyring'][alias]
+    xnat_session = xnat.connect(keyring['URL'],
+                                keyring['USERNAME'],
+                                keyring['PASSWORD'])
+    return xnat_session
+
+
 @net.retry(max_attempts=5)
-def sync_xnatpy(Lochness, subject, dry=False):
+def sync_xnatpy(Lochness, subject, dry=False, session=None):
     """A new sync function with XNATpy"""
     logger.debug('exploring {0}/{1}'.format(subject.study, subject.id))
 
@@ -236,10 +245,8 @@ def sync_xnatpy(Lochness, subject, dry=False):
             logger.warning(f'Failed to remove {tmp_file}: {e}')
 
     for alias, xnat_uids in iter(subject.xnat.items()):
-        keyring = Lochness['keyring'][alias]
-        session = xnat.connect(keyring['URL'],
-                               keyring['USERNAME'],
-                               keyring['PASSWORD'])
+        if session is None:
+            session = get_xnat_session(Lochness, subject)
         '''
         pull XNAT data agnostic to the case of subject IDs loop over lower and
         upper case IDs if the data for one ID do not exist, experiments(auth,
diff --git a/scripts/sync.py b/scripts/sync.py
index ac26e20a..a9f574ae 100755
--- a/scripts/sync.py
+++ b/scripts/sync.py
@@ -34,6 +34,7 @@
 from datetime import datetime, date
 from lochness.cleaner import rm_transferred_files_under_phoenix
 from lochness.utils.source_check import check_source
+from lochness.xnat import get_xnat_session
 # import dpanonymize
 
 SOURCES = {
@@ -229,6 +230,7 @@ def do(args, Lochness):
                                      multiple_site, upenn_redcap)
 
     n = 0
+    xnat_session = None
     for subject in lochness.read_phoenix_metadata(Lochness, args.studies):
         if n == 0:
             save_redcap_metadata(Lochness, subject)
@@ -249,8 +251,11 @@ def do(args, Lochness):
         else:
             for source, Module in zip(args.input_sources, args.source):
                 if source == 'xnat':
+                    if xnat_session is None:
+                        xnat_session = get_xnat_session(Lochness, subject)
                     lochness.attempt(Module.sync_xnatpy, Lochness,
-                                     subject, dry=args.dry)
+                                     subject, dry=args.dry,
+                                     session=xnat_session)
                 else:
                     lochness.attempt(Module.sync, Lochness,
                                      subject, dry=args.dry)