From 53885e4bdaeea498ef8962ccbc407703d8fc1dc8 Mon Sep 17 00:00:00 2001
From: leavesster <11785335+leavesster@users.noreply.github.com>
Date: Fri, 10 Oct 2025 12:48:15 +0800
Subject: [PATCH] refactor: query auth use CredentialInput instead of str

---
 executor/python_executor/credential.py | 7 ++-----
 oocana/oocana/__init__.py              | 3 ++-
 oocana/oocana/context.py               | 5 +++--
 oocana/oocana/credential.py            | 4 ++++
 4 files changed, 11 insertions(+), 8 deletions(-)
 create mode 100644 oocana/oocana/credential.py

diff --git a/executor/python_executor/credential.py b/executor/python_executor/credential.py
index 299cd7a5..6f42481c 100644
--- a/executor/python_executor/credential.py
+++ b/executor/python_executor/credential.py
@@ -1,9 +1,6 @@
 from typing import Any, Dict
-from oocana import InputHandleDef
-class CredentialInput:
-    def __init__(self, type: str, id: str):
-        self.type = type
-        self.id = id
+from oocana import InputHandleDef, CredentialInput
+
 
 def generate_credential_input(credential_path: str) -> CredentialInput | None:
     """Generate a CredentialInput from a credential path string.
diff --git a/oocana/oocana/__init__.py b/oocana/oocana/__init__.py
index c2441a53..c88bc436 100644
--- a/oocana/oocana/__init__.py
+++ b/oocana/oocana/__init__.py
@@ -6,4 +6,5 @@
 from .extra import *  # noqa: F403
 from .schema import *  # noqa: F403
 from .mainframe import Mainframe as Mainframe  # noqa: F403
-from .serialization import setup_dataframe_serialization, CompressionOptions  # noqa: F403
\ No newline at end of file
+from .serialization import setup_dataframe_serialization, CompressionOptions  # noqa: F403
+from .credential import CredentialInput
\ No newline at end of file
diff --git a/oocana/oocana/context.py b/oocana/oocana/context.py
index 07606f67..ea922e61 100644
--- a/oocana/oocana/context.py
+++ b/oocana/oocana/context.py
@@ -11,6 +11,7 @@
 from .preview import PreviewPayload, DataFrame, PreviewPayloadInternal, ShapeDataFrame
 from .data import EXECUTOR_NAME
 from .internal import random_string, InternalAPI
+from .credential import CredentialInput
 import os.path
 import logging
 import hashlib
@@ -680,7 +681,7 @@ def response_callback(payload: Dict[str, Any]):
 
         return await f
     
-    async def query_auth(self, id: str) -> Dict[str, Any]:
+    async def query_auth(self, credential: CredentialInput) -> Dict[str, Any]:
         request_id = random_string(16)
         loop = asyncio.get_running_loop()
         f: asyncio.Future[Dict[str, Any]] = loop.create_future()
@@ -699,7 +700,7 @@ def response_callback(payload: Dict[str, Any]):
         self.__mainframe.send(self.job_info, {
             "type": "BlockRequest",
             "action": "QueryAuth",
-            "id": id,
+            "id": credential.id,
             "session_id": self.session_id,
             "job_id": self.job_id,
             "request_id": request_id,
diff --git a/oocana/oocana/credential.py b/oocana/oocana/credential.py
new file mode 100644
index 00000000..fdb3268b
--- /dev/null
+++ b/oocana/oocana/credential.py
@@ -0,0 +1,4 @@
+class CredentialInput:
+    def __init__(self, type: str, id: str):
+        self.type = type
+        self.id = id
\ No newline at end of file