From 317fa4ede79f0669289d8158d504e75eb28c6e5b Mon Sep 17 00:00:00 2001
From: Hai Nguyen <quanghai.ng1512@gmail.com>
Date: Mon, 2 Jan 2023 22:16:13 +0700
Subject: [PATCH] feat: Make UI runnable behind proxy

Signed-off-by: Hai Nguyen <quanghai.ng1512@gmail.com>
---
 sdk/python/feast/cli.py           | 17 +++++++++++++++--
 sdk/python/feast/feature_store.py |  8 +++++++-
 sdk/python/feast/ui_server.py     | 12 ++++++++++--
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
index f677584fe5..c978ca1023 100644
--- a/sdk/python/feast/cli.py
+++ b/sdk/python/feast/cli.py
@@ -151,12 +151,24 @@ def version():
     "--registry_ttl_sec",
     "-r",
     help="Number of seconds after which the registry is refreshed",
-    type=int,
+    type=click.INT,
     default=5,
     show_default=True,
 )
+@click.option(
+    "--root_path",
+    help="Provide root path to make the UI working behind proxy",
+    type=click.STRING,
+    default="",
+)
 @click.pass_context
-def ui(ctx: click.Context, host: str, port: int, registry_ttl_sec: int):
+def ui(
+    ctx: click.Context,
+    host: str,
+    port: int,
+    registry_ttl_sec: int,
+    root_path: str = "",
+):
     """
     Shows the Feast UI over the current directory
     """
@@ -170,6 +182,7 @@ def ui(ctx: click.Context, host: str, port: int, registry_ttl_sec: int):
         port=port,
         get_registry_dump=registry_dump,
         registry_ttl_sec=registry_ttl_sec,
+        root_path=root_path,
     )
 
 
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
index c47b06c047..d3656e1355 100644
--- a/sdk/python/feast/feature_store.py
+++ b/sdk/python/feast/feature_store.py
@@ -2319,7 +2319,12 @@ def get_feature_server_endpoint(self) -> Optional[str]:
 
     @log_exceptions_and_usage
     def serve_ui(
-        self, host: str, port: int, get_registry_dump: Callable, registry_ttl_sec: int
+        self,
+        host: str,
+        port: int,
+        get_registry_dump: Callable,
+        registry_ttl_sec: int,
+        root_path: Optional[str] = "",
     ) -> None:
         """Start the UI server locally"""
         if flags_helper.is_test():
@@ -2335,6 +2340,7 @@ def serve_ui(
             get_registry_dump=get_registry_dump,
             project_id=self.config.project,
             registry_ttl_sec=registry_ttl_sec,
+            root_path=root_path,
         )
 
     @log_exceptions_and_usage
diff --git a/sdk/python/feast/ui_server.py b/sdk/python/feast/ui_server.py
index f79030e8d3..d69c133289 100644
--- a/sdk/python/feast/ui_server.py
+++ b/sdk/python/feast/ui_server.py
@@ -101,6 +101,14 @@ def start_server(
     get_registry_dump: Callable,
     project_id: str,
     registry_ttl_sec: int,
+    root_path: Optional[str] = "",
 ):
-    app = get_app(store, get_registry_dump, project_id, registry_ttl_sec, host, port)
-    uvicorn.run(app, host=host, port=port)
+    app = get_app(
+        store,
+        get_registry_dump,
+        project_id,
+        registry_ttl_sec,
+        host,
+        port,
+    )
+    uvicorn.run(app, host=host, port=port, root_path=root_path)