Skip to content

Commit 91fe868

Browse files
committed
add support for pat in snowflake
1 parent 4cdaeda commit 91fe868

3 files changed

Lines changed: 46 additions & 1 deletion

File tree

sqlit/domains/connections/providers/snowflake/adapter.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@ def connect(self, config: ConnectionConfig) -> Any:
9797
connect_args["private_key_file_pwd"] = extras["private_key_file_pwd"]
9898
if "oauth_token" in extras:
9999
connect_args["token"] = extras["oauth_token"]
100+
if "pat_token" in extras:
101+
connect_args["token"] = extras["pat_token"]
100102

101103
# Pass through any extra_options to the driver
102104
connect_args.update(config.extra_options)

sqlit/domains/connections/providers/snowflake/schema.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ def _get_snowflake_auth_options() -> tuple[SelectOption, ...]:
1717
SelectOption("externalbrowser", "SSO (Browser)"),
1818
SelectOption("snowflake_jwt", "Key Pair (JWT)"),
1919
SelectOption("oauth", "OAuth Token"),
20+
SelectOption("PROGRAMMATIC_ACCESS_TOKEN", "Programmatic Access Token"),
2021
)
2122

2223

@@ -26,6 +27,8 @@ def _get_snowflake_auth_options() -> tuple[SelectOption, ...]:
2627
_AUTH_NEEDS_PRIVATE_KEY = {"snowflake_jwt"}
2728
# Auth types that need OAuth token
2829
_AUTH_NEEDS_OAUTH = {"oauth"}
30+
# Auth types that need Programmatic Access Token
31+
_AUTH_NEEDS_PAT = {"PROGRAMMATIC_ACCESS_TOKEN"}
2932

3033

3134
SCHEMA = ConnectionSchema(
@@ -81,6 +84,14 @@ def _get_snowflake_auth_options() -> tuple[SelectOption, ...]:
8184
required=False,
8285
visible_when=lambda v: v.get("authenticator") in _AUTH_NEEDS_OAUTH,
8386
),
87+
SchemaField(
88+
name="pat_token",
89+
label="PAT",
90+
field_type=FieldType.PASSWORD,
91+
placeholder="PROGRAMMATIC_ACCESS_TOKEN",
92+
required=False,
93+
visible_when=lambda v: v.get("authenticator") in _AUTH_NEEDS_PAT,
94+
),
8495
_database_field(),
8596
SchemaField(
8697
name="warehouse",

tests/unit/test_extra_options_passthrough.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,37 @@ def test_snowflake_jwt_auth_options(self):
8080
assert call_kwargs.get("private_key_file") == "/path/to/key.p8"
8181
assert call_kwargs.get("private_key_file_pwd") == "secret"
8282

83+
def test_snowflake_pat_auth_options(self):
84+
"""Test Snowflake PAT authentication options are passed."""
85+
from sqlit.domains.connections.domain.config import ConnectionConfig, TcpEndpoint
86+
from sqlit.domains.connections.providers.snowflake.adapter import SnowflakeAdapter
87+
88+
mock_sf = MagicMock()
89+
mock_conn = MagicMock()
90+
mock_sf.connect.return_value = mock_conn
91+
92+
with patch.dict("sys.modules", {"snowflake.connector": mock_sf}):
93+
adapter = SnowflakeAdapter()
94+
config = ConnectionConfig(
95+
name="test_sf_pat",
96+
db_type="snowflake",
97+
endpoint=TcpEndpoint(
98+
host="account.snowflakecomputing.com",
99+
username="user",
100+
database="db",
101+
),
102+
options={
103+
"authenticator": "PROGRAMMATIC_ACCESS_TOKEN",
104+
"pat_token": "test_pat_token",
105+
},
106+
)
107+
108+
adapter.connect(config)
109+
110+
call_kwargs = mock_sf.connect.call_args[1]
111+
assert call_kwargs.get("authenticator") == "PROGRAMMATIC_ACCESS_TOKEN"
112+
assert call_kwargs.get("token") == "test_pat_token"
113+
83114
def test_postgresql_passes_extra_options(self):
84115
"""Test PostgreSQL adapter passes extra_options to driver."""
85116
from sqlit.domains.connections.domain.config import ConnectionConfig, TcpEndpoint
@@ -162,12 +193,13 @@ def test_snowflake_schema_has_auth_dropdown(self):
162193
break
163194

164195
assert auth_field is not None, "Snowflake schema should have authenticator field"
165-
assert len(auth_field.options) == 4
196+
assert len(auth_field.options) == 5
166197
auth_values = [opt.value for opt in auth_field.options]
167198
assert "default" in auth_values
168199
assert "externalbrowser" in auth_values
169200
assert "snowflake_jwt" in auth_values
170201
assert "oauth" in auth_values
202+
assert "PROGRAMMATIC_ACCESS_TOKEN" in auth_values
171203

172204
def test_snowflake_schema_has_private_key_fields(self):
173205
"""Test Snowflake schema includes private key fields for JWT auth."""

0 commit comments

Comments
 (0)