From a4d8375c42fb61498dcb73a63eed0a14559bd7bb Mon Sep 17 00:00:00 2001
From: Kevin Smith <Kevin.Smith@octobergray.com>
Date: Wed, 21 Aug 2024 10:35:34 -0500
Subject: [PATCH] Add ability to disable Fusion commands

---
 singlestoredb/fusion/handler.py      |  7 +++++++
 singlestoredb/fusion/handlers/job.py | 16 ++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/singlestoredb/fusion/handler.py b/singlestoredb/fusion/handler.py
index b61ac18c..03186532 100644
--- a/singlestoredb/fusion/handler.py
+++ b/singlestoredb/fusion/handler.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import abc
 import functools
+import os
 import re
 import sys
 import textwrap
@@ -543,6 +544,7 @@ class SQLHandler(NodeVisitor):
 
     _grammar: str = CORE_GRAMMAR
     _is_compiled: bool = False
+    _enabled: bool = True
 
     def __init__(self, connection: Connection):
         self.connection = connection
@@ -581,6 +583,11 @@ def register(cls, overwrite: bool = False) -> None:
             Overwrite an existing command with the same name?
 
         """
+        if not cls._enabled and \
+                os.environ.get('SINGLESTOREDB_FUSION_ENABLE_HIDDEN', '0').lower() not in \
+                ['1', 't', 'true', 'y', 'yes']:
+            return
+
         from . import registry
         cls.compile()
         registry.register_handler(cls, overwrite=overwrite)
diff --git a/singlestoredb/fusion/handlers/job.py b/singlestoredb/fusion/handlers/job.py
index 8098fc5b..78c79182 100644
--- a/singlestoredb/fusion/handlers/job.py
+++ b/singlestoredb/fusion/handlers/job.py
@@ -122,6 +122,8 @@ class ScheduleJobHandler(SQLHandler):
         ;
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('JobID', result.STRING)
@@ -222,6 +224,8 @@ class RunJobHandler(SQLHandler):
 
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('JobID', result.STRING)
@@ -284,6 +288,8 @@ class WaitOnJobsHandler(SQLHandler):
 
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('Success', result.BOOL)
@@ -353,6 +359,8 @@ class ShowJobsHandler(SQLHandler):
 
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('JobID', result.STRING)
@@ -484,6 +492,8 @@ class ShowJobExecutionsHandler(SQLHandler):
           EXTENDED;
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('ExecutionID', result.STRING)
@@ -554,6 +564,8 @@ class ShowJobParametersHandler(SQLHandler):
         SHOW JOB PARAMETERS FOR 'job1';
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('Name', result.STRING)
@@ -594,6 +606,8 @@ class ShowJobRuntimesHandler(SQLHandler):
         SHOW JOB RUNTIMES;
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('Name', result.STRING)
@@ -639,6 +653,8 @@ class DropJobHandler(SQLHandler):
         DROP JOBS 'job1', 'job2';
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         res = FusionSQLResult()
         res.add_field('JobID', result.STRING)