Skip to content

Commit 91e6e60

Browse files
authored
support auth key from content and from file (apache#40390)
* support auth key from content and from file * add logs * fix static check
1 parent d0e4b8d commit 91e6e60

File tree

4 files changed

+49
-10
lines changed

4 files changed

+49
-10
lines changed

airflow/providers/ydb/provider.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ versions:
3030
dependencies:
3131
- apache-airflow>=2.7.0
3232
- apache-airflow-providers-common-sql>=1.3.1
33-
- ydb>=3.11.3
33+
- ydb>=3.12.1
3434

3535
integrations:
3636
- integration-name: YDB

airflow/providers/ydb/utils/credentials.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,17 @@
1616
# under the License.
1717
from __future__ import annotations
1818

19+
import logging
1920
from typing import TYPE_CHECKING, Any
2021

2122
import ydb
2223
import ydb.iam.auth as auth
2324

24-
from airflow.exceptions import AirflowException
25-
2625
if TYPE_CHECKING:
2726
from airflow.models.connection import Connection
2827

28+
log = logging.getLogger(__name__)
29+
2930

3031
def get_credentials_from_connection(
3132
endpoint: str, database: str, connection: Connection, connection_extra: dict[str, Any] | None = None
@@ -54,23 +55,29 @@ def get_credentials_from_connection(
5455
database=database,
5556
)
5657

58+
log.info("using login as credentials")
5759
return ydb.StaticCredentials(driver_config, user=connection.login, password=connection.password)
5860

5961
connection_extra = connection_extra or {}
6062
token = connection_extra.get("token")
6163
if token:
64+
log.info("using token as credentials")
6265
return ydb.AccessTokenCredentials(token)
6366

6467
service_account_json_path = connection_extra.get("service_account_json_path")
6568
if service_account_json_path:
66-
return auth.BaseJWTCredentials.from_file(auth.ServiceAccountCredentials, service_account_json_path)
69+
log.info("using service_account_json_path as credentials")
70+
return auth.ServiceAccountCredentials.from_file(service_account_json_path)
6771

6872
service_account_json = connection_extra.get("service_account_json")
6973
if service_account_json:
70-
raise AirflowException("service_account_json parameter is not supported yet")
74+
log.info("using service_account_json as credentials")
75+
return auth.ServiceAccountCredentials.from_content(service_account_json)
7176

7277
use_vm_metadata = connection_extra.get("use_vm_metadata", False)
7378
if use_vm_metadata:
79+
log.info("using vm metadata as credentials")
7480
return auth.MetadataUrlCredentials()
7581

82+
log.info("using anonymous access")
7683
return ydb.AnonymousCredentials()

generated/provider_dependencies.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1343,7 +1343,7 @@
13431343
"deps": [
13441344
"apache-airflow-providers-common-sql>=1.3.1",
13451345
"apache-airflow>=2.7.0",
1346-
"ydb>=3.11.3"
1346+
"ydb>=3.12.1"
13471347
],
13481348
"devel-deps": [],
13491349
"plugins": [],

tests/providers/ydb/utils/test_credentials.py

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,20 @@ def test_vm_metadata_creds(mock):
6868
mock.assert_called_once()
6969

7070

71+
@patch("ydb.iam.auth.BaseJWTCredentials.from_content")
72+
def test_service_account_json_creds(mock):
73+
mock.return_value = MAGIC_CONST
74+
c = Connection(conn_type="ydb", host="localhost")
75+
76+
credentials = get_credentials_from_connection(
77+
TEST_ENDPOINT, TEST_DATABASE, c, {"service_account_json": "my_json"}
78+
)
79+
assert credentials == MAGIC_CONST
80+
mock.assert_called_once()
81+
82+
assert mock.call_args.args == ("my_json",)
83+
84+
7185
@patch("ydb.iam.auth.BaseJWTCredentials.from_file")
7286
def test_service_account_json_path_creds(mock):
7387
mock.return_value = MAGIC_CONST
@@ -79,8 +93,7 @@ def test_service_account_json_path_creds(mock):
7993
assert credentials == MAGIC_CONST
8094
mock.assert_called_once()
8195

82-
assert len(mock.call_args.args) == 2
83-
assert mock.call_args.args[1] == "my_path"
96+
assert mock.call_args.args == ("my_path",)
8497

8598

8699
def test_creds_priority():
@@ -93,6 +106,7 @@ def test_creds_priority():
93106
TEST_DATABASE,
94107
c,
95108
{
109+
"service_account_json": "my_json",
96110
"service_account_json_path": "my_path",
97111
"use_vm_metadata": True,
98112
"token": "my_token",
@@ -110,6 +124,7 @@ def test_creds_priority():
110124
TEST_DATABASE,
111125
c,
112126
{
127+
"service_account_json": "my_json",
113128
"service_account_json_path": "my_path",
114129
"use_vm_metadata": True,
115130
"token": "my_token",
@@ -127,14 +142,31 @@ def test_creds_priority():
127142
TEST_DATABASE,
128143
c,
129144
{
145+
"service_account_json": "my_json",
130146
"service_account_json_path": "my_path",
131147
"use_vm_metadata": True,
132148
},
133149
)
134150
assert credentials == MAGIC_CONST
135151
mock.assert_called_once()
136152

137-
# 4. vm metadata
153+
# 4. service account json
154+
with patch("ydb.iam.auth.BaseJWTCredentials.from_content") as mock:
155+
c = Connection(conn_type="ydb", host="localhost")
156+
mock.return_value = MAGIC_CONST
157+
credentials = get_credentials_from_connection(
158+
TEST_ENDPOINT,
159+
TEST_DATABASE,
160+
c,
161+
{
162+
"service_account_json": "my_json",
163+
"use_vm_metadata": True,
164+
},
165+
)
166+
assert credentials == MAGIC_CONST
167+
mock.assert_called_once()
168+
169+
# 5. vm metadata
138170
with patch("ydb.iam.auth.MetadataUrlCredentials") as mock:
139171
c = Connection(conn_type="ydb", host="localhost")
140172
mock.return_value = MAGIC_CONST
@@ -149,7 +181,7 @@ def test_creds_priority():
149181
assert credentials == MAGIC_CONST
150182
mock.assert_called_once()
151183

152-
# 5. anonymous
184+
# 6. anonymous
153185
with patch("ydb.AnonymousCredentials") as mock:
154186
c = Connection(conn_type="ydb", host="localhost")
155187
mock.return_value = MAGIC_CONST

0 commit comments

Comments
 (0)