UI: Preserve raw connection extra JSON

airflow-core/newsfragments/69741.bugfix.rst

@@ -0,0 +1 @@
+Preserve raw connection extra JSON when editing provider-specific connection fields in the UI.

airflow-core/src/airflow/ui/src/components/FlexibleForm/FlexibleForm.tsx

@@ -77,6 +77,7 @@ export type FlexibleFormProps = {
   readonly key?: string;
   readonly namespace?: string;
   readonly noAccordion?: boolean;
+  readonly preserveRawConf?: boolean;
   readonly setError: (error: boolean) => void;
   readonly subHeader?: string;
 };
@@ -89,31 +90,55 @@ export const FlexibleForm = ({
   isHITL,
   namespace = "default",
   noAccordion,
+  preserveRawConf = false,
   setError,
   subHeader,
 }: FlexibleFormProps) => {
-  const { paramsDict: params, setDisabled, setInitialParamDict, setParamsDict } = useParamStore(namespace);
+  const {
+    initParamsDictFromConf,
+    paramsDict: params,
+    setDisabled,
+    setInitialParamDict,
+    setMergeParamUpdatesIntoConf,
+    setParamsDict,
+  } = useParamStore(namespace);
   const processedSections = new Map();
   const [sectionError, setSectionError] = useState<Map<string, boolean>>(new Map());
   const [hasValidationError, setHasValidationError] = useState(false);
 
+  useEffect(() => {
+    setMergeParamUpdatesIntoConf(preserveRawConf);
+  }, [preserveRawConf, setMergeParamUpdatesIntoConf]);
+
   useEffect(() => {
     // Initialize paramsDict and initialParamDict when modal opens
     if (Object.keys(initialParamsDict.paramsDict).length > 0 && Object.keys(params).length === 0) {
       const paramsCopy = structuredClone(initialParamsDict.paramsDict);
 
-      setParamsDict(paramsCopy);
-      setInitialParamDict(initialParamsDict.paramsDict);
+      if (preserveRawConf) {
+        initParamsDictFromConf(paramsCopy);
+      } else {
+        setParamsDict(paramsCopy);
+        setInitialParamDict(initialParamsDict.paramsDict);
+      }
     }
-  }, [initialParamsDict, params, setParamsDict, setInitialParamDict]);
+  }, [
+    initParamsDictFromConf,
+    initialParamsDict,
+    params,
+    preserveRawConf,
+    setParamsDict,
+    setInitialParamDict,
+  ]);
 
   useEffect(
     () => () => {
       // Clear paramsDict and initialParamDict when the component is unmounted or modal closes
       setParamsDict({});
       setInitialParamDict({});
+      setMergeParamUpdatesIntoConf(false);
     },
-    [setParamsDict, setInitialParamDict],
+    [setParamsDict, setInitialParamDict, setMergeParamUpdatesIntoConf],
   );
 
   useEffect(() => {

airflow-core/src/airflow/ui/src/pages/Connections/ConnectionForm.test.tsx

@@ -0,0 +1,145 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { render, screen, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+
+import type { ParamsSpec } from "src/queries/useDagParams";
+import { Wrapper } from "src/utils/Wrapper";
+
+import ConnectionForm from "./ConnectionForm";
+import type { ConnectionBody } from "./Connections";
+
+const { snowflakeExtraFields } = vi.hoisted(() => ({
+  snowflakeExtraFields: {
+    account: {
+      description: null,
+      schema: {
+        const: undefined,
+        description_md: undefined,
+        enum: undefined,
+        examples: undefined,
+        format: undefined,
+        items: undefined,
+        maximum: undefined,
+        maxLength: undefined,
+        minimum: undefined,
+        minLength: undefined,
+        section: undefined,
+        title: "Account",
+        type: ["string", "null"],
+        values_display: undefined,
+      },
+      value: undefined,
+    },
+    insecure_mode: {
+      description: "Turns off OCSP certificate checks",
+      schema: {
+        const: undefined,
+        description_md: undefined,
+        enum: undefined,
+        examples: undefined,
+        format: undefined,
+        items: undefined,
+        maximum: undefined,
+        maxLength: undefined,
+        minimum: undefined,
+        minLength: undefined,
+        section: undefined,
+        title: "Insecure Mode",
+        type: ["boolean", "null"],
+        values_display: undefined,
+      },
+      value: false,
+    },
+  } satisfies ParamsSpec,
+}));
+
+vi.mock("src/components/JsonEditor", () => ({
+  JsonEditor: ({
+    onBlur,
+    onChange,
+    value,
+  }: {
+    readonly onBlur?: () => void;
+    readonly onChange?: (value: string) => void;
+    readonly value?: string;
+  }) => (
+    <textarea
+      aria-label="extra-json"
+      onBlur={onBlur}
+      onChange={(event) => onChange?.(event.target.value)}
+      value={value ?? ""}
+    />
+  ),
+}));
+
+vi.mock("src/queries/useConfig.tsx", () => ({
+  useConfig: () => false,
+}));
+
+vi.mock("src/queries/useConnectionTypeMeta", () => ({
+  useConnectionTypeMeta: () => ({
+    formattedData: {
+      snowflake: {
+        connection_type: "snowflake",
+        default_conn_name: undefined,
+        extra_fields: snowflakeExtraFields,
+        hook_class_name: "airflow.providers.snowflake.hooks.snowflake.SnowflakeHook",
+        hook_name: "Snowflake",
+        standard_fields: {},
+      },
+    },
+    hookNames: { snowflake: "Snowflake" },
+    isPending: false,
+    keysList: ["snowflake"],
+  }),
+}));
+
+const initialConnection: ConnectionBody = {
+  conn_type: "snowflake",
+  connection_id: "snowflake_default",
+  description: "",
+  extra: '{"account":"1234"}',
+  host: "",
+  login: "",
+  password: "",
+  port: "",
+  schema: "",
+  team_name: "",
+};
+
+describe("ConnectionForm", () => {
+  it("does not rewrite raw Snowflake extras with untouched provider defaults", async () => {
+    render(
+      <ConnectionForm
+        error={undefined}
+        initialConnection={initialConnection}
+        isEditMode
+        isPending={false}
+        mutateConnection={vi.fn()}
+      />,
+      { wrapper: Wrapper },
+    );
+
+    const extraJson = await screen.findByLabelText("extra-json");
+
+    await waitFor(() => expect((extraJson as HTMLTextAreaElement).value).toBe('{\n  "account": "1234"\n}'));
+    expect((extraJson as HTMLTextAreaElement).value).not.toContain("insecure_mode");
+  });
+});

airflow-core/src/airflow/ui/src/pages/Connections/ConnectionForm.tsx

@@ -35,6 +35,8 @@ import { useParamStore } from "src/queries/useParamStore";
 import StandardFields from "./ConnectionStandardFields";
 import type { ConnectionBody } from "./Connections";
 
+const CONNECTION_EXTRA_NAMESPACE = "connection-extra";
+
 type AddConnectionFormProps = {
   readonly error: unknown;
   readonly initialConnection: ConnectionBody;
@@ -57,7 +59,7 @@ const ConnectionForm = ({
     isPending: isMetaPending,
     keysList: connectionTypes,
   } = useConnectionTypeMeta();
-  const { conf: extra, setConf } = useParamStore();
+  const { conf: extra, setConf } = useParamStore(CONNECTION_EXTRA_NAMESPACE);
   const {
     control,
     formState: { isDirty, isValid },
@@ -223,6 +225,8 @@ const ConnectionForm = ({
               flexibleFormDefaultSection={translate("connections.form.extraFields")}
               initialParamsDict={paramsDic}
               key={selectedConnType}
+              namespace={CONNECTION_EXTRA_NAMESPACE}
+              preserveRawConf
               setError={setFormErrors}
               subHeader={isEditMode ? translate("connections.form.helperTextForRedactedFields") : undefined}
             />

airflow-core/src/airflow/ui/src/queries/useParamStore.test.ts

@@ -0,0 +1,155 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { act, renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+
+import type { ParamsSpec, ParamSpec } from "src/queries/useDagParams";
+
+import { useParamStore } from "./useParamStore";
+
+const baseSchema = {
+  const: undefined,
+  description_md: undefined,
+  enum: undefined,
+  examples: undefined,
+  format: undefined,
+  items: undefined,
+  maximum: undefined,
+  maxLength: undefined,
+  minimum: undefined,
+  minLength: undefined,
+  section: undefined,
+  values_display: undefined,
+};
+
+let namespaceId = 0;
+
+const getStore = () => {
+  const namespace = `use-param-store-test-${namespaceId}`;
+
+  namespaceId += 1;
+
+  return renderHook(() => useParamStore(namespace));
+};
+
+const createParam = (title: string, value: unknown, type: ParamSpec["schema"]["type"]): ParamSpec => ({
+  description: null,
+  schema: {
+    ...baseSchema,
+    title,
+    type,
+  },
+  value,
+});
+
+const createSnowflakeParams = (): ParamsSpec => ({
+  account: createParam("Account", undefined, ["string", "null"]),
+  insecure_mode: createParam("Insecure Mode", false, ["boolean", "null"]),
+  warehouse: createParam("Warehouse", undefined, ["string", "null"]),
+});
+
+const setParamValue = (params: ParamsSpec, key: string, value: unknown) => {
+  const param = params[key];
+
+  if (param === undefined) {
+    throw new Error(`Missing test param: ${key}`);
+  }
+
+  param.value = value;
+};
+
+describe("useParamStore", () => {
+  it("keeps the default params-to-conf serialization behavior", () => {
+    const { result } = getStore();
+    const params = createSnowflakeParams();
+
+    setParamValue(params, "account", "1234");
+
+    act(() => result.current.setParamsDict(params));
+
+    expect(JSON.parse(result.current.conf) as Record<string, unknown>).toEqual({
+      account: "1234",
+      insecure_mode: false,
+    });
+  });
+
+  it("initializes provider params from raw conf without injecting provider defaults", () => {
+    const { result } = getStore();
+
+    act(() => {
+      result.current.setMergeParamUpdatesIntoConf(true);
+      result.current.setConf('{"account":"1234","unknown":"keep"}');
+      result.current.initParamsDictFromConf(createSnowflakeParams());
+    });
+
+    expect(result.current.paramsDict.account?.value).toBe("1234");
+    expect(result.current.paramsDict.insecure_mode?.value).toBeUndefined();
+    expect(result.current.paramsDict.unknown).toBeUndefined();
+    expect(JSON.parse(result.current.conf) as Record<string, unknown>).toEqual({
+      account: "1234",
+      unknown: "keep",
+    });
+  });
+
+  it("merges changed provider fields into raw conf while preserving unknown keys", () => {
+    const { result } = getStore();
+
+    act(() => {
+      result.current.setMergeParamUpdatesIntoConf(true);
+      result.current.setConf('{"account":"1234","unknown":"keep"}');
+      result.current.initParamsDictFromConf(createSnowflakeParams());
+    });
+
+    const nextParams = structuredClone(result.current.paramsDict);
+
+    setParamValue(nextParams, "warehouse", "test_warehouse");
+
+    act(() => result.current.setParamsDict(nextParams));
+
+    const parsedConf = JSON.parse(result.current.conf) as Record<string, unknown>;
+
+    expect(parsedConf).toEqual({
+      account: "1234",
+      unknown: "keep",
+      warehouse: "test_warehouse",
+    });
+    expect(parsedConf).not.toHaveProperty("insecure_mode");
+  });
+
+  it("removes cleared provider fields from raw conf without dropping unrelated keys", () => {
+    const { result } = getStore();
+
+    act(() => {
+      result.current.setMergeParamUpdatesIntoConf(true);
+      result.current.setConf('{"account":"1234","unknown":"keep","warehouse":"test_warehouse"}');
+      result.current.initParamsDictFromConf(createSnowflakeParams());
+    });
+
+    const nextParams = structuredClone(result.current.paramsDict);
+
+    setParamValue(nextParams, "account", null);
+
+    act(() => result.current.setParamsDict(nextParams));
+
+    expect(JSON.parse(result.current.conf) as Record<string, unknown>).toEqual({
+      unknown: "keep",
+      warehouse: "test_warehouse",
+    });
+  });
+});