From 424ec61f6ae4ad987add925989cef46f9d341726 Mon Sep 17 00:00:00 2001 From: Lenar Sharipov Date: Sun, 29 Mar 2026 12:43:39 +0200 Subject: [PATCH] Add OAuth authentication method for Snowflake --- src/databao_context_engine/__init__.py | 2 ++ .../plugins/databases/snowflake/config_file.py | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/databao_context_engine/__init__.py b/src/databao_context_engine/__init__.py index 50323cf4..5760a774 100644 --- a/src/databao_context_engine/__init__.py +++ b/src/databao_context_engine/__init__.py @@ -57,6 +57,7 @@ SnowflakeConfigFile, SnowflakeConnectionProperties, SnowflakeKeyPairAuth, + SnowflakeOAuthAuth, SnowflakePasswordAuth, SnowflakeSSOAuth, ) @@ -134,6 +135,7 @@ "SnowflakeConfigFile", "SnowflakeConnectionProperties", "SnowflakeSSOAuth", + "SnowflakeOAuthAuth", "SnowflakeKeyPairAuth", "SnowflakePasswordAuth", "SQLiteConfigFile", diff --git a/src/databao_context_engine/plugins/databases/snowflake/config_file.py b/src/databao_context_engine/plugins/databases/snowflake/config_file.py index 80cf9f34..541f05a4 100644 --- a/src/databao_context_engine/plugins/databases/snowflake/config_file.py +++ b/src/databao_context_engine/plugins/databases/snowflake/config_file.py @@ -22,13 +22,17 @@ class SnowflakeSSOAuth(BaseModel): authenticator: str = Field(description='e.g. "externalbrowser"') +class SnowflakeOAuthAuth(BaseModel): + token: Annotated[str, ConfigPropertyAnnotation(secret=True)] + + class SnowflakeConnectionProperties(BaseModel): account: Annotated[str, ConfigPropertyAnnotation(required=True)] warehouse: str | None = None database: str | None = None user: str | None = None role: str | None = None - auth: SnowflakePasswordAuth | SnowflakeKeyPairAuth | SnowflakeSSOAuth + auth: SnowflakePasswordAuth | SnowflakeKeyPairAuth | SnowflakeSSOAuth | SnowflakeOAuthAuth additional_properties: dict[str, Any] = {} def to_snowflake_kwargs(self) -> dict[str, Any]: