From 8b0d3cc5d715f460f9e1a1d415b674b466664dd9 Mon Sep 17 00:00:00 2001
From: ddl-rliu <140021987+ddl-rliu@users.noreply.github.com>
Date: Thu, 25 Jul 2024 12:29:38 -0700
Subject: [PATCH] fix: Add CODE_SNIPPET_USE_AUTO_CONFIG environment variable
 (#883)

* Add CODE_SNIPPET_USE_AUTO_CONFIG environment variable

Signed-off-by: ddl-rliu <richard.liu@dominodatalab.com>

* Fix eslint

Signed-off-by: ddl-rliu <richard.liu@dominodatalab.com>

---------

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

diff --git a/packages/common/src/environment/index.ts b/packages/common/src/environment/index.ts
index 82e275c52..78ccdd302 100644
--- a/packages/common/src/environment/index.ts
+++ b/packages/common/src/environment/index.ts
@@ -33,6 +33,12 @@ export interface Env extends NodeJS.ProcessEnv {
    * @example MAINTENANCE_MODE="We are currently down for maintenance.\n\nPlease try again later."
    */
   MAINTENANCE_MODE?: string;
+
+  /**
+   * Controls whether the config object in the inputs/outputs Python code snippet is automatically constructed
+   * Leave unset to use the config object for endpoint window.location.host
+   */
+  CODE_SNIPPET_USE_AUTO_CONFIG?: 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..48fd07283 100644
--- a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeURL.tsx
+++ b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeURL.tsx
@@ -3,6 +3,7 @@ import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
 import { prism } from 'react-syntax-highlighter/dist/esm/styles/prism';
 import FileCopyIcon from '@mui/icons-material/FileCopy';
 import copyToClipboard from 'copy-to-clipboard';
+import { env } from '@clients/common/environment';
 import { errorBackgroundColor } from '@clients/theme/CommonStyles/constants';
 import Button from '@mui/material/Button';
 import Typography from '@mui/material/Typography';
@@ -31,19 +32,19 @@ export const ExecutionNodeURL: React.FC<{
   const isHttps = /^https:/.test(window.location.href);
   const ref = React.useRef<HTMLDivElement>(null);
 
-  const code = isHttps
-    ? // https snippet
-      `from flytekit.remote.remote import FlyteRemote
+  const config =
+    // eslint-disable-next-line no-nested-ternary
+    env.CODE_SNIPPET_USE_AUTO_CONFIG === "true"
+      ? 'Config.auto()'
+      : isHttps
+      ? // https snippet
+        `Config.for_endpoint("${window.location.host}")`
+      : // http snippet
+        `Config.for_endpoint("${window.location.host}", True)`;
+    const code = `from flytekit.remote.remote import FlyteRemote
 from flytekit.configuration import Config
 remote = FlyteRemote(
-    Config.for_endpoint("${window.location.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},
 )
 remote.get("${dataSourceURI}")`;
 
diff --git a/website/console/env/index.ts b/website/console/env/index.ts
index 4114cb9a8..16772c23a 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 whether the config object in the inputs/outputs Python code snippet is automatically constructed
+ */
+const CODE_SNIPPET_USE_AUTO_CONFIG = process.env.CODE_SNIPPET_USE_AUTO_CONFIG || '';
+
 const processEnv = {
   NODE_ENV,
   PORT,
@@ -86,6 +91,7 @@ const processEnv = {
   BASE_HREF,
   DISABLE_CONSOLE_ROUTE_PREFIX,
   MAINTENANCE_MODE,
+  CODE_SNIPPET_USE_AUTO_CONFIG,
 };
 
 export {
@@ -101,5 +107,6 @@ export {
   ADMIN_API,
   LOCAL_DEV_HOST,
   MAINTENANCE_MODE,
+  CODE_SNIPPET_USE_AUTO_CONFIG,
   processEnv,
 };