浏览代码

[api] Add error handling for missing connector in notebook setup install_example API (#4120)

Wraps connector fetching logic in try-except block to gracefully handle cases where the connector doesn't exist or when queries fail. This prevents the API from sending 500 response when setting up notebook examples with invalid connector IDs.

Adds corresponding test cases to verify both scenarios:
- When connector lookup returns None
- When connector lookup raises an exception
Harsh Gupta 7 月之前
父节点
当前提交
8c522a8bf7
共有 2 个文件被更改,包括 17 次插入2 次删除
  1. 5 1
      desktop/core/src/desktop/api2.py
  2. 12 1
      desktop/core/src/desktop/api2_tests.py

+ 5 - 1
desktop/core/src/desktop/api2.py

@@ -1489,7 +1489,11 @@ def _setup_oozie_examples(request):
 
 
 def _setup_notebook_examples(request):
-  connector = Connector.objects.get(id=request.POST.get('connector_id'))
+  try:
+    connector = Connector.objects.get(id=request.POST.get('connector_id'))
+  except Exception as e:
+    LOG.error(f'Error getting connector: {e}')
+    connector = None
 
   if connector:
     dialect = connector.dialect

+ 12 - 1
desktop/core/src/desktop/api2_tests.py

@@ -737,7 +737,6 @@ class TestDocumentGist(object):
 
 
 class TestAvailableAppExamplesAPI:
-
   # Using custom MockApp instead of Mock to avoid conflicts with Mock's built in 'name' attribute.
   @dataclass
   class MockApp:
@@ -895,6 +894,18 @@ class TestInstallAppExampleAPI:
           assert not mock_beeswax_install_command.called
           mock_notebook_setup_command.assert_called_once_with(dialect='spark', user=request.user)
 
+  def test_setup_notebook_examples_connector_exception(self):
+    with patch('desktop.api2.Connector.objects.get') as mock_get_connector:
+      with patch('desktop.api2.beeswax_install_examples.Command.handle') as mock_beeswax_install_command:
+        with patch('desktop.api2.notebook_setup.Command.handle') as mock_notebook_setup_command:
+          request = Mock(method='POST', POST={'app_name': 'notebook', 'dialect': 'spark'}, user=Mock())
+          mock_get_connector.side_effect = Exception('Connector matching query does not exist.')
+
+          _setup_notebook_examples(request)
+
+          assert not mock_beeswax_install_command.called
+          mock_notebook_setup_command.assert_called_once_with(dialect='spark', user=request.user)
+
   def test_setup_search_examples_with_log_analytics_demo(self):
     with patch('desktop.api2.search_setup.Command.handle') as mock_search_setup_command:
       with patch('desktop.api2.indexer_setup.Command.handle') as mock_indexer_setup_command: