From 44f2c8decd03d271ca2600c615c0f14f6594dadc Mon Sep 17 00:00:00 2001
From: Martin Bonnin <martin@mbonnin.net>
Date: Wed, 7 Aug 2024 16:36:26 +0200
Subject: [PATCH] add NetworkTransport.closeConnection() (#6105)

---
 .../apollographql/apollo/exception/Exceptions.kt  |  8 +++++++-
 .../websocket/WebSocketNetworkTransport.kt        | 15 ++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/libraries/apollo-api/src/commonMain/kotlin/com/apollographql/apollo/exception/Exceptions.kt b/libraries/apollo-api/src/commonMain/kotlin/com/apollographql/apollo/exception/Exceptions.kt
index bfecbb6475d..9cbe814546e 100644
--- a/libraries/apollo-api/src/commonMain/kotlin/com/apollographql/apollo/exception/Exceptions.kt
+++ b/libraries/apollo-api/src/commonMain/kotlin/com/apollographql/apollo/exception/Exceptions.kt
@@ -78,7 +78,7 @@ class RouterError(
 ) : ApolloException(message = "Router error(s) (first: '${errors.firstOrNull()?.message}')")
 
 /**
- * A WebSocket connection could not be established: e.g., expired token
+ * A WebSocket close frame was received from the server
  */
 class ApolloWebSocketClosedException(
     val code: Int,
@@ -86,6 +86,12 @@ class ApolloWebSocketClosedException(
     cause: Throwable? = null,
 ) : ApolloException(message = "WebSocket Closed code='$code' reason='$reason'", cause = cause)
 
+/**
+ * `closeConnection()` was called to force closing the websocket
+ */
+@ApolloExperimental
+data object ApolloWebSocketForceCloseException : ApolloException(message = "closeConnection() was called")
+
 /**
  * The response was received but the response code was not 200
  *
diff --git a/libraries/apollo-runtime/src/commonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketNetworkTransport.kt b/libraries/apollo-runtime/src/commonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketNetworkTransport.kt
index 2f369a72484..7dec7222547 100644
--- a/libraries/apollo-runtime/src/commonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketNetworkTransport.kt
+++ b/libraries/apollo-runtime/src/commonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketNetworkTransport.kt
@@ -9,6 +9,7 @@ import com.apollographql.apollo.api.json.ApolloJsonElement
 import com.apollographql.apollo.api.json.jsonReader
 import com.apollographql.apollo.api.toApolloResponse
 import com.apollographql.apollo.exception.ApolloException
+import com.apollographql.apollo.exception.ApolloWebSocketForceCloseException
 import com.apollographql.apollo.exception.DefaultApolloException
 import com.apollographql.apollo.exception.SubscriptionOperationException
 import com.apollographql.apollo.internal.DeferredJsonMerger
@@ -279,7 +280,7 @@ private fun Map<String, Any?>.isDeferred(): Boolean {
  *
  * [exception] is passed down to [ApolloResponse.exception] so you can decide how to handle the exception for active subscriptions.
  *
- * ```
+ * ```kotlin
  * apolloClient.subscriptionNetworkTransport.closeConnection(DefaultApolloException("oh no!"))
  * ```
  *
@@ -292,3 +293,15 @@ fun NetworkTransport.closeConnection(exception: ApolloException) {
 
   webSocketNetworkTransport.closeConnection(exception)
 }
+
+/**
+ * Closes the websocket connection if the transport is a [WebSocketNetworkTransport].
+ *
+ * A response is emitted with [ApolloResponse.exception] set to [ApolloWebSocketForceCloseException].
+ */
+@ApolloExperimental
+fun NetworkTransport.closeConnection() {
+  val webSocketNetworkTransport = (this as? WebSocketNetworkTransport) ?: throw IllegalArgumentException("'$this' is not an instance of com.apollographql.apollo.websocket.WebSocketNetworkTransport")
+
+  webSocketNetworkTransport.closeConnection(ApolloWebSocketForceCloseException)
+}