From d5be1bf7476e1a6aef39635364034f3aa9f08a76 Mon Sep 17 00:00:00 2001
From: guy-har <60321938+guy-har@users.noreply.github.com>
Date: Thu, 4 Apr 2024 09:26:36 +0300
Subject: [PATCH] Feature/python wrapper sts (#7620)

---
 clients/python-wrapper/lakefs/client.py | 22 +++++++++++++++++++++-
 clients/python-wrapper/requirements.txt |  2 +-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/clients/python-wrapper/lakefs/client.py b/clients/python-wrapper/lakefs/client.py
index e3e00b7401e..2665d19aca7 100644
--- a/clients/python-wrapper/lakefs/client.py
+++ b/clients/python-wrapper/lakefs/client.py
@@ -13,7 +13,7 @@
 from lakefs_sdk.client import LakeFSClient
 
 from lakefs.config import ClientConfig
-from lakefs.exceptions import NotAuthorizedException, ServerException
+from lakefs.exceptions import NotAuthorizedException, ServerException, api_exception_handler
 from lakefs.models import ServerStorageConfiguration
 
 
@@ -106,6 +106,26 @@ def version(self) -> str:
         return self._server_conf.version
 
 
+def from_web_identity(code: str, state: str, redirect_uri: str, ttl_seconds: int = 3600, **kwargs) -> Client:
+    """
+    Authenticate against lakeFS using a code received from an identity provider
+
+    :param code: The code received from the identity provider
+    :param state: The state received from the identity provider
+    :param redirect_uri: The redirect URI used in the authentication process
+    :param ttl_seconds: The token's time-to-live in seconds
+    :param kwargs: Remaining arguments for the Client object
+    :return: The authenticated Client object
+    :raise NotAuthorizedException: if user is not authorized to perform this operation
+    """
+    client = Client(**kwargs)
+    sts_requests = lakefs_sdk.StsAuthRequest(code=code, state=state, redirect_uri=redirect_uri, ttl_seconds=ttl_seconds)
+    with api_exception_handler():
+        auth_token = client.sdk_client.experimental_api.sts_login(sts_requests)
+    client.config.access_token = auth_token.token
+    return client
+
+
 class _BaseLakeFSObject:
     """
     Base class for all lakeFS SDK objects, holds the client object and handles errors where no authentication method
diff --git a/clients/python-wrapper/requirements.txt b/clients/python-wrapper/requirements.txt
index 0f2e8730027..0df424a0e2f 100644
--- a/clients/python-wrapper/requirements.txt
+++ b/clients/python-wrapper/requirements.txt
@@ -1,5 +1,5 @@
 aenum~=3.1.15
-lakefs-sdk==1.10.0
+lakefs-sdk>=1.16.0
 pydantic >= 1.10.5, < 2
 python-dateutil~=2.8.2
 PyYAML~=6.0.1