From 6b07e0e8d66e9c0ab7f43f817109ffd94960006f Mon Sep 17 00:00:00 2001
From: Harry <xh001x@hotmail.com>
Date: Mon, 11 Aug 2025 11:25:36 +0800
Subject: [PATCH] feat: add expiration for OAuth credentials in datasource
 provider

---
 .../console/datasets/rag_pipeline/datasource_auth.py          | 2 ++
 api/models/oauth.py                                           | 1 +
 api/services/datasource_provider_service.py                   | 4 ++++
 3 files changed, 7 insertions(+)

diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py
index d67af182cd..192d7ffdf6 100644
--- a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py
+++ b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py
@@ -110,6 +110,7 @@ class DatasourceOAuthCallback(Resource):
                 provider_id=datasource_provider_id,
                 avatar_url=oauth_response.metadata.get("avatar_url") or None,
                 name=oauth_response.metadata.get("name") or None,
+                expire_at=oauth_response.expires_at,
                 credentials=dict(oauth_response.credentials),
                 credential_id=context.get("credential_id"),
             )
@@ -119,6 +120,7 @@ class DatasourceOAuthCallback(Resource):
                 provider_id=datasource_provider_id,
                 avatar_url=oauth_response.metadata.get("avatar_url") or None,
                 name=oauth_response.metadata.get("name") or None,
+                expire_at=oauth_response.expires_at,
                 credentials=dict(oauth_response.credentials),
             )
         return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback")
diff --git a/api/models/oauth.py b/api/models/oauth.py
index 8e661051a7..4848b86af2 100644
--- a/api/models/oauth.py
+++ b/api/models/oauth.py
@@ -37,6 +37,7 @@ class DatasourceProvider(Base):
     encrypted_credentials: Mapped[dict] = db.Column(JSONB, nullable=False)
     avatar_url: Mapped[str] = db.Column(db.String(255), nullable=True, default="default")
     is_default: Mapped[bool] = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
+    expires_at: Mapped[int] = db.Column(db.Integer, nullable=False, default=-1)
 
     created_at: Mapped[datetime] = db.Column(db.DateTime, nullable=False, default=datetime.now)
     updated_at: Mapped[datetime] = db.Column(db.DateTime, nullable=False, default=datetime.now)
diff --git a/api/services/datasource_provider_service.py b/api/services/datasource_provider_service.py
index a0b61c758e..ef29654a35 100644
--- a/api/services/datasource_provider_service.py
+++ b/api/services/datasource_provider_service.py
@@ -383,6 +383,7 @@ class DatasourceProviderService:
         tenant_id: str,
         provider_id: DatasourceProviderID,
         avatar_url: str | None,
+        expire_at: int,
         credentials: dict,
         credential_id: str,
     ) -> None:
@@ -433,6 +434,7 @@ class DatasourceProviderService:
                     if key in provider_credential_secret_variables:
                         credentials[key] = encrypter.encrypt_token(tenant_id, value)
 
+                target_provider.expires_at = expire_at
                 target_provider.encrypted_credentials = credentials
                 target_provider.avatar_url = avatar_url or target_provider.avatar_url
                 session.commit()
@@ -443,6 +445,7 @@ class DatasourceProviderService:
         tenant_id: str,
         provider_id: DatasourceProviderID,
         avatar_url: str | None,
+        expire_at: int,
         credentials: dict,
     ) -> None:
         """
@@ -500,6 +503,7 @@ class DatasourceProviderService:
                     auth_type=credential_type.value,
                     encrypted_credentials=credentials,
                     avatar_url=avatar_url or "default",
+                    expires_at=expire_at,
                 )
                 session.add(datasource_provider)
                 session.commit()