From 95d5e5f4ba1e4c285b6c91133151a4776fd5a1b0 Mon Sep 17 00:00:00 2001
From: Christine Banek <cbanek@lsst.org>
Date: Wed, 29 Nov 2023 16:41:06 -0700
Subject: [PATCH] [DM-41830] Add some sanity checking on siav2 labels

---
 src/lsst/rsp/__init__.py |  3 +++
 src/lsst/rsp/service.py  | 50 ++++++++++++++++++++++++++++++++++++----
 src/lsst/rsp/utils.py    |  2 +-
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/lsst/rsp/__init__.py b/src/lsst/rsp/__init__.py
index be2e6b9..86ed6d1 100644
--- a/src/lsst/rsp/__init__.py
+++ b/src/lsst/rsp/__init__.py
@@ -12,6 +12,7 @@
 )
 from .forwarder import Forwarder
 from .log import IPythonHandler, forward_lsst_log
+from .service import get_datalink_result, get_siav2_service
 from .utils import (
     format_bytes,
     get_access_token,
@@ -19,6 +20,7 @@
     get_hostname,
     get_node,
     get_pod,
+    get_pyvo_auth,
     show_with_bokeh_server,
 )
 
@@ -47,5 +49,6 @@
     "get_obstap_service",
     "retrieve_query",
     "get_hostname",
+    "get_pyvo_auth",
     "show_with_bokeh_server",
 ]
diff --git a/src/lsst/rsp/service.py b/src/lsst/rsp/service.py
index 348b6c2..9e57e19 100644
--- a/src/lsst/rsp/service.py
+++ b/src/lsst/rsp/service.py
@@ -1,12 +1,52 @@
+import os
+
 import pyvo
 from pyvo.dal import SIA2Service
+from pyvo.dal.adhoc import DatalinkResults
+
+from utils import get_pyvo_auth, get_service_url
+
+LSST_CLOUD = [
+    "https://data.lsst.cloud",
+    "https://data-int.lsst.cloud", 
+    "https://data-dev.lsst.cloud"
+]
+
+USDF = [
+    "https://usdf-rsp.slac.stanford.edu",
+    "https://usdf-rsp-dev.slac.stanford.edu"
+]
+
 
-from .utils import get_pyvo_auth, get_service_url
+def get_datalink_result(result) -> DatalinkResults:
+    """Helper function to return the datalink part of the result."""
+    return DatalinkResults.from_result_url(result.getdataurl(),session=get_pyvo_auth())
+
+
+def get_siav2_service(*args: str) -> SIA2Service:
+    """Return a configured SIA2Service object that is ready to use."""
+    fqdn = os.environ["EXTERNAL_INSTANCE_URL"]
 
-def get_siav2_service(*args: str) -> pyvo.dal.SIA2Service:
     if len(args) == 0:
-        ds = "dp0.2"
-    elif args == "latiss":
-        ds = "latiss"
+        # Use the default for each environment
+        if fqdn in LSST_CLOUD:
+            label = "dp0.2"
+        elif fqdn in USDF:
+            label = "latiss"
+        else:
+            raise Exception("Unknown environment " + fqdn)
+    else:
+        label = args[0]
+
+    # If a label is passed, check that.
+    if label not in ["dp0.2", "latiss"]:
+        raise Exception(label + " is not a valid siav2 label")
+
+    if label == "latiss" and fqdn not in USDF:
+        raise Exception(label + " data not available at your location")
+    if label == "dp0.2" and fqdn not in LSST_CLOUD:
+        raise Exception(label + " data not available at your location")
 
+    # No matter what, we've only got one sia server per environment
+    # so for now just do some checking.
     return SIA2Service(get_service_url("siav2"), get_pyvo_auth())
diff --git a/src/lsst/rsp/utils.py b/src/lsst/rsp/utils.py
index 591300e..3cd29ed 100644
--- a/src/lsst/rsp/utils.py
+++ b/src/lsst/rsp/utils.py
@@ -65,7 +65,7 @@ def get_service_url(name: str, env_name=None) -> str:
         return url
 
     fqdn = os.getenv("EXTERNAL_INSTANCE_URL") or ""
-    path = os.getenv(f"{env_name}_ROUTE") or f"/api/{name}"
+    path = os.getenv(f"{env_name}_ROUTE") or f"api/{name}"
     return f"{fqdn}/{path}"