opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-08 02:43:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(services): migrate datasource_provider_service to SQLAlchemy 2.0 select() API (#34974)

Browse Source
This commit is contained in:
wdeveloper16 2026-04-12 03:23:24 +02:00 committed by GitHub
parent 7515eee0a8
commit 7bd5e80323
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 205 additions and 170 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

286
api/services/datasource_provider_service.py
View File

@ -4,7 +4,7 @@ from collections.abc import Mapping
from typing import Any from typing import Any
from graphon.model_runtime.entities.provider_entities import FormType from graphon.model_runtime.entities.provider_entities import FormType
from sqlalchemy import func, select from sqlalchemy import delete, func, select, update
from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config from configs import dify_config
@ -54,11 +54,13 @@ class DatasourceProviderService:
remove oauth custom client params remove oauth custom client params
""" """
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
session.query(DatasourceOauthTenantParamConfig).filter_by( session.execute(
tenant_id=tenant_id, delete(DatasourceOauthTenantParamConfig).where(
provider=datasource_provider_id.provider_name, DatasourceOauthTenantParamConfig.tenant_id == tenant_id,
plugin_id=datasource_provider_id.plugin_id, DatasourceOauthTenantParamConfig.provider == datasource_provider_id.provider_name,
).delete() DatasourceOauthTenantParamConfig.plugin_id == datasource_provider_id.plugin_id,
)
)
def decrypt_datasource_provider_credentials( def decrypt_datasource_provider_credentials(
self, self,
@ -110,15 +112,21 @@ class DatasourceProviderService:
""" """
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
if credential_id: if credential_id:
datasource_provider = ( datasource_provider = session.scalar(
session.query(DatasourceProvider).filter_by(tenant_id=tenant_id, id=credential_id).first() select(DatasourceProvider)
.where(DatasourceProvider.tenant_id == tenant_id, DatasourceProvider.id == credential_id)
.limit(1)
) )
else: else:
datasource_provider = ( datasource_provider = session.scalar(
session.query(DatasourceProvider) select(DatasourceProvider)
.filter_by(tenant_id=tenant_id, provider=provider, plugin_id=plugin_id) .where(
DatasourceProvider.tenant_id == tenant_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
.order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc())
.first() .limit(1)
) )
if not datasource_provider: if not datasource_provider:
return {} return {}
@ -173,12 +181,15 @@ class DatasourceProviderService:
get all datasource credentials by provider get all datasource credentials by provider
""" """
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
datasource_providers = ( datasource_providers = session.scalars(
session.query(DatasourceProvider) select(DatasourceProvider)
.filter_by(tenant_id=tenant_id, provider=provider, plugin_id=plugin_id) .where(
DatasourceProvider.tenant_id == tenant_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
.order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc())
.all() ).all()
)
if not datasource_providers: if not datasource_providers:
return [] return []
current_user = get_current_user() current_user = get_current_user()
@ -232,15 +243,15 @@ class DatasourceProviderService:
update datasource provider name update datasource provider name
""" """
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
target_provider = ( target_provider = session.scalar(
session.query(DatasourceProvider) select(DatasourceProvider)
.filter_by( .where(
tenant_id=tenant_id, DatasourceProvider.tenant_id == tenant_id,
id=credential_id, DatasourceProvider.id == credential_id,
provider=datasource_provider_id.provider_name, DatasourceProvider.provider == datasource_provider_id.provider_name,
plugin_id=datasource_provider_id.plugin_id, DatasourceProvider.plugin_id == datasource_provider_id.plugin_id,
) )
.first() .limit(1)
) )
if target_provider is None: if target_provider is None:
raise ValueError("provider not found") raise ValueError("provider not found")
@ -250,16 +261,16 @@ class DatasourceProviderService:
# check name is exist # check name is exist
if ( if (
session.query(DatasourceProvider) session.scalar(
.filter_by( select(func.count(DatasourceProvider.id)).where(
tenant_id=tenant_id, DatasourceProvider.tenant_id == tenant_id,
name=name, DatasourceProvider.name == name,
provider=datasource_provider_id.provider_name, DatasourceProvider.provider == datasource_provider_id.provider_name,
plugin_id=datasource_provider_id.plugin_id, DatasourceProvider.plugin_id == datasource_provider_id.plugin_id,
)
) )
.count() or 0
> 0 ) > 0:
):
raise ValueError("Authorization name is already exists") raise ValueError("Authorization name is already exists")
target_provider.name = name target_provider.name = name
@ -273,26 +284,31 @@ class DatasourceProviderService:
""" """
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
# get provider # get provider
target_provider = ( target_provider = session.scalar(
session.query(DatasourceProvider) select(DatasourceProvider)
.filter_by( .where(
tenant_id=tenant_id, DatasourceProvider.tenant_id == tenant_id,
id=credential_id, DatasourceProvider.id == credential_id,
provider=datasource_provider_id.provider_name, DatasourceProvider.provider == datasource_provider_id.provider_name,
plugin_id=datasource_provider_id.plugin_id, DatasourceProvider.plugin_id == datasource_provider_id.plugin_id,
) )
.first() .limit(1)
) )
if target_provider is None: if target_provider is None:
raise ValueError("provider not found") raise ValueError("provider not found")
# clear default provider # clear default provider
session.query(DatasourceProvider).filter_by( session.execute(
tenant_id=tenant_id, update(DatasourceProvider)
provider=target_provider.provider, .where(
plugin_id=target_provider.plugin_id, DatasourceProvider.tenant_id == tenant_id,
is_default=True, DatasourceProvider.provider == target_provider.provider,
).update({"is_default": False}) DatasourceProvider.plugin_id == target_provider.plugin_id,
DatasourceProvider.is_default.is_(True),
)
.values(is_default=False)
.execution_options(synchronize_session=False)
)
# set new default provider # set new default provider
target_provider.is_default = True target_provider.is_default = True
@ -311,14 +327,14 @@ class DatasourceProviderService:
if client_params is None and enabled is None: if client_params is None and enabled is None:
return return
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
tenant_oauth_client_params = ( tenant_oauth_client_params = session.scalar(
session.query(DatasourceOauthTenantParamConfig) select(DatasourceOauthTenantParamConfig)
.filter_by( .where(
tenant_id=tenant_id, DatasourceOauthTenantParamConfig.tenant_id == tenant_id,
provider=datasource_provider_id.provider_name, DatasourceOauthTenantParamConfig.provider == datasource_provider_id.provider_name,
plugin_id=datasource_provider_id.plugin_id, DatasourceOauthTenantParamConfig.plugin_id == datasource_provider_id.plugin_id,
) )
.first() .limit(1)
) )
if not tenant_oauth_client_params: if not tenant_oauth_client_params:
@ -351,9 +367,14 @@ class DatasourceProviderService:
""" """
with Session(db.engine).no_autoflush as session: with Session(db.engine).no_autoflush as session:
return ( return (
session.query(DatasourceOauthParamConfig) session.scalar(
.filter_by(provider=datasource_provider_id.provider_name, plugin_id=datasource_provider_id.plugin_id) select(DatasourceOauthParamConfig)
.first() .where(
DatasourceOauthParamConfig.provider == datasource_provider_id.provider_name,
DatasourceOauthParamConfig.plugin_id == datasource_provider_id.plugin_id,
)
.limit(1)
)
is not None is not None
) )
@ -423,15 +444,15 @@ class DatasourceProviderService:
plugin_id = datasource_provider_id.plugin_id plugin_id = datasource_provider_id.plugin_id
with Session(db.engine).no_autoflush as session: with Session(db.engine).no_autoflush as session:
# get tenant oauth client params # get tenant oauth client params
tenant_oauth_client_params = ( tenant_oauth_client_params = session.scalar(
session.query(DatasourceOauthTenantParamConfig) select(DatasourceOauthTenantParamConfig)
.filter_by( .where(
tenant_id=tenant_id, DatasourceOauthTenantParamConfig.tenant_id == tenant_id,
provider=provider, DatasourceOauthTenantParamConfig.provider == provider,
plugin_id=plugin_id, DatasourceOauthTenantParamConfig.plugin_id == plugin_id,
enabled=True, DatasourceOauthTenantParamConfig.enabled.is_(True),
) )
.first() .limit(1)
) )
if tenant_oauth_client_params: if tenant_oauth_client_params:
encrypter, _ = self.get_oauth_encrypter(tenant_id, datasource_provider_id) encrypter, _ = self.get_oauth_encrypter(tenant_id, datasource_provider_id)
@ -443,8 +464,13 @@ class DatasourceProviderService:
is_verified = PluginService.is_plugin_verified(tenant_id, provider_controller.plugin_unique_identifier) is_verified = PluginService.is_plugin_verified(tenant_id, provider_controller.plugin_unique_identifier)
if is_verified: if is_verified:
# fallback to system oauth client params # fallback to system oauth client params
oauth_client_params = ( oauth_client_params = session.scalar(
session.query(DatasourceOauthParamConfig).filter_by(provider=provider, plugin_id=plugin_id).first() select(DatasourceOauthParamConfig)
.where(
DatasourceOauthParamConfig.provider == provider,
DatasourceOauthParamConfig.plugin_id == plugin_id,
)
.limit(1)
) )
if oauth_client_params: if oauth_client_params:
return oauth_client_params.system_credentials return oauth_client_params.system_credentials
@ -455,15 +481,13 @@ class DatasourceProviderService:
def generate_next_datasource_provider_name( def generate_next_datasource_provider_name(
session: Session, tenant_id: str, provider_id: DatasourceProviderID, credential_type: CredentialType session: Session, tenant_id: str, provider_id: DatasourceProviderID, credential_type: CredentialType
) -> str: ) -> str:
db_providers = ( db_providers = session.scalars(
session.query(DatasourceProvider) select(DatasourceProvider).where(
.filter_by( DatasourceProvider.tenant_id == tenant_id,
tenant_id=tenant_id, DatasourceProvider.provider == provider_id.provider_name,
provider=provider_id.provider_name, DatasourceProvider.plugin_id == provider_id.plugin_id,
plugin_id=provider_id.plugin_id,
) )
.all() ).all()
)
return generate_incremental_name( return generate_incremental_name(
[provider.name for provider in db_providers], [provider.name for provider in db_providers],
f"{credential_type.get_name()}", f"{credential_type.get_name()}",
@ -485,8 +509,10 @@ class DatasourceProviderService:
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{CredentialType.OAUTH2.value}" lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{CredentialType.OAUTH2.value}"
with redis_client.lock(lock, timeout=20): with redis_client.lock(lock, timeout=20):
target_provider = ( target_provider = session.scalar(
session.query(DatasourceProvider).filter_by(id=credential_id, tenant_id=tenant_id).first() select(DatasourceProvider)
.where(DatasourceProvider.id == credential_id, DatasourceProvider.tenant_id == tenant_id)
.limit(1)
) )
if target_provider is None: if target_provider is None:
raise ValueError("provider not found") raise ValueError("provider not found")
@ -496,25 +522,28 @@ class DatasourceProviderService:
db_provider_name = target_provider.name db_provider_name = target_provider.name
else: else:
name_conflict = ( name_conflict = (
session.query(DatasourceProvider) session.scalar(
.filter_by( select(func.count(DatasourceProvider.id)).where(
tenant_id=tenant_id, DatasourceProvider.tenant_id == tenant_id,
name=db_provider_name, DatasourceProvider.name == db_provider_name,
provider=provider_id.provider_name, DatasourceProvider.provider == provider_id.provider_name,
plugin_id=provider_id.plugin_id, DatasourceProvider.plugin_id == provider_id.plugin_id,
auth_type=CredentialType.OAUTH2.value, DatasourceProvider.auth_type == CredentialType.OAUTH2.value,
)
) )
.count() or 0
) )
if name_conflict > 0: if name_conflict > 0:
db_provider_name = generate_incremental_name( db_provider_name = generate_incremental_name(
[ [
provider.name provider.name
for provider in session.query(DatasourceProvider).filter_by( for provider in session.scalars(
tenant_id=tenant_id, select(DatasourceProvider).where(
provider=provider_id.provider_name, DatasourceProvider.tenant_id == tenant_id,
plugin_id=provider_id.plugin_id, DatasourceProvider.provider == provider_id.provider_name,
) DatasourceProvider.plugin_id == provider_id.plugin_id,
)
).all()
], ],
db_provider_name, db_provider_name,
) )
@ -556,25 +585,27 @@ class DatasourceProviderService:
) )
else: else:
if ( if (
session.query(DatasourceProvider) session.scalar(
.filter_by( select(func.count(DatasourceProvider.id)).where(
tenant_id=tenant_id, DatasourceProvider.tenant_id == tenant_id,
name=db_provider_name, DatasourceProvider.name == db_provider_name,
provider=provider_id.provider_name, DatasourceProvider.provider == provider_id.provider_name,
plugin_id=provider_id.plugin_id, DatasourceProvider.plugin_id == provider_id.plugin_id,
auth_type=credential_type.value, DatasourceProvider.auth_type == credential_type.value,
)
) )
.count() or 0
> 0 ) > 0:
):
db_provider_name = generate_incremental_name( db_provider_name = generate_incremental_name(
[ [
provider.name provider.name
for provider in session.query(DatasourceProvider).filter_by( for provider in session.scalars(
tenant_id=tenant_id, select(DatasourceProvider).where(
provider=provider_id.provider_name, DatasourceProvider.tenant_id == tenant_id,
plugin_id=provider_id.plugin_id, DatasourceProvider.provider == provider_id.provider_name,
) DatasourceProvider.plugin_id == provider_id.plugin_id,
)
).all()
], ],
db_provider_name, db_provider_name,
) )
@ -627,11 +658,16 @@ class DatasourceProviderService:
# check name is exist # check name is exist
if ( if (
session.query(DatasourceProvider) session.scalar(
.filter_by(tenant_id=tenant_id, plugin_id=plugin_id, provider=provider_name, name=db_provider_name) select(func.count(DatasourceProvider.id)).where(
.count() DatasourceProvider.tenant_id == tenant_id,
> 0 DatasourceProvider.plugin_id == plugin_id,
): DatasourceProvider.provider == provider_name,
DatasourceProvider.name == db_provider_name,
)
)
or 0
) > 0:
raise ValueError("Authorization name is already exists") raise ValueError("Authorization name is already exists")
try: try:
@ -918,21 +954,31 @@ class DatasourceProviderService:
""" """
with sessionmaker(bind=db.engine).begin() as session: with sessionmaker(bind=db.engine).begin() as session:
datasource_provider = ( datasource_provider = session.scalar(
session.query(DatasourceProvider) select(DatasourceProvider)
.filter_by(tenant_id=tenant_id, id=auth_id, provider=provider, plugin_id=plugin_id) .where(
.first() DatasourceProvider.tenant_id == tenant_id,
DatasourceProvider.id == auth_id,
DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
.limit(1)
) )
if not datasource_provider: if not datasource_provider:
raise ValueError("Datasource provider not found") raise ValueError("Datasource provider not found")
# update name # update name
if name and name != datasource_provider.name: if name and name != datasource_provider.name:
if ( if (
session.query(DatasourceProvider) session.scalar(
.filter_by(tenant_id=tenant_id, name=name, provider=provider, plugin_id=plugin_id) select(func.count(DatasourceProvider.id)).where(
.count() DatasourceProvider.tenant_id == tenant_id,
> 0 DatasourceProvider.name == name,
): DatasourceProvider.provider == provider,
DatasourceProvider.plugin_id == plugin_id,
)
)
or 0
) > 0:
raise ValueError("Authorization name is already exists") raise ValueError("Authorization name is already exists")
datasource_provider.name = name datasource_provider.name = name

89
api/tests/unit_tests/services/test_datasource_provider_service.py
View File

@ -179,11 +179,11 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_return_true_when_system_oauth_params_exist(self, service, mock_db_session): def test_should_return_true_when_system_oauth_params_exist(self, service, mock_db_session):
mock_db_session.query().first.return_value = MagicMock() mock_db_session.scalar.return_value = MagicMock()
assert service.is_system_oauth_params_exist(make_id()) is True assert service.is_system_oauth_params_exist(make_id()) is True
def test_should_return_false_when_system_oauth_params_missing(self, service, mock_db_session): def test_should_return_false_when_system_oauth_params_missing(self, service, mock_db_session):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
assert service.is_system_oauth_params_exist(make_id()) is False assert service.is_system_oauth_params_exist(make_id()) is False
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
@ -205,7 +205,7 @@ class TestDatasourceProviderService:
def test_should_delete_tenant_config_when_removing_oauth_params(self, service, mock_db_session): def test_should_delete_tenant_config_when_removing_oauth_params(self, service, mock_db_session):
service.remove_oauth_custom_client_params("t1", make_id()) service.remove_oauth_custom_client_params("t1", make_id())
mock_db_session.query().delete.assert_called_once() mock_db_session.execute.assert_called_once()
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
# setup_oauth_custom_client_params (315-351) # setup_oauth_custom_client_params (315-351)
@ -217,14 +217,14 @@ class TestDatasourceProviderService:
mock_db_session.add.assert_not_called() mock_db_session.add.assert_not_called()
def test_should_create_new_config_when_none_exists(self, service, mock_db_session): def test_should_create_new_config_when_none_exists(self, service, mock_db_session):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)): with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)):
service.setup_oauth_custom_client_params("t1", make_id(), {"k": "v"}, True) service.setup_oauth_custom_client_params("t1", make_id(), {"k": "v"}, True)
mock_db_session.add.assert_called_once() mock_db_session.add.assert_called_once()
def test_should_update_existing_config_when_record_found(self, service, mock_db_session): def test_should_update_existing_config_when_record_found(self, service, mock_db_session):
existing = MagicMock() existing = MagicMock()
mock_db_session.query().first.return_value = existing mock_db_session.scalar.return_value = existing
with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)): with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)):
service.setup_oauth_custom_client_params("t1", make_id(), {"k": "v"}, False) service.setup_oauth_custom_client_params("t1", make_id(), {"k": "v"}, False)
mock_db_session.add.assert_not_called() # update in place, no add mock_db_session.add.assert_not_called() # update in place, no add
@ -255,7 +255,7 @@ class TestDatasourceProviderService:
def test_should_return_empty_dict_when_credential_not_found(self, service, mock_db_session, mock_user): def test_should_return_empty_dict_when_credential_not_found(self, service, mock_db_session, mock_user):
with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): with patch("services.datasource_provider_service.get_current_user", return_value=mock_user):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
assert service.get_datasource_credentials("t1", "prov", "org/plug") == {} assert service.get_datasource_credentials("t1", "prov", "org/plug") == {}
def test_should_refresh_oauth_tokens_when_expired(self, service, mock_db_session, mock_user): def test_should_refresh_oauth_tokens_when_expired(self, service, mock_db_session, mock_user):
@ -264,7 +264,7 @@ class TestDatasourceProviderService:
p.auth_type = "oauth2" p.auth_type = "oauth2"
p.expires_at = 0 # expired p.expires_at = 0 # expired
p.encrypted_credentials = {"tok": "x"} p.encrypted_credentials = {"tok": "x"}
mock_db_session.query().first.return_value = p mock_db_session.scalar.return_value = p
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service, "get_oauth_client", return_value={"oc": "v"}), patch.object(service, "get_oauth_client", return_value={"oc": "v"}),
@ -278,7 +278,7 @@ class TestDatasourceProviderService:
p.auth_type = "api_key" p.auth_type = "api_key"
p.expires_at = -1 # sentinel: never expires p.expires_at = -1 # sentinel: never expires
p.encrypted_credentials = {"k": "v"} p.encrypted_credentials = {"k": "v"}
mock_db_session.query().first.return_value = p mock_db_session.scalar.return_value = p
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service, "decrypt_datasource_provider_credentials", return_value={"k": "plain"}), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"k": "plain"}),
@ -292,7 +292,7 @@ class TestDatasourceProviderService:
p.auth_type = "api_key" p.auth_type = "api_key"
p.expires_at = -1 p.expires_at = -1
p.encrypted_credentials = {} p.encrypted_credentials = {}
mock_db_session.query().first.return_value = p mock_db_session.scalar.return_value = p
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service, "decrypt_datasource_provider_credentials", return_value={"k": "v"}), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"k": "v"}),
@ -306,7 +306,7 @@ class TestDatasourceProviderService:
def test_should_return_empty_list_when_no_provider_credentials_exist(self, service, mock_db_session, mock_user): def test_should_return_empty_list_when_no_provider_credentials_exist(self, service, mock_db_session, mock_user):
with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): with patch("services.datasource_provider_service.get_current_user", return_value=mock_user):
mock_db_session.query().all.return_value = [] mock_db_session.scalars.return_value.all.return_value = []
assert service.get_all_datasource_credentials_by_provider("t1", "prov", "org/plug") == [] assert service.get_all_datasource_credentials_by_provider("t1", "prov", "org/plug") == []
def test_should_refresh_and_return_credentials_when_oauth_expired(self, service, mock_db_session, mock_user): def test_should_refresh_and_return_credentials_when_oauth_expired(self, service, mock_db_session, mock_user):
@ -314,7 +314,7 @@ class TestDatasourceProviderService:
p.auth_type = "oauth2" p.auth_type = "oauth2"
p.expires_at = 0 p.expires_at = 0
p.encrypted_credentials = {"t": "x"} p.encrypted_credentials = {"t": "x"}
mock_db_session.query().all.return_value = [p] mock_db_session.scalars.return_value.all.return_value = [p]
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service, "get_oauth_client", return_value={"oc": "v"}), patch.object(service, "get_oauth_client", return_value={"oc": "v"}),
@ -328,22 +328,21 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_raise_value_error_when_provider_not_found_on_name_update(self, service, mock_db_session): def test_should_raise_value_error_when_provider_not_found_on_name_update(self, service, mock_db_session):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
with pytest.raises(ValueError, match="not found"): with pytest.raises(ValueError, match="not found"):
service.update_datasource_provider_name("t1", make_id(), "new", "cred-id") service.update_datasource_provider_name("t1", make_id(), "new", "cred-id")
def test_should_return_early_when_new_name_matches_current(self, service, mock_db_session): def test_should_return_early_when_new_name_matches_current(self, service, mock_db_session):
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
p.name = "same" p.name = "same"
mock_db_session.query().first.return_value = p mock_db_session.scalar.return_value = p
service.update_datasource_provider_name("t1", make_id(), "same", "cred-id") service.update_datasource_provider_name("t1", make_id(), "same", "cred-id")
def test_should_raise_value_error_when_name_already_exists(self, service, mock_db_session): def test_should_raise_value_error_when_name_already_exists(self, service, mock_db_session):
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
p.name = "old_name" p.name = "old_name"
p.is_default = False p.is_default = False
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 1] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 1 # conflict
with pytest.raises(ValueError, match="already exists"): with pytest.raises(ValueError, match="already exists"):
service.update_datasource_provider_name("t1", make_id(), "new_name", "some-id") service.update_datasource_provider_name("t1", make_id(), "new_name", "some-id")
@ -351,8 +350,7 @@ class TestDatasourceProviderService:
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
p.name = "old_name" p.name = "old_name"
p.is_default = False p.is_default = False
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 0] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 0
service.update_datasource_provider_name("t1", make_id(), "new_name", "some-id") service.update_datasource_provider_name("t1", make_id(), "new_name", "some-id")
assert p.name == "new_name" assert p.name == "new_name"
@ -361,7 +359,7 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_raise_value_error_when_target_provider_not_found(self, service, mock_db_session): def test_should_raise_value_error_when_target_provider_not_found(self, service, mock_db_session):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
with pytest.raises(ValueError, match="not found"): with pytest.raises(ValueError, match="not found"):
service.set_default_datasource_provider("t1", make_id(), "bad-id") service.set_default_datasource_provider("t1", make_id(), "bad-id")
@ -369,7 +367,7 @@ class TestDatasourceProviderService:
target = MagicMock(spec=DatasourceProvider) target = MagicMock(spec=DatasourceProvider)
target.provider = "provider" target.provider = "provider"
target.plugin_id = "org/plug" target.plugin_id = "org/plug"
mock_db_session.query().first.return_value = target mock_db_session.scalar.return_value = target
service.set_default_datasource_provider("t1", make_id(), "new-id") service.set_default_datasource_provider("t1", make_id(), "new-id")
assert target.is_default is True assert target.is_default is True
@ -428,13 +426,13 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_use_tenant_config_when_available(self, service, mock_db_session): def test_should_use_tenant_config_when_available(self, service, mock_db_session):
mock_db_session.query().first.return_value = MagicMock(client_params={"k": "v"}) mock_db_session.scalar.return_value = MagicMock(client_params={"k": "v"})
with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)): with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)):
result = service.get_oauth_client("t1", make_id()) result = service.get_oauth_client("t1", make_id())
assert result == {"k": "dec"} assert result == {"k": "dec"}
def test_should_fallback_to_system_credentials_when_tenant_config_missing(self, service, mock_db_session): def test_should_fallback_to_system_credentials_when_tenant_config_missing(self, service, mock_db_session):
mock_db_session.query().first.side_effect = [None, MagicMock(system_credentials={"k": "sys"})] mock_db_session.scalar.side_effect = [None, MagicMock(system_credentials={"k": "sys"})]
with ( with (
patch.object(service.provider_manager, "fetch_datasource_provider"), patch.object(service.provider_manager, "fetch_datasource_provider"),
patch("services.datasource_provider_service.PluginService.is_plugin_verified", return_value=True), patch("services.datasource_provider_service.PluginService.is_plugin_verified", return_value=True),
@ -444,7 +442,7 @@ class TestDatasourceProviderService:
def test_should_raise_value_error_when_no_oauth_config_available(self, service, mock_db_session): def test_should_raise_value_error_when_no_oauth_config_available(self, service, mock_db_session):
"""Neither tenant nor system credentials → raises ValueError.""" """Neither tenant nor system credentials → raises ValueError."""
mock_db_session.query().first.side_effect = [None, None] mock_db_session.scalar.side_effect = [None, None]
with ( with (
patch.object(service.provider_manager, "fetch_datasource_provider"), patch.object(service.provider_manager, "fetch_datasource_provider"),
patch("services.datasource_provider_service.PluginService.is_plugin_verified", return_value=False), patch("services.datasource_provider_service.PluginService.is_plugin_verified", return_value=False),
@ -457,15 +455,14 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_add_oauth_provider_successfully_when_name_is_unique(self, service, mock_db_session): def test_should_add_oauth_provider_successfully_when_name_is_unique(self, service, mock_db_session):
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
with patch.object(service, "extract_secret_variables", return_value=[]): with patch.object(service, "extract_secret_variables", return_value=[]):
service.add_datasource_oauth_provider("new", "t1", make_id(), "http://cb", 9999, {}) service.add_datasource_oauth_provider("new", "t1", make_id(), "http://cb", 9999, {})
mock_db_session.add.assert_called_once() mock_db_session.add.assert_called_once()
def test_should_auto_rename_when_oauth_provider_name_conflicts(self, service, mock_db_session): def test_should_auto_rename_when_oauth_provider_name_conflicts(self, service, mock_db_session):
"""Conflict on name results in auto-incremented name, not an error.""" """Conflict on name results in auto-incremented name, not an error."""
mock_db_session.query().count.return_value = 1 # conflict first, then auto-named mock_db_session.scalar.return_value = 1 # conflict first, then auto-named
mock_db_session.query().all.return_value = []
with ( with (
patch.object(service, "extract_secret_variables", return_value=[]), patch.object(service, "extract_secret_variables", return_value=[]),
patch.object(service, "generate_next_datasource_provider_name", return_value="new_gen"), patch.object(service, "generate_next_datasource_provider_name", return_value="new_gen"),
@ -475,8 +472,7 @@ class TestDatasourceProviderService:
def test_should_auto_generate_name_when_none_provided_for_oauth(self, service, mock_db_session): def test_should_auto_generate_name_when_none_provided_for_oauth(self, service, mock_db_session):
"""name=None causes auto-generation via generate_next_datasource_provider_name.""" """name=None causes auto-generation via generate_next_datasource_provider_name."""
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
mock_db_session.query().all.return_value = []
with ( with (
patch.object(service, "extract_secret_variables", return_value=[]), patch.object(service, "extract_secret_variables", return_value=[]),
patch.object(service, "generate_next_datasource_provider_name", return_value="auto"), patch.object(service, "generate_next_datasource_provider_name", return_value="auto"),
@ -485,13 +481,13 @@ class TestDatasourceProviderService:
mock_db_session.add.assert_called_once() mock_db_session.add.assert_called_once()
def test_should_encrypt_secret_fields_when_adding_oauth_provider(self, service, mock_db_session): def test_should_encrypt_secret_fields_when_adding_oauth_provider(self, service, mock_db_session):
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
with patch.object(service, "extract_secret_variables", return_value=["secret_key"]): with patch.object(service, "extract_secret_variables", return_value=["secret_key"]):
service.add_datasource_oauth_provider("nm", "t1", make_id(), "http://cb", 9999, {"secret_key": "value"}) service.add_datasource_oauth_provider("nm", "t1", make_id(), "http://cb", 9999, {"secret_key": "value"})
self._enc.encrypt_token.assert_called() self._enc.encrypt_token.assert_called()
def test_should_acquire_redis_lock_when_adding_oauth_provider(self, service, mock_db_session): def test_should_acquire_redis_lock_when_adding_oauth_provider(self, service, mock_db_session):
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
with patch.object(service, "extract_secret_variables", return_value=[]): with patch.object(service, "extract_secret_variables", return_value=[]):
service.add_datasource_oauth_provider("nm", "t1", make_id(), "http://cb", 9999, {}) service.add_datasource_oauth_provider("nm", "t1", make_id(), "http://cb", 9999, {})
self._redis.lock.assert_called() self._redis.lock.assert_called()
@ -501,23 +497,21 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_raise_value_error_when_credential_id_not_found_on_reauth(self, service, mock_db_session): def test_should_raise_value_error_when_credential_id_not_found_on_reauth(self, service, mock_db_session):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
with patch.object(service, "extract_secret_variables", return_value=[]): with patch.object(service, "extract_secret_variables", return_value=[]):
with pytest.raises(ValueError, match="not found"): with pytest.raises(ValueError, match="not found"):
service.reauthorize_datasource_oauth_provider("n", "t1", make_id(), "u", 1, {}, "bad-id") service.reauthorize_datasource_oauth_provider("n", "t1", make_id(), "u", 1, {}, "bad-id")
def test_should_reauthorize_and_commit_when_credential_found(self, service, mock_db_session): def test_should_reauthorize_and_commit_when_credential_found(self, service, mock_db_session):
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 0] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 0
with patch.object(service, "extract_secret_variables", return_value=[]): with patch.object(service, "extract_secret_variables", return_value=[]):
service.reauthorize_datasource_oauth_provider("n", "t1", make_id(), "u", 1, {}, "oid") service.reauthorize_datasource_oauth_provider("n", "t1", make_id(), "u", 1, {}, "oid")
def test_should_auto_rename_when_reauth_name_conflicts(self, service, mock_db_session): def test_should_auto_rename_when_reauth_name_conflicts(self, service, mock_db_session):
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 1] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 1 # conflict mock_db_session.scalars.return_value.all.return_value = []
mock_db_session.query().all.return_value = []
with patch.object(service, "extract_secret_variables", return_value=["tok"]): with patch.object(service, "extract_secret_variables", return_value=["tok"]):
service.reauthorize_datasource_oauth_provider( service.reauthorize_datasource_oauth_provider(
"conflict_name", "t1", make_id(), "u", 9999, {"tok": "v"}, "cred-id" "conflict_name", "t1", make_id(), "u", 9999, {"tok": "v"}, "cred-id"
@ -525,16 +519,14 @@ class TestDatasourceProviderService:
def test_should_encrypt_secret_fields_when_reauthorizing(self, service, mock_db_session): def test_should_encrypt_secret_fields_when_reauthorizing(self, service, mock_db_session):
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 0] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 0
with patch.object(service, "extract_secret_variables", return_value=["tok"]): with patch.object(service, "extract_secret_variables", return_value=["tok"]):
service.reauthorize_datasource_oauth_provider(None, "t1", make_id(), "u", 9999, {"tok": "val"}, "cred-id") service.reauthorize_datasource_oauth_provider(None, "t1", make_id(), "u", 9999, {"tok": "val"}, "cred-id")
self._enc.encrypt_token.assert_called() self._enc.encrypt_token.assert_called()
def test_should_acquire_redis_lock_when_reauthorizing(self, service, mock_db_session): def test_should_acquire_redis_lock_when_reauthorizing(self, service, mock_db_session):
p = MagicMock(spec=DatasourceProvider) p = MagicMock(spec=DatasourceProvider)
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 0] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 0
with patch.object(service, "extract_secret_variables", return_value=[]): with patch.object(service, "extract_secret_variables", return_value=[]):
service.reauthorize_datasource_oauth_provider("n", "t1", make_id(), "u", 1, {}, "oid") service.reauthorize_datasource_oauth_provider("n", "t1", make_id(), "u", 1, {}, "oid")
self._redis.lock.assert_called() self._redis.lock.assert_called()
@ -545,13 +537,13 @@ class TestDatasourceProviderService:
def test_should_raise_value_error_when_api_key_name_already_exists(self, service, mock_db_session, mock_user): def test_should_raise_value_error_when_api_key_name_already_exists(self, service, mock_db_session, mock_user):
"""explicit name supplied + conflict → raises ValueError immediately.""" """explicit name supplied + conflict → raises ValueError immediately."""
mock_db_session.query().count.return_value = 1 mock_db_session.scalar.return_value = 1
with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): with patch("services.datasource_provider_service.get_current_user", return_value=mock_user):
with pytest.raises(ValueError, match="already exists"): with pytest.raises(ValueError, match="already exists"):
service.add_datasource_api_key_provider("clash", "t1", make_id(), {"sk": "v"}) service.add_datasource_api_key_provider("clash", "t1", make_id(), {"sk": "v"})
def test_should_raise_value_error_when_credentials_validation_fails(self, service, mock_db_session, mock_user): def test_should_raise_value_error_when_credentials_validation_fails(self, service, mock_db_session, mock_user):
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service.provider_manager, "validate_provider_credentials", side_effect=Exception("bad cred")), patch.object(service.provider_manager, "validate_provider_credentials", side_effect=Exception("bad cred")),
@ -561,7 +553,7 @@ class TestDatasourceProviderService:
service.add_datasource_api_key_provider("nm", "t1", make_id(), {"k": "v"}) service.add_datasource_api_key_provider("nm", "t1", make_id(), {"k": "v"})
def test_should_add_api_key_provider_and_commit_when_valid(self, service, mock_db_session, mock_user): def test_should_add_api_key_provider_and_commit_when_valid(self, service, mock_db_session, mock_user):
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service.provider_manager, "validate_provider_credentials"), patch.object(service.provider_manager, "validate_provider_credentials"),
@ -571,7 +563,7 @@ class TestDatasourceProviderService:
mock_db_session.add.assert_called_once() mock_db_session.add.assert_called_once()
def test_should_acquire_redis_lock_when_adding_api_key_provider(self, service, mock_db_session, mock_user): def test_should_acquire_redis_lock_when_adding_api_key_provider(self, service, mock_db_session, mock_user):
mock_db_session.query().count.return_value = 0 mock_db_session.scalar.return_value = 0
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service.provider_manager, "validate_provider_credentials"), patch.object(service.provider_manager, "validate_provider_credentials"),
@ -694,7 +686,7 @@ class TestDatasourceProviderService:
# ----------------------------------------------------------------------- # -----------------------------------------------------------------------
def test_should_raise_value_error_when_credential_not_found_on_update(self, service, mock_db_session, mock_user): def test_should_raise_value_error_when_credential_not_found_on_update(self, service, mock_db_session, mock_user):
mock_db_session.query().first.return_value = None mock_db_session.scalar.return_value = None
with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): with patch("services.datasource_provider_service.get_current_user", return_value=mock_user):
with pytest.raises(ValueError, match="not found"): with pytest.raises(ValueError, match="not found"):
service.update_datasource_credentials("t1", "id", "prov", "org/plug", {}, "name") service.update_datasource_credentials("t1", "id", "prov", "org/plug", {}, "name")
@ -704,8 +696,7 @@ class TestDatasourceProviderService:
p.name = "old_name" p.name = "old_name"
p.auth_type = "api_key" p.auth_type = "api_key"
p.encrypted_credentials = {"sk": "e"} p.encrypted_credentials = {"sk": "e"}
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 1] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 1
with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): with patch("services.datasource_provider_service.get_current_user", return_value=mock_user):
with pytest.raises(ValueError, match="already exists"): with pytest.raises(ValueError, match="already exists"):
service.update_datasource_credentials("t1", "id", "prov", "org/plug", {}, "new_name") service.update_datasource_credentials("t1", "id", "prov", "org/plug", {}, "new_name")
@ -717,8 +708,7 @@ class TestDatasourceProviderService:
p.name = "old_name" p.name = "old_name"
p.auth_type = "api_key" p.auth_type = "api_key"
p.encrypted_credentials = {"sk": "e"} p.encrypted_credentials = {"sk": "e"}
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 0] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 0
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service, "extract_secret_variables", return_value=["sk"]), patch.object(service, "extract_secret_variables", return_value=["sk"]),
@ -733,8 +723,7 @@ class TestDatasourceProviderService:
p.name = "old_name" p.name = "old_name"
p.auth_type = "api_key" p.auth_type = "api_key"
p.encrypted_credentials = {"sk": "old_enc"} p.encrypted_credentials = {"sk": "old_enc"}
mock_db_session.query().first.return_value = p mock_db_session.scalar.side_effect = [p, 0] # first: fetch provider, second: name conflict count
mock_db_session.query().count.return_value = 0
with ( with (
patch("services.datasource_provider_service.get_current_user", return_value=mock_user), patch("services.datasource_provider_service.get_current_user", return_value=mock_user),
patch.object(service, "extract_secret_variables", return_value=["sk"]), patch.object(service, "extract_secret_variables", return_value=["sk"]),