Skip to content

Commit 255c327

Browse files
committed
Ability to pass credentials by string
1 parent 492dab6 commit 255c327

File tree

2 files changed

+34
-5
lines changed

2 files changed

+34
-5
lines changed

ydb_dbapi/connections.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from .errors import NotSupportedError
2121
from .utils import handle_ydb_errors
2222
from .utils import maybe_get_current_trace_id
23+
from .utils import prepare_credentials
2324

2425

2526
class IsolationLevel(str, Enum):
@@ -69,13 +70,13 @@ def __init__(
6970
port: str = "",
7071
database: str = "",
7172
ydb_table_path_prefix: str = "",
72-
credentials: ydb.AbstractCredentials | None = None,
73+
credentials: ydb.Credentials | dict | str | None = None,
7374
ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
7475
**kwargs: dict,
7576
) -> None:
76-
self.endpoint = f"grpc://{host}:{port}"
77+
self.endpoint = f"grpcs://{host}:{port}"
7778
self.database = database
78-
self.credentials = credentials
79+
self.credentials = prepare_credentials(credentials)
7980
self.table_path_prefix = ydb_table_path_prefix
8081

8182
self.connection_kwargs: dict = kwargs
@@ -170,7 +171,7 @@ def __init__(
170171
port: str = "",
171172
database: str = "",
172173
ydb_table_path_prefix: str = "",
173-
credentials: ydb.AbstractCredentials | None = None,
174+
credentials: ydb.Credentials | None = None,
174175
ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
175176
**kwargs: dict,
176177
) -> None:
@@ -333,7 +334,7 @@ def __init__(
333334
port: str = "",
334335
database: str = "",
335336
ydb_table_path_prefix: str = "",
336-
credentials: ydb.AbstractCredentials | None = None,
337+
credentials: ydb.Credentials | None = None,
337338
ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
338339
**kwargs: dict,
339340
) -> None:

ydb_dbapi/utils.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import functools
44
import importlib.util
5+
import json
56
from enum import Enum
67
from inspect import iscoroutinefunction
78
from typing import Any
@@ -117,3 +118,30 @@ def maybe_get_current_trace_id() -> str | None:
117118

118119
# Return None if OpenTelemetry is not available or trace ID is invalid
119120
return None
121+
122+
123+
def prepare_credentials(
124+
credentials: ydb.Credentials | dict | str | None,
125+
) -> ydb.Credentials | None:
126+
if not credentials:
127+
return None
128+
129+
if isinstance(credentials, ydb.Credentials):
130+
return credentials
131+
132+
if isinstance(credentials, str):
133+
credentials = json.loads(credentials)
134+
135+
if isinstance(credentials, dict):
136+
credentials = credentials or {}
137+
token = credentials.get("token")
138+
if token:
139+
return ydb.AccessTokenCredentials(token)
140+
141+
service_account_json = credentials.get("service_account_json")
142+
if service_account_json:
143+
return ydb.iam.ServiceAccountCredentials.from_content(
144+
json.dumps(service_account_json)
145+
)
146+
147+
return ydb.AnonymousCredentials()

0 commit comments

Comments
 (0)