From 95d5e054660c6cfbbfa1fce54137eb31a738b24f Mon Sep 17 00:00:00 2001
From: Advitya Gemawat <agemawat@microsoft.com>
Date: Fri, 19 Jan 2024 10:22:37 -0500
Subject: [PATCH] Private Link URL support on rai_core_flask (#2493)

* pl support ckpt

* decoupled changes

* python lint fixes
---
 .../environments/azure_nb_environment.py             | 12 +++++++++---
 rai_core_flask/rai_core_flask/flask_helper.py        |  4 +++-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/rai_core_flask/rai_core_flask/environments/azure_nb_environment.py b/rai_core_flask/rai_core_flask/environments/azure_nb_environment.py
index 81264a84cd..d2d4ecfab4 100644
--- a/rai_core_flask/rai_core_flask/environments/azure_nb_environment.py
+++ b/rai_core_flask/rai_core_flask/environments/azure_nb_environment.py
@@ -36,14 +36,20 @@ def __init__(self, service):
             else:
                 instance_name = self.nbvm["instance"]
                 domain_suffix = self.nbvm["domainsuffix"]
+                if service.is_private_link:
+                    url_format = f"{instance_name}.{domain_suffix}\
+                        :{service.port}"
+                else:
+                    url_format = f"{instance_name}-{service.port}\
+                        .{domain_suffix}"
                 self.base_url = \
-                    f"https://{instance_name}-{service.port}.{domain_suffix}"
+                    f"https://{url_format}"
                 self.successfully_detected = True
                 self.nbvm_origins = [
                     f"https://{instance_name}.{domain_suffix}",
-                    f"https://{instance_name}-{service.port}.{domain_suffix}",
+                    f"https://{url_format}",
                     f"wss://{instance_name}.{domain_suffix}",
-                    f"wss://{instance_name}-{service.port}.{domain_suffix}",
+                    f"wss://{url_format}",
                     "https://ml.azure.com",
                     "https://dev.ml.azure.com",
                     "https://int.ml.azure.com"
diff --git a/rai_core_flask/rai_core_flask/flask_helper.py b/rai_core_flask/rai_core_flask/flask_helper.py
index 3aef18b6fe..6158901e6f 100644
--- a/rai_core_flask/rai_core_flask/flask_helper.py
+++ b/rai_core_flask/rai_core_flask/flask_helper.py
@@ -24,13 +24,15 @@
 class FlaskHelper(object):
     """FlaskHelper is a class for common Flask utilities used in dashboards."""
 
-    def __init__(self, ip=None, port=None, with_credentials=False):
+    def __init__(self, ip=None, port=None, with_credentials=False,
+                 is_private_link=False):
         # The name passed to Flask needs to be unique per instance.
         self.app = Flask(uuid.uuid4().hex)
 
         self.port = port
         self.ip = ip
         self.with_credentials = with_credentials
+        self.is_private_link = is_private_link
         # dictionary to store arbitrary state for use by consuming classes
         self.shared_state = {}
         if self.ip is None: