support auth key from content and from file (#40390)

* support auth key from content and from file

* add logs

* fix static check

airflow/providers/ydb/provider.yaml

@@ -30,7 +30,7 @@ versions:
 dependencies:
  - apache-airflow>=2.7.0
  - apache-airflow-providers-common-sql>=1.3.1
- - ydb>=3.11.3
+ - ydb>=3.12.1
 
 integrations:
  - integration-name: YDB

airflow/providers/ydb/utils/credentials.py

@@ -16,16 +16,17 @@
 # under the License.
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING, Any
 
 import ydb
 import ydb.iam.auth as auth
 
-from airflow.exceptions import AirflowException
-
 if TYPE_CHECKING:
  from airflow.models.connection import Connection
 
+log = logging.getLogger(__name__)
+
 
 def get_credentials_from_connection(
  endpoint: str, database: str, connection: Connection, connection_extra: dict[str, Any] | None = None
@@ -54,23 +55,29 @@ def get_credentials_from_connection(
  database=database,
  )
 
+ log.info("using login as credentials")
  return ydb.StaticCredentials(driver_config, user=connection.login, password=connection.password)
 
  connection_extra = connection_extra or {}
  token = connection_extra.get("token")
  if token:
+ log.info("using token as credentials")
  return ydb.AccessTokenCredentials(token)
 
  service_account_json_path = connection_extra.get("service_account_json_path")
  if service_account_json_path:
- return auth.BaseJWTCredentials.from_file(auth.ServiceAccountCredentials, service_account_json_path)
+ log.info("using service_account_json_path as credentials")
+ return auth.ServiceAccountCredentials.from_file(service_account_json_path)
 
  service_account_json = connection_extra.get("service_account_json")
  if service_account_json:
- raise AirflowException("service_account_json parameter is not supported yet")
+ log.info("using service_account_json as credentials")
+ return auth.ServiceAccountCredentials.from_content(service_account_json)
 
  use_vm_metadata = connection_extra.get("use_vm_metadata", False)
  if use_vm_metadata:
+ log.info("using vm metadata as credentials")
  return auth.MetadataUrlCredentials()
 
+ log.info("using anonymous access")
  return ydb.AnonymousCredentials()

generated/provider_dependencies.json

@@ -1343,7 +1343,7 @@
  "deps": [
  "apache-airflow-providers-common-sql>=1.3.1",
  "apache-airflow>=2.7.0",
- "ydb>=3.11.3"
+ "ydb>=3.12.1"
  ],
  "devel-deps": [],
  "plugins": [],

tests/providers/ydb/utils/test_credentials.py

@@ -68,6 +68,20 @@ def test_vm_metadata_creds(mock):
  mock.assert_called_once()
 
 
+@patch("ydb.iam.auth.BaseJWTCredentials.from_content")
+def test_service_account_json_creds(mock):
+ mock.return_value = MAGIC_CONST
+ c = Connection(conn_type="ydb", host="localhost")
+
+ credentials = get_credentials_from_connection(
+ TEST_ENDPOINT, TEST_DATABASE, c, {"service_account_json": "my_json"}
+ )
+ assert credentials == MAGIC_CONST
+ mock.assert_called_once()
+
+ assert mock.call_args.args == ("my_json",)
+
+
 @patch("ydb.iam.auth.BaseJWTCredentials.from_file")
 def test_service_account_json_path_creds(mock):
  mock.return_value = MAGIC_CONST
@@ -79,8 +93,7 @@ def test_service_account_json_path_creds(mock):
  assert credentials == MAGIC_CONST
  mock.assert_called_once()
 
- assert len(mock.call_args.args) == 2
- assert mock.call_args.args[1] == "my_path"
+ assert mock.call_args.args == ("my_path",)
 
 
 def test_creds_priority():
@@ -93,6 +106,7 @@ def test_creds_priority():
  TEST_DATABASE,
  c,
  {
+ "service_account_json": "my_json",
  "service_account_json_path": "my_path",
  "use_vm_metadata": True,
  "token": "my_token",
@@ -110,6 +124,7 @@ def test_creds_priority():
  TEST_DATABASE,
  c,
  {
+ "service_account_json": "my_json",
  "service_account_json_path": "my_path",
  "use_vm_metadata": True,
  "token": "my_token",
@@ -127,14 +142,31 @@ def test_creds_priority():
  TEST_DATABASE,
  c,
  {
+ "service_account_json": "my_json",
  "service_account_json_path": "my_path",
  "use_vm_metadata": True,
  },
  )
  assert credentials == MAGIC_CONST
  mock.assert_called_once()
 
- # 4. vm metadata
+ # 4. service account json
+ with patch("ydb.iam.auth.BaseJWTCredentials.from_content") as mock:
+ c = Connection(conn_type="ydb", host="localhost")
+ mock.return_value = MAGIC_CONST
+ credentials = get_credentials_from_connection(
+ TEST_ENDPOINT,
+ TEST_DATABASE,
+ c,
+ {
+ "service_account_json": "my_json",
+ "use_vm_metadata": True,
+ },
+ )
+ assert credentials == MAGIC_CONST
+ mock.assert_called_once()
+
+ # 5. vm metadata
  with patch("ydb.iam.auth.MetadataUrlCredentials") as mock:
  c = Connection(conn_type="ydb", host="localhost")
  mock.return_value = MAGIC_CONST
@@ -149,7 +181,7 @@ def test_creds_priority():
  assert credentials == MAGIC_CONST
  mock.assert_called_once()
 
- # 5. anonymous
+ # 6. anonymous
  with patch("ydb.AnonymousCredentials") as mock:
  c = Connection(conn_type="ydb", host="localhost")
  mock.return_value = MAGIC_CONST