feat(server): support session close

resolves #176

fakesnow/server.py

@@ -52,7 +52,7 @@ async def login_request(request: Request) -> JSONResponse:
 
 async def query_request(request: Request) -> JSONResponse:
     try:
-        conn = to_conn(request)
+        conn = to_conn(to_token(request))
 
         body = await request.body()
         if request.headers.get("Content-Encoding") == "gzip":
@@ -107,24 +107,46 @@ async def query_request(request: Request) -> JSONResponse:
         )
 
 
-def to_conn(request: Request) -> FakeSnowflakeConnection:
+def to_token(request: Request) -> str:
     if not (auth := request.headers.get("Authorization")):
-        raise ServerError(status_code=401, code="390103", message="Session token not found in the request data.")
+        raise ServerError(status_code=401, code="390101", message="Authorization header not found in the request data.")
 
-    token = auth[17:-1]
+    return auth[17:-1]
 
+
+def to_conn(token: str) -> FakeSnowflakeConnection:
     if not (conn := sessions.get(token)):
         raise ServerError(status_code=401, code="390104", message="User must login again to access the service.")
 
     return conn
 
 
+async def session(request: Request) -> JSONResponse:
+    try:
+        token = to_token(request)
+        _ = to_conn(token)
+
+        if bool(request.query_params.get("delete")):
+            del sessions[token]
+
+        return JSONResponse(
+            {"data": None, "code": None, "message": None, "success": True},
+        )
+
+    except ServerError as e:
+        return JSONResponse(
+            {"data": None, "code": e.code, "message": e.message, "success": False, "headers": None},
+            status_code=e.status_code,
+        )
+
+
 routes = [
     Route(
         "/session/v1/login-request",
         login_request,
         methods=["POST"],
     ),
+    Route("/session", session, methods=["POST"]),
     Route(
         "/queries/v1/query-request",
         query_request,

tests/test_server.py

@@ -7,6 +7,7 @@
 import pytz
 import requests
 import snowflake.connector
+import snowflake.connector.network
 from dirty_equals import IsUUID
 from snowflake.connector.cursor import ResultMetadata
 
@@ -111,6 +112,21 @@ def test_server_abort_request(server: dict) -> None:
         cur.execute("select 'will abort'")
 
 
+def test_server_close(server: dict) -> None:
+    conn = snowflake.connector.connect(**server)
+
+    # conn.close() ignores errors so we call the endpoint directly
+    assert conn.rest and conn.rest.token
+    response = requests.post(
+        f"http://{server['host']}:{server['port']}/session?delete=true",
+        headers={"Authorization": f'Snowflake Token="{conn.rest.token}"'},
+        timeout=5,
+        json={},
+    )
+    assert response.status_code == 200
+    assert response.json()["success"]
+
+
 def test_server_no_gzip(server: dict) -> None:
     # mimic the go snowflake connector which does not gzip requests
     headers = {