From ec1091f2ee76a5b6746be2d127b5f333d3089173 Mon Sep 17 00:00:00 2001
From: Sergey Fedoseev <fedoseev.sergey@quiltdata.io>
Date: Thu, 15 Aug 2024 15:18:40 +0200
Subject: [PATCH] Fix SSO config removal in quilt3.admin.sso_config.set()
 (#4098)

---
 .../quilt3/admin/_graphql_client/client.py      | 10 ++++++----
 .../admin/_graphql_client/sso_config_set.py     | 17 +++++++++++------
 api/python/quilt3/admin/sso_config.py           |  7 ++++---
 api/python/tests/test_admin_api.py              |  1 +
 docs/CHANGELOG.md                               |  2 +-
 docs/api-reference/Admin.md                     |  4 ++--
 shared/graphql/schema.graphql                   |  2 +-
 7 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/api/python/quilt3/admin/_graphql_client/client.py b/api/python/quilt3/admin/_graphql_client/client.py
index 6218be65f9c..3a77062d12f 100644
--- a/api/python/quilt3/admin/_graphql_client/client.py
+++ b/api/python/quilt3/admin/_graphql_client/client.py
@@ -924,10 +924,12 @@ def sso_config_get(self, **kwargs: Any) -> Optional[SsoConfigGetAdminSsoConfig]:
 
     def sso_config_set(
         self, config: Union[Optional[str], UnsetType] = UNSET, **kwargs: Any
-    ) -> Union[
-        SsoConfigSetAdminSetSsoConfigSsoConfig,
-        SsoConfigSetAdminSetSsoConfigInvalidInput,
-        SsoConfigSetAdminSetSsoConfigOperationError,
+    ) -> Optional[
+        Union[
+            SsoConfigSetAdminSetSsoConfigSsoConfig,
+            SsoConfigSetAdminSetSsoConfigInvalidInput,
+            SsoConfigSetAdminSetSsoConfigOperationError,
+        ]
     ]:
         query = gql(
             """
diff --git a/api/python/quilt3/admin/_graphql_client/sso_config_set.py b/api/python/quilt3/admin/_graphql_client/sso_config_set.py
index e66706fc267..af7e32af9fe 100644
--- a/api/python/quilt3/admin/_graphql_client/sso_config_set.py
+++ b/api/python/quilt3/admin/_graphql_client/sso_config_set.py
@@ -1,7 +1,7 @@
 # Generated by ariadne-codegen
 # Source: queries.graphql
 
-from typing import Literal, Union
+from typing import Annotated, Literal, Optional, Union
 
 from pydantic import Field
 
@@ -18,11 +18,16 @@ class SsoConfigSet(BaseModel):
 
 
 class SsoConfigSetAdmin(BaseModel):
-    set_sso_config: Union[
-        "SsoConfigSetAdminSetSsoConfigSsoConfig",
-        "SsoConfigSetAdminSetSsoConfigInvalidInput",
-        "SsoConfigSetAdminSetSsoConfigOperationError",
-    ] = Field(alias="setSsoConfig", discriminator="typename__")
+    set_sso_config: Optional[
+        Annotated[
+            Union[
+                "SsoConfigSetAdminSetSsoConfigSsoConfig",
+                "SsoConfigSetAdminSetSsoConfigInvalidInput",
+                "SsoConfigSetAdminSetSsoConfigOperationError",
+            ],
+            Field(discriminator="typename__"),
+        ]
+    ] = Field(alias="setSsoConfig")
 
 
 class SsoConfigSetAdminSetSsoConfigSsoConfig(SsoConfigSelection):
diff --git a/api/python/quilt3/admin/sso_config.py b/api/python/quilt3/admin/sso_config.py
index cf4f5c448be..9e2ef36314f 100644
--- a/api/python/quilt3/admin/sso_config.py
+++ b/api/python/quilt3/admin/sso_config.py
@@ -11,8 +11,9 @@ def get() -> T.Optional[types.SSOConfig]:
     return None if result is None else types.SSOConfig(**result.model_dump())
 
 
-def set(config: T.Optional[str]) -> types.SSOConfig:
+def set(config: T.Optional[str]) -> T.Optional[types.SSOConfig]:
     """
-    Set the SSO configuration.
+    Set the SSO configuration. Pass `None` to remove SSO configuration.
     """
-    return types.SSOConfig(**util.handle_errors(util.get_client().sso_config_set(config)).model_dump())
+    result = util.get_client().sso_config_set(config)
+    return None if result is None else types.SSOConfig(**util.handle_errors(result).model_dump())
diff --git a/api/python/tests/test_admin_api.py b/api/python/tests/test_admin_api.py
index 87cd44a5259..41c845dfd67 100644
--- a/api/python/tests/test_admin_api.py
+++ b/api/python/tests/test_admin_api.py
@@ -348,6 +348,7 @@ def test_sso_config_get(data, result):
     "data,result",
     [
         (SSO_CONFIG, admin.SSOConfig(**_as_dataclass_kwargs(SSO_CONFIG))),
+        (None, None),
         *MUTATION_ERRORS,
     ],
 )
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index d192e4a71e2..8dfd30f09a7 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -18,7 +18,7 @@ Entries inside each section should be ordered by type:
 # unreleased - YYYY-MM-DD
 ## Python API
 
-* [Added] New `quilt3.admin.sso_config` sub-module for management of SSO configuration ([#4065](https://github.com/quiltdata/quilt/pull/4065))
+* [Added] New `quilt3.admin.sso_config` sub-module for management of SSO configuration ([#4065](https://github.com/quiltdata/quilt/pull/4065), [#4098](https://github.com/quiltdata/quilt/pull/4098))
 
 ## CLI
 
diff --git a/docs/api-reference/Admin.md b/docs/api-reference/Admin.md
index 2d6b08cc194..63724f08de9 100644
--- a/docs/api-reference/Admin.md
+++ b/docs/api-reference/Admin.md
@@ -140,7 +140,7 @@ __Arguments__
 Get the current SSO configuration.
 
 
-## set(config: Optional[str]) -> quilt3.admin.types.SSOConfig  {#set}
+## set(config: Optional[str]) -> Optional[quilt3.admin.types.SSOConfig]  {#set}
 
-Set the SSO configuration.
+Set the SSO configuration. Pass `None` to remove SSO configuration.
 
diff --git a/shared/graphql/schema.graphql b/shared/graphql/schema.graphql
index c6d5ade223a..41ebe809fe5 100644
--- a/shared/graphql/schema.graphql
+++ b/shared/graphql/schema.graphql
@@ -884,7 +884,7 @@ union SetSsoConfigResult = SsoConfig | InvalidInput | OperationError
 
 type AdminMutations {
   user: UserAdminMutations!
-  setSsoConfig(config: String): SetSsoConfigResult!
+  setSsoConfig(config: String): SetSsoConfigResult
 }
 
 type Mutation {