From 2bcc1d86fc3a9c1f966dec56f393751cd601a227 Mon Sep 17 00:00:00 2001
From: Levente Balogh <balogh.levente.hu@gmail.com>
Date: Wed, 25 Jun 2025 16:02:35 +0200
Subject: [PATCH] PluginExtensions: fix `setJsonData()` race condition in
 EditDataSource extension point (#107102)

* fix: EditDataSource extension point - setJsonData()

* fix(EditDatasource): make it possible to call `setSecureJsonData()` first
---
 .../components/EditDataSource.test.tsx        | 98 ++++++++++++++++++-
 .../datasources/components/EditDataSource.tsx | 23 +++--
 2 files changed, 113 insertions(+), 8 deletions(-)

diff --git a/public/app/features/datasources/components/EditDataSource.test.tsx b/public/app/features/datasources/components/EditDataSource.test.tsx
index be4779e68bb..6a90180c43b 100644
--- a/public/app/features/datasources/components/EditDataSource.test.tsx
+++ b/public/app/features/datasources/components/EditDataSource.test.tsx
@@ -1,7 +1,8 @@
 import { screen, render } from '@testing-library/react';
+import { useEffect } from 'react';
 import { Provider } from 'react-redux';
 
-import { PluginState } from '@grafana/data';
+import { DataSourceJsonData, PluginExtensionDataSourceConfigContext, PluginState } from '@grafana/data';
 import { setPluginComponentsHook } from '@grafana/runtime';
 import { createComponentWithMeta } from 'app/features/plugins/extensions/usePluginComponents';
 import { configureStore } from 'app/store/configureStore';
@@ -12,6 +13,8 @@ import { missingRightsMessage } from './DataSourceMissingRightsMessage';
 import { readOnlyMessage } from './DataSourceReadOnlyMessage';
 import { EditDataSourceView, ViewProps } from './EditDataSource';
 
+const onOptionsChange = jest.fn();
+
 jest.mock('@grafana/runtime', () => {
   return {
     ...jest.requireActual('@grafana/runtime'),
@@ -38,7 +41,7 @@ const setup = (props?: Partial<ViewProps>) => {
         onDelete={jest.fn()}
         onDefaultChange={jest.fn()}
         onNameChange={jest.fn()}
-        onOptionsChange={jest.fn()}
+        onOptionsChange={onOptionsChange}
         onTest={jest.fn()}
         onUpdate={jest.fn()}
         {...props}
@@ -50,6 +53,7 @@ const setup = (props?: Partial<ViewProps>) => {
 describe('<EditDataSource>', () => {
   beforeEach(() => {
     setPluginComponentsHook(jest.fn().mockReturnValue({ isLoading: false, components: [] }));
+    onOptionsChange.mockClear();
   });
 
   describe('On loading errors', () => {
@@ -359,4 +363,94 @@ describe('<EditDataSource>', () => {
       expect(props.context.testingStatus).toBeDefined();
     });
   });
+
+  it('should be possible to update the `jsonData` first and `secureJsonData` directly afterwards from the extension component', () => {
+    const message = "I'm a UI extension component!";
+    const component = ({ context }: { context: PluginExtensionDataSourceConfigContext }) => {
+      useEffect(() => {
+        context.setJsonData({ test: 'test' } as unknown as DataSourceJsonData);
+        context.setSecureJsonData({ test: 'test' });
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+      }, []);
+
+      return <div>{message}</div>;
+    };
+
+    setPluginComponentsHook(
+      jest.fn().mockReturnValue({
+        isLoading: false,
+        components: [
+          createComponentWithMeta(
+            {
+              pluginId: 'grafana-pdc-app',
+              title: 'Example component',
+              description: 'Example description',
+              component: component as unknown as React.ComponentType<{}>,
+            },
+            '1'
+          ),
+        ],
+      })
+    );
+
+    setup({
+      dataSourceRights: {
+        readOnly: false,
+        hasDeleteRights: true,
+        hasWriteRights: true,
+      },
+    });
+
+    expect(onOptionsChange).toHaveBeenCalledTimes(2);
+    expect(onOptionsChange).toHaveBeenCalledWith({
+      ...getMockDataSource(),
+      jsonData: { ...getMockDataSource().jsonData, test: 'test' },
+      secureJsonData: { test: 'test' },
+    });
+  });
+
+  it('should be possible to update the `secureJsonData` first and `jsonData` directly afterwards from the extension component', () => {
+    const message = "I'm a UI extension component!";
+    const component = ({ context }: { context: PluginExtensionDataSourceConfigContext }) => {
+      useEffect(() => {
+        context.setSecureJsonData({ test: 'test' });
+        context.setJsonData({ test: 'test' } as unknown as DataSourceJsonData);
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+      }, []);
+
+      return <div>{message}</div>;
+    };
+
+    setPluginComponentsHook(
+      jest.fn().mockReturnValue({
+        isLoading: false,
+        components: [
+          createComponentWithMeta(
+            {
+              pluginId: 'grafana-pdc-app',
+              title: 'Example component',
+              description: 'Example description',
+              component: component as unknown as React.ComponentType<{}>,
+            },
+            '1'
+          ),
+        ],
+      })
+    );
+
+    setup({
+      dataSourceRights: {
+        readOnly: false,
+        hasDeleteRights: true,
+        hasWriteRights: true,
+      },
+    });
+
+    expect(onOptionsChange).toHaveBeenCalledTimes(2);
+    expect(onOptionsChange).toHaveBeenCalledWith({
+      ...getMockDataSource(),
+      jsonData: { ...getMockDataSource().jsonData, test: 'test' },
+      secureJsonData: { test: 'test' },
+    });
+  });
 });
diff --git a/public/app/features/datasources/components/EditDataSource.tsx b/public/app/features/datasources/components/EditDataSource.tsx
index e8c56430ec1..5b8f8fc8d36 100644
--- a/public/app/features/datasources/components/EditDataSource.tsx
+++ b/public/app/features/datasources/components/EditDataSource.tsx
@@ -117,6 +117,11 @@ export function EditDataSourceView({
   const { readOnly, hasWriteRights, hasDeleteRights } = dataSourceRights;
   const hasDataSource = dataSource.id > 0;
   const { components, isLoading } = useDataSourceConfigPluginExtensions();
+  // This is a workaround to avoid race-conditions between the `setSecureJsonData()` and `setJsonData()` calls instantiated by the extension components.
+  // Both those exposed functions are calling `onOptionsChange()` with the new jsonData and secureJsonData, and if they are called in the same tick, the Redux store
+  // (which provides the `datasource` object) won't be updated yet, and they override each others `jsonData` value.
+  let currentJsonData = dataSource.jsonData;
+  let currentSecureJsonData = dataSource.secureJsonData;
 
   const dsi = getDataSourceSrv()?.getInstanceSettings(dataSource.uid);
 
@@ -201,16 +206,22 @@ export function EditDataSourceView({
                 dataSource,
                 dataSourceMeta,
                 testingStatus,
-                setJsonData: (jsonData) =>
+                setJsonData: (jsonData) => {
+                  currentJsonData = { ...currentJsonData, ...jsonData };
                   onOptionsChange({
                     ...dataSource,
-                    jsonData: { ...dataSource.jsonData, ...jsonData },
-                  }),
-                setSecureJsonData: (secureJsonData) =>
+                    secureJsonData: { ...currentSecureJsonData },
+                    jsonData: currentJsonData,
+                  });
+                },
+                setSecureJsonData: (secureJsonData) => {
+                  currentSecureJsonData = { ...currentSecureJsonData, ...secureJsonData };
                   onOptionsChange({
                     ...dataSource,
-                    secureJsonData: { ...dataSource.secureJsonData, ...secureJsonData },
-                  }),
+                    jsonData: { ...currentJsonData },
+                    secureJsonData: currentSecureJsonData,
+                  });
+                },
               }}
             />
           </div>