Add ability to disable Fusion commands

singlestore-labs · Aug 21, 2024 · a4d8375 · a4d8375
1 parent 4fa6cad
commit a4d8375
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git 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
@@ -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)