From 4aa12eddbdc67e4880ac5eda3b3664cb855bacfe Mon Sep 17 00:00:00 2001
From: ddl-rliu <richard.liu@dominodatalab.com>
Date: Wed, 17 Jul 2024 12:28:02 -0700
Subject: [PATCH] Add CODE_SNIPPET_URL to config

Signed-off-by: ddl-rliu <richard.liu@dominodatalab.com>
---
 packages/common/src/environment/index.ts                  | 5 +++++
 .../Executions/ExecutionDetails/ExecutionNodeURL.tsx      | 8 +++++---
 website/console/env/index.ts                              | 7 +++++++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/packages/common/src/environment/index.ts b/packages/common/src/environment/index.ts
index 82e275c52..650f19a3f 100644
--- a/packages/common/src/environment/index.ts
+++ b/packages/common/src/environment/index.ts
@@ -33,6 +33,11 @@ export interface Env extends NodeJS.ProcessEnv {
    * @example MAINTENANCE_MODE="We are currently down for maintenance.\n\nPlease try again later."
    */
   MAINTENANCE_MODE?: string;
+
+  /**
+   * Controls the endpoint used in Python code snippet, leave unset to use the default window.location.host
+   */
+  CODE_SNIPPET_URL?: string;
 }
 
 /** Represents a plain object where string keys map to values of the same type */
diff --git a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeURL.tsx b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeURL.tsx
index 8cda4de12..e2408c949 100644
--- a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeURL.tsx
+++ b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeURL.tsx
@@ -11,6 +11,7 @@ import Grid from '@mui/material/Grid';
 import Link from '@mui/material/Link';
 import { RowExpander } from '../Tables/RowExpander';
 import { ScrollableMonospaceText } from '../../common/ScrollableMonospaceText';
+import { env } from '@clients/common/environment';
 
 const StyledScrollableMonospaceText = styled(ScrollableMonospaceText)(({ theme }) => ({
   '&>div': {
@@ -28,7 +29,8 @@ export const ExecutionNodeURL: React.FC<{
   copyUrlText: string;
 }> = ({ dataSourceURI, copyUrlText }) => {
   const [expanded, setExpanded] = React.useState<boolean>(false);
-  const isHttps = /^https:/.test(window.location.href);
+  const codeSnippetUrl = new URL(env.CODE_SNIPPET_URL || window.location.href);
+  const isHttps = /^https:/.test(codeSnippetUrl.protocol);
   const ref = React.useRef<HTMLDivElement>(null);
 
   const code = isHttps
@@ -36,14 +38,14 @@ export const ExecutionNodeURL: React.FC<{
       `from flytekit.remote.remote import FlyteRemote
 from flytekit.configuration import Config
 remote = FlyteRemote(
-    Config.for_endpoint("${window.location.host}"),
+    Config.for_endpoint("${codeSnippetUrl.host}"),
 )
 remote.get("${dataSourceURI}")`
     : // http snippet
       `from flytekit.remote.remote import FlyteRemote
 from flytekit.configuration import Config
 remote = FlyteRemote(
-    Config.for_endpoint("${window.location.host}", True),
+    Config.for_endpoint("${codeSnippetUrl.host}", True),
 )
 remote.get("${dataSourceURI}")`;
 
diff --git a/website/console/env/index.ts b/website/console/env/index.ts
index 4114cb9a8..0c884f9cf 100644
--- a/website/console/env/index.ts
+++ b/website/console/env/index.ts
@@ -77,6 +77,11 @@ const ASSETS_PATH = `${BASE_URL}/assets/`;
  */
 const MAINTENANCE_MODE = process.env.MAINTENANCE_MODE || '';
 
+/**
+ * Controls the endpoint used in Python code snippet
+ */
+const CODE_SNIPPET_URL = process.env.CODE_SNIPPET_URL || '';
+
 const processEnv = {
   NODE_ENV,
   PORT,
@@ -86,6 +91,7 @@ const processEnv = {
   BASE_HREF,
   DISABLE_CONSOLE_ROUTE_PREFIX,
   MAINTENANCE_MODE,
+  CODE_SNIPPET_URL,
 };
 
 export {
@@ -101,5 +107,6 @@ export {
   ADMIN_API,
   LOCAL_DEV_HOST,
   MAINTENANCE_MODE,
+  CODE_SNIPPET_URL,
   processEnv,
 };