mirror of
				https://github.com/open-metadata/OpenMetadata.git
				synced 2025-11-03 20:19:31 +00:00 
			
		
		
		
	Fix bot creation and edition logic (#7796)
* Fix bot creation and edition logic * Fix minor error creating user * Fix failing test * Minor fixes * Add missing tests for new flows * Fix put_failIfBotHasARelationshipToAnotherUser test * Changes after manual testing * Move where auth_provider is retrieved in the secret manager
This commit is contained in:
		
							parent
							
								
									16e8778993
								
							
						
					
					
						commit
						593ca3a4a0
					
				@ -17,6 +17,7 @@ working with OpenMetadata entities.
 | 
			
		||||
import traceback
 | 
			
		||||
from typing import Dict, Generic, Iterable, List, Optional, Type, TypeVar, Union
 | 
			
		||||
 | 
			
		||||
from metadata.generated.schema.entity.bot import BotType
 | 
			
		||||
from metadata.ingestion.ometa.mixins.dashboard_mixin import OMetaDashboardMixin
 | 
			
		||||
from metadata.ingestion.ometa.mixins.patch_mixin import OMetaPatchMixin
 | 
			
		||||
from metadata.ingestion.ometa.ssl_registry import (
 | 
			
		||||
@ -178,7 +179,9 @@ class OpenMetadata(
 | 
			
		||||
        )
 | 
			
		||||
 | 
			
		||||
        # Load auth provider config from Secret Manager if necessary
 | 
			
		||||
        self.secrets_manager_client.add_auth_provider_security_config(self.config)
 | 
			
		||||
        self.secrets_manager_client.add_auth_provider_security_config(
 | 
			
		||||
            self.config, BotType.ingestion_bot.value
 | 
			
		||||
        )
 | 
			
		||||
 | 
			
		||||
        # Load the auth provider init from the registry
 | 
			
		||||
        auth_provider_fn = auth_provider_registry.registry.get(
 | 
			
		||||
 | 
			
		||||
@ -33,6 +33,7 @@ from metadata.generated.schema.security.client.openMetadataJWTClientConfig impor
 | 
			
		||||
from metadata.utils.logger import utils_logger
 | 
			
		||||
from metadata.utils.secrets.secrets_manager import (
 | 
			
		||||
    AUTH_PROVIDER_MAPPING,
 | 
			
		||||
    AUTH_PROVIDER_PREFIX,
 | 
			
		||||
    BOT_PREFIX,
 | 
			
		||||
    DBT_SOURCE_CONFIG_SECRET_PREFIX,
 | 
			
		||||
    TEST_CONNECTION_TEMP_SECRET_PREFIX,
 | 
			
		||||
@ -84,7 +85,9 @@ class ExternalSecretsManager(SecretsManager, ABC):
 | 
			
		||||
        )
 | 
			
		||||
        return ServiceConnection(__root__=service_connection)
 | 
			
		||||
 | 
			
		||||
    def add_auth_provider_security_config(self, config: OpenMetadataConnection) -> None:
 | 
			
		||||
    def add_auth_provider_security_config(
 | 
			
		||||
        self, config: OpenMetadataConnection, bot_name: str
 | 
			
		||||
    ) -> None:
 | 
			
		||||
        """
 | 
			
		||||
        Add the auth provider security config from the AWS client store to a given OpenMetadata connection object.
 | 
			
		||||
        :param config: OpenMetadataConnection object
 | 
			
		||||
@ -93,18 +96,20 @@ class ExternalSecretsManager(SecretsManager, ABC):
 | 
			
		||||
        logger.debug(
 | 
			
		||||
            f"Adding auth provider security config using {self.provider} secrets' manager"
 | 
			
		||||
        )
 | 
			
		||||
        if config.authProvider != AuthProvider.no_auth:
 | 
			
		||||
            secret_id = self.build_secret_id(BOT_PREFIX, config.workflowBot)
 | 
			
		||||
 | 
			
		||||
        if (
 | 
			
		||||
            config.authProvider != AuthProvider.no_auth
 | 
			
		||||
            and config.securityConfig is None
 | 
			
		||||
        ):
 | 
			
		||||
            auth_provider_secret_id = self.build_secret_id(
 | 
			
		||||
                BOT_PREFIX, bot_name, AUTH_PROVIDER_PREFIX
 | 
			
		||||
            )
 | 
			
		||||
            auth_provider_secret = self.get_string_value(auth_provider_secret_id)
 | 
			
		||||
            config.authProvider = AuthProvider(json.loads(auth_provider_secret))
 | 
			
		||||
            secret_id = self.build_secret_id(BOT_PREFIX, bot_name)
 | 
			
		||||
            auth_config_json = self.get_string_value(secret_id)
 | 
			
		||||
            try:
 | 
			
		||||
                config_object = json.loads(auth_config_json)
 | 
			
		||||
                if config.authProvider == AuthProvider.openmetadata:
 | 
			
		||||
                    auth_mechanism: JWTAuthMechanism = JWTAuthMechanism.parse_obj(
 | 
			
		||||
                        config_object
 | 
			
		||||
                    )
 | 
			
		||||
                    config_object = OpenMetadataJWTClientConfig(
 | 
			
		||||
                        jwtToken=auth_mechanism.JWTToken
 | 
			
		||||
                    ).dict()
 | 
			
		||||
                config.securityConfig = AUTH_PROVIDER_MAPPING.get(
 | 
			
		||||
                    config.authProvider
 | 
			
		||||
                ).parse_obj(config_object)
 | 
			
		||||
 | 
			
		||||
@ -38,7 +38,7 @@ class NoopSecretsManager(SecretsManager):
 | 
			
		||||
    provider: str = SecretsManagerProvider.noop.name
 | 
			
		||||
 | 
			
		||||
    def add_auth_provider_security_config(
 | 
			
		||||
        self, open_metadata_connection: OpenMetadataConnection
 | 
			
		||||
        self, open_metadata_connection: OpenMetadataConnection, bot_name: str
 | 
			
		||||
    ) -> None:
 | 
			
		||||
        """
 | 
			
		||||
        The LocalSecretsManager does not modify the OpenMetadataConnection object
 | 
			
		||||
 | 
			
		||||
@ -98,6 +98,8 @@ DBT_SOURCE_CONFIG_SECRET_PREFIX: str = "database-metadata-pipeline"
 | 
			
		||||
 | 
			
		||||
BOT_PREFIX: str = "bot"
 | 
			
		||||
 | 
			
		||||
AUTH_PROVIDER_PREFIX: str = "auth-provider"
 | 
			
		||||
 | 
			
		||||
TEST_CONNECTION_TEMP_SECRET_PREFIX: str = "test-connection-temp"
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@ -128,10 +130,13 @@ class SecretsManager(metaclass=Singleton):
 | 
			
		||||
        pass
 | 
			
		||||
 | 
			
		||||
    @abstractmethod
 | 
			
		||||
    def add_auth_provider_security_config(self, config: OpenMetadataConnection) -> None:
 | 
			
		||||
    def add_auth_provider_security_config(
 | 
			
		||||
        self, config: OpenMetadataConnection, bot_name: str
 | 
			
		||||
    ) -> None:
 | 
			
		||||
        """
 | 
			
		||||
        Add the auth provider security config from the secret manager to a given OpenMetadata connection object.
 | 
			
		||||
        :param config: OpenMetadataConnection object
 | 
			
		||||
        :param bot_name: Bot name with the credentials
 | 
			
		||||
        """
 | 
			
		||||
        pass
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@ -17,6 +17,7 @@ from copy import deepcopy
 | 
			
		||||
from typing import Any, Dict
 | 
			
		||||
from unittest.mock import Mock, patch
 | 
			
		||||
 | 
			
		||||
from metadata.generated.schema.entity.bot import BotType
 | 
			
		||||
from metadata.generated.schema.entity.services.connections.serviceConnection import (
 | 
			
		||||
    ServiceConnection,
 | 
			
		||||
)
 | 
			
		||||
@ -79,15 +80,21 @@ class AWSBasedSecretsManager(object):
 | 
			
		||||
        @patch("metadata.clients.aws_client.AWSClient.get_client")
 | 
			
		||||
        def test_aws_manager_add_auth_provider_security_config(self, mocked_get_client):
 | 
			
		||||
            aws_manager = self.build_secret_manager(
 | 
			
		||||
                mocked_get_client, self.build_response_value(AUTH_PROVIDER_CONFIG)
 | 
			
		||||
                mocked_get_client,
 | 
			
		||||
                self.build_response_value("google"),
 | 
			
		||||
                self.build_response_value(AUTH_PROVIDER_CONFIG),
 | 
			
		||||
            )
 | 
			
		||||
            actual_om_connection = deepcopy(self.om_connection)
 | 
			
		||||
            actual_om_connection.securityConfig = None
 | 
			
		||||
 | 
			
		||||
            aws_manager.add_auth_provider_security_config(actual_om_connection)
 | 
			
		||||
            aws_manager.add_auth_provider_security_config(
 | 
			
		||||
                actual_om_connection, BotType.ingestion_bot.value
 | 
			
		||||
            )
 | 
			
		||||
 | 
			
		||||
            self.assert_client_called_once(
 | 
			
		||||
                aws_manager, "/openmetadata/bot/ingestion-bot"
 | 
			
		||||
                aws_manager,
 | 
			
		||||
                "/openmetadata/bot/ingestion-bot/auth-provider",
 | 
			
		||||
                "/openmetadata/bot/ingestion-bot",
 | 
			
		||||
            )
 | 
			
		||||
            self.assertEqual(
 | 
			
		||||
                self.auth_provider_config, actual_om_connection.securityConfig
 | 
			
		||||
@ -103,7 +110,9 @@ class AWSBasedSecretsManager(object):
 | 
			
		||||
            aws_manager = self.build_secret_manager(mocked_get_client, {})
 | 
			
		||||
 | 
			
		||||
            with self.assertRaises(ValueError) as value_error:
 | 
			
		||||
                aws_manager.add_auth_provider_security_config(self.om_connection)
 | 
			
		||||
                aws_manager.add_auth_provider_security_config(
 | 
			
		||||
                    self.om_connection, BotType.ingestion_bot.value
 | 
			
		||||
                )
 | 
			
		||||
            self.assertTrue(
 | 
			
		||||
                "/openmetadata/bot/ingestion-bot" in str(value_error.exception)
 | 
			
		||||
            )
 | 
			
		||||
@ -186,18 +195,23 @@ class AWSBasedSecretsManager(object):
 | 
			
		||||
 | 
			
		||||
        @abstractmethod
 | 
			
		||||
        def build_secret_manager(
 | 
			
		||||
            self, mocked_get_client: Mock, expected_json: Dict[str, Any]
 | 
			
		||||
            self,
 | 
			
		||||
            mocked_get_client: Mock,
 | 
			
		||||
            expected_json_2: Dict[str, Any],
 | 
			
		||||
            expected_json_1: Dict[str, Any],
 | 
			
		||||
        ) -> AWSBasedSecretsManager:
 | 
			
		||||
            pass
 | 
			
		||||
 | 
			
		||||
        @staticmethod
 | 
			
		||||
        @abstractmethod
 | 
			
		||||
        def assert_client_called_once(
 | 
			
		||||
            aws_manager: AWSBasedSecretsManager, expected_call: str
 | 
			
		||||
            aws_manager: AWSBasedSecretsManager,
 | 
			
		||||
            expected_call_1: str,
 | 
			
		||||
            expected_call_2: str,
 | 
			
		||||
        ) -> None:
 | 
			
		||||
            pass
 | 
			
		||||
 | 
			
		||||
        @staticmethod
 | 
			
		||||
        @abstractmethod
 | 
			
		||||
        def build_response_value(json_value: dict):
 | 
			
		||||
        def build_response_value(json_value: Any):
 | 
			
		||||
            pass
 | 
			
		||||
 | 
			
		||||
@ -14,7 +14,7 @@ Test AWS Secrets Manager
 | 
			
		||||
"""
 | 
			
		||||
import json
 | 
			
		||||
from abc import ABC
 | 
			
		||||
from typing import Any, Dict
 | 
			
		||||
from typing import Any, Dict, List
 | 
			
		||||
from unittest.mock import Mock
 | 
			
		||||
 | 
			
		||||
from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
 | 
			
		||||
@ -25,9 +25,14 @@ from .test_aws_based_secrets_manager import AWSBasedSecretsManager
 | 
			
		||||
 | 
			
		||||
class TestAWSSecretsManager(AWSBasedSecretsManager.TestCase, ABC):
 | 
			
		||||
    def build_secret_manager(
 | 
			
		||||
        self, mocked_get_client: Mock, expected_json: Dict[str, Any]
 | 
			
		||||
        self,
 | 
			
		||||
        mocked_get_client: Mock,
 | 
			
		||||
        expected_json_1: Dict[str, Any],
 | 
			
		||||
        expected_json_2: Dict[str, Any] = None,
 | 
			
		||||
    ) -> AWSSecretsManager:
 | 
			
		||||
        self.init_mocked_get_client(mocked_get_client, expected_json)
 | 
			
		||||
        self.init_mocked_get_client(
 | 
			
		||||
            mocked_get_client, [expected_json_1, expected_json_2]
 | 
			
		||||
        )
 | 
			
		||||
        return AWSSecretsManager(
 | 
			
		||||
            AWSCredentials(
 | 
			
		||||
                awsAccessKeyId="fake_key",
 | 
			
		||||
@ -39,19 +44,24 @@ class TestAWSSecretsManager(AWSBasedSecretsManager.TestCase, ABC):
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def init_mocked_get_client(
 | 
			
		||||
        get_client_mock: Mock, client_return: Dict[str, Any]
 | 
			
		||||
        get_client_mock: Mock, client_return: List[Dict[str, Any]]
 | 
			
		||||
    ) -> None:
 | 
			
		||||
        mocked_secret_manager = Mock()
 | 
			
		||||
        mocked_secret_manager.get_secret_value = Mock(return_value=client_return)
 | 
			
		||||
        mocked_secret_manager.get_secret_value = Mock(side_effect=client_return)
 | 
			
		||||
        get_client_mock.return_value = mocked_secret_manager
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def assert_client_called_once(
 | 
			
		||||
        aws_manager: AWSSecretsManager, expected_call: str
 | 
			
		||||
        aws_manager: AWSSecretsManager,
 | 
			
		||||
        expected_call_1: str,
 | 
			
		||||
        expected_call_2: str = None,
 | 
			
		||||
    ) -> None:
 | 
			
		||||
        expected_call = {"SecretId": expected_call}
 | 
			
		||||
        aws_manager.client.get_secret_value.assert_called_once_with(**expected_call)
 | 
			
		||||
        expected_call = {"SecretId": expected_call_1}
 | 
			
		||||
        aws_manager.client.get_secret_value.assert_any_call(**expected_call)
 | 
			
		||||
        if expected_call_2:
 | 
			
		||||
            expected_call = {"SecretId": expected_call_2}
 | 
			
		||||
            aws_manager.client.get_secret_value.assert_any_call(**expected_call)
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def build_response_value(json_value: dict):
 | 
			
		||||
    def build_response_value(json_value: Any):
 | 
			
		||||
        return {"SecretString": json.dumps(json_value)}
 | 
			
		||||
 | 
			
		||||
@ -14,7 +14,7 @@ Test AWS SSM Secrets Manager
 | 
			
		||||
"""
 | 
			
		||||
import json
 | 
			
		||||
from abc import ABC
 | 
			
		||||
from typing import Any, Dict
 | 
			
		||||
from typing import Any, Dict, List
 | 
			
		||||
from unittest.mock import Mock
 | 
			
		||||
 | 
			
		||||
from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
 | 
			
		||||
@ -26,9 +26,14 @@ from .test_aws_based_secrets_manager import AWSBasedSecretsManager
 | 
			
		||||
 | 
			
		||||
class TestAWSSecretsManager(AWSBasedSecretsManager.TestCase, ABC):
 | 
			
		||||
    def build_secret_manager(
 | 
			
		||||
        self, mocked_get_client: Mock, expected_json: Dict[str, Any]
 | 
			
		||||
        self,
 | 
			
		||||
        mocked_get_client: Mock,
 | 
			
		||||
        expected_json_1: Dict[str, Any],
 | 
			
		||||
        expected_json_2: Dict[str, Any] = None,
 | 
			
		||||
    ) -> AWSSSMSecretsManager:
 | 
			
		||||
        self.init_mocked_get_client(mocked_get_client, expected_json)
 | 
			
		||||
        self.init_mocked_get_client(
 | 
			
		||||
            mocked_get_client, [expected_json_1, expected_json_2]
 | 
			
		||||
        )
 | 
			
		||||
        return AWSSSMSecretsManager(
 | 
			
		||||
            AWSCredentials(
 | 
			
		||||
                awsAccessKeyId="fake_key",
 | 
			
		||||
@ -39,17 +44,24 @@ class TestAWSSecretsManager(AWSBasedSecretsManager.TestCase, ABC):
 | 
			
		||||
        )
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def init_mocked_get_client(get_client_mock: Mock, client_return: Dict[str, Any]):
 | 
			
		||||
    def init_mocked_get_client(
 | 
			
		||||
        get_client_mock: Mock, client_return: List[Dict[str, Any]]
 | 
			
		||||
    ):
 | 
			
		||||
        mocked_secret_manager = Mock()
 | 
			
		||||
        mocked_secret_manager.get_parameter = Mock(return_value=client_return)
 | 
			
		||||
        mocked_secret_manager.get_parameter = Mock(side_effect=client_return)
 | 
			
		||||
        get_client_mock.return_value = mocked_secret_manager
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def assert_client_called_once(
 | 
			
		||||
        aws_manager: AWSSecretsManager, expected_call: str
 | 
			
		||||
        aws_manager: AWSSecretsManager,
 | 
			
		||||
        expected_call_1: str,
 | 
			
		||||
        expected_call_2: str = None,
 | 
			
		||||
    ) -> None:
 | 
			
		||||
        expected_call = {"Name": expected_call, "WithDecryption": True}
 | 
			
		||||
        aws_manager.client.get_parameter.assert_called_once_with(**expected_call)
 | 
			
		||||
        expected_call = {"Name": expected_call_1, "WithDecryption": True}
 | 
			
		||||
        aws_manager.client.get_parameter.assert_any_call(**expected_call)
 | 
			
		||||
        if expected_call_2:
 | 
			
		||||
            expected_call = {"Name": expected_call_2, "WithDecryption": True}
 | 
			
		||||
            aws_manager.client.get_parameter.assert_any_call(**expected_call)
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def build_response_value(json_value: dict):
 | 
			
		||||
 | 
			
		||||
@ -14,6 +14,7 @@ Test Local Secrets Manager
 | 
			
		||||
"""
 | 
			
		||||
from copy import deepcopy
 | 
			
		||||
 | 
			
		||||
from metadata.generated.schema.entity.bot import BotType
 | 
			
		||||
from metadata.generated.schema.entity.services.connections.metadata.secretsManagerProvider import (
 | 
			
		||||
    SecretsManagerProvider,
 | 
			
		||||
)
 | 
			
		||||
@ -55,7 +56,9 @@ class TestLocalSecretsManager(TestSecretsManager.External):
 | 
			
		||||
        actual_om_connection = deepcopy(self.om_connection)
 | 
			
		||||
        actual_om_connection.securityConfig = self.auth_provider_config
 | 
			
		||||
 | 
			
		||||
        noop_manager.add_auth_provider_security_config(actual_om_connection)
 | 
			
		||||
        noop_manager.add_auth_provider_security_config(
 | 
			
		||||
            actual_om_connection, BotType.ingestion_bot.value
 | 
			
		||||
        )
 | 
			
		||||
 | 
			
		||||
        self.assertEqual(self.auth_provider_config, actual_om_connection.securityConfig)
 | 
			
		||||
        assert id(self.auth_provider_config) == id(actual_om_connection.securityConfig)
 | 
			
		||||
 | 
			
		||||
@ -34,7 +34,6 @@ import org.openmetadata.schema.api.slackChat.SlackChatConfiguration;
 | 
			
		||||
import org.openmetadata.schema.email.SmtpSettings;
 | 
			
		||||
import org.openmetadata.service.migration.MigrationConfiguration;
 | 
			
		||||
import org.openmetadata.service.secrets.SecretsManagerConfiguration;
 | 
			
		||||
import org.openmetadata.service.validators.AirflowConfigValidation;
 | 
			
		||||
 | 
			
		||||
@Getter
 | 
			
		||||
@Setter
 | 
			
		||||
@ -62,7 +61,6 @@ public class OpenMetadataApplicationConfig extends Configuration {
 | 
			
		||||
  @JsonProperty("eventHandlerConfiguration")
 | 
			
		||||
  private EventHandlerConfiguration eventHandlerConfiguration;
 | 
			
		||||
 | 
			
		||||
  @AirflowConfigValidation
 | 
			
		||||
  @NotNull
 | 
			
		||||
  @Valid
 | 
			
		||||
  @JsonProperty("airflowConfiguration")
 | 
			
		||||
 | 
			
		||||
@ -15,17 +15,25 @@ package org.openmetadata.service.jdbi3;
 | 
			
		||||
 | 
			
		||||
import java.io.IOException;
 | 
			
		||||
import org.openmetadata.schema.entity.Bot;
 | 
			
		||||
import org.openmetadata.schema.entity.BotType;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.User;
 | 
			
		||||
import org.openmetadata.schema.type.EntityReference;
 | 
			
		||||
import org.openmetadata.schema.type.Include;
 | 
			
		||||
import org.openmetadata.schema.type.Relationship;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.resources.bots.BotResource;
 | 
			
		||||
import org.openmetadata.service.secrets.SecretsManager;
 | 
			
		||||
import org.openmetadata.service.util.EntityUtil.Fields;
 | 
			
		||||
 | 
			
		||||
public class BotRepository extends EntityRepository<Bot> {
 | 
			
		||||
  public BotRepository(CollectionDAO dao) {
 | 
			
		||||
    super(BotResource.COLLECTION_PATH, Entity.BOT, Bot.class, dao.botDAO(), dao, "", "");
 | 
			
		||||
 | 
			
		||||
  static final String BOT_UPDATE_FIELDS = "botUser";
 | 
			
		||||
 | 
			
		||||
  SecretsManager secretsManager;
 | 
			
		||||
 | 
			
		||||
  public BotRepository(CollectionDAO dao, SecretsManager secretsManager) {
 | 
			
		||||
    super(BotResource.COLLECTION_PATH, Entity.BOT, Bot.class, dao.botDAO(), dao, "", BOT_UPDATE_FIELDS);
 | 
			
		||||
    this.secretsManager = secretsManager;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
@ -45,6 +53,9 @@ public class BotRepository extends EntityRepository<Bot> {
 | 
			
		||||
    EntityReference botUser = entity.getBotUser();
 | 
			
		||||
    entity.withBotUser(null);
 | 
			
		||||
    store(entity.getId(), entity, update);
 | 
			
		||||
    if (!BotType.BOT.equals(entity.getBotType())) {
 | 
			
		||||
      secretsManager.encryptOrDecryptBotCredentials(entity.getBotType().value(), botUser.getName(), true);
 | 
			
		||||
    }
 | 
			
		||||
    entity.withBotUser(botUser);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -72,5 +83,23 @@ public class BotRepository extends EntityRepository<Bot> {
 | 
			
		||||
    public BotUpdater(Bot original, Bot updated, Operation operation) {
 | 
			
		||||
      super(original, updated, operation);
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    @Override
 | 
			
		||||
    public void entitySpecificUpdate() throws IOException {
 | 
			
		||||
      updateUser(original, updated);
 | 
			
		||||
      if (original.getBotType() != null) {
 | 
			
		||||
        updated.setBotType(original.getBotType());
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    private void updateUser(Bot original, Bot updated) throws IOException {
 | 
			
		||||
      deleteTo(original.getBotUser().getId(), Entity.USER, Relationship.CONTAINS, Entity.BOT);
 | 
			
		||||
      addRelationship(updated.getId(), updated.getBotUser().getId(), Entity.BOT, Entity.USER, Relationship.CONTAINS);
 | 
			
		||||
      if (original.getBotUser() == null
 | 
			
		||||
          || updated.getBotUser() == null
 | 
			
		||||
          || !updated.getBotUser().getId().equals(original.getBotUser().getId())) {
 | 
			
		||||
        recordChange("botUser", original.getBotUser(), updated.getBotUser());
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -119,7 +119,7 @@ public class UserRepository extends EntityRepository<User> {
 | 
			
		||||
    if (secretsManager != null && Boolean.TRUE.equals(user.getIsBot()) && user.getAuthenticationMechanism() != null) {
 | 
			
		||||
      user.getAuthenticationMechanism()
 | 
			
		||||
          .setConfig(
 | 
			
		||||
              secretsManager.encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
              secretsManager.encryptOrDecryptBotUserCredentials(
 | 
			
		||||
                  user.getName(), user.getAuthenticationMechanism().getConfig(), true));
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
@ -233,7 +233,7 @@ public class UserRepository extends EntityRepository<User> {
 | 
			
		||||
    if (user.getAuthenticationMechanism() != null) {
 | 
			
		||||
      user.getAuthenticationMechanism()
 | 
			
		||||
          .withConfig(
 | 
			
		||||
              this.secretsManager.encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
              this.secretsManager.encryptOrDecryptBotUserCredentials(
 | 
			
		||||
                  user.getName(), user.getAuthenticationMechanism().getConfig(), false));
 | 
			
		||||
    }
 | 
			
		||||
    return user;
 | 
			
		||||
@ -349,15 +349,11 @@ public class UserRepository extends EntityRepository<User> {
 | 
			
		||||
      AuthenticationMechanism updatedAuthMechanism = updated.getAuthenticationMechanism();
 | 
			
		||||
      if (origAuthMechanism == null && updatedAuthMechanism != null) {
 | 
			
		||||
        recordChange("authenticationMechanism", original.getAuthenticationMechanism(), "new-encrypted-value");
 | 
			
		||||
      } else if (hasConfig(origAuthMechanism) && hasConfig(updatedAuthMechanism)) {
 | 
			
		||||
      } else if (origAuthMechanism != null && updatedAuthMechanism != null) {
 | 
			
		||||
        if (!JsonUtils.areEquals(origAuthMechanism, updatedAuthMechanism)) {
 | 
			
		||||
          recordChange("authenticationMechanism", "old-encrypted-value", "new-encrypted-value");
 | 
			
		||||
        }
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    private boolean hasConfig(AuthenticationMechanism authenticationMechanism) {
 | 
			
		||||
      return authenticationMechanism != null && authenticationMechanism.getConfig() != null;
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -47,15 +47,21 @@ import javax.ws.rs.core.SecurityContext;
 | 
			
		||||
import javax.ws.rs.core.UriInfo;
 | 
			
		||||
import org.openmetadata.schema.api.CreateBot;
 | 
			
		||||
import org.openmetadata.schema.entity.Bot;
 | 
			
		||||
import org.openmetadata.schema.entity.BotType;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.User;
 | 
			
		||||
import org.openmetadata.schema.type.EntityHistory;
 | 
			
		||||
import org.openmetadata.schema.type.Include;
 | 
			
		||||
import org.openmetadata.schema.type.Relationship;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.jdbi3.BotRepository;
 | 
			
		||||
import org.openmetadata.service.jdbi3.CollectionDAO;
 | 
			
		||||
import org.openmetadata.service.jdbi3.ListFilter;
 | 
			
		||||
import org.openmetadata.service.jdbi3.UserRepository;
 | 
			
		||||
import org.openmetadata.service.resources.Collection;
 | 
			
		||||
import org.openmetadata.service.resources.EntityResource;
 | 
			
		||||
import org.openmetadata.service.secrets.SecretsManager;
 | 
			
		||||
import org.openmetadata.service.security.Authorizer;
 | 
			
		||||
import org.openmetadata.service.util.EntityUtil;
 | 
			
		||||
import org.openmetadata.service.util.ResultList;
 | 
			
		||||
 | 
			
		||||
@Path("/v1/bots")
 | 
			
		||||
@ -66,8 +72,11 @@ import org.openmetadata.service.util.ResultList;
 | 
			
		||||
public class BotResource extends EntityResource<Bot, BotRepository> {
 | 
			
		||||
  public static final String COLLECTION_PATH = "/v1/bots/";
 | 
			
		||||
 | 
			
		||||
  public BotResource(CollectionDAO dao, Authorizer authorizer) {
 | 
			
		||||
    super(Bot.class, new BotRepository(dao), authorizer);
 | 
			
		||||
  SecretsManager secretsManager;
 | 
			
		||||
 | 
			
		||||
  public BotResource(CollectionDAO dao, Authorizer authorizer, SecretsManager secretsManager) {
 | 
			
		||||
    super(Bot.class, new BotRepository(dao, secretsManager), authorizer);
 | 
			
		||||
    this.secretsManager = secretsManager;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
@ -236,7 +245,7 @@ public class BotResource extends EntityResource<Bot, BotRepository> {
 | 
			
		||||
      })
 | 
			
		||||
  public Response create(@Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateBot create)
 | 
			
		||||
      throws IOException {
 | 
			
		||||
    Bot bot = getBot(create, securityContext.getUserPrincipal().getName());
 | 
			
		||||
    Bot bot = getBot(securityContext, create);
 | 
			
		||||
    return create(uriInfo, securityContext, bot, false);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -255,8 +264,14 @@ public class BotResource extends EntityResource<Bot, BotRepository> {
 | 
			
		||||
      })
 | 
			
		||||
  public Response createOrUpdate(
 | 
			
		||||
      @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateBot create) throws IOException {
 | 
			
		||||
    Bot bot = getBot(create, securityContext.getUserPrincipal().getName());
 | 
			
		||||
    return createOrUpdate(uriInfo, securityContext, bot, false);
 | 
			
		||||
    Bot bot = getBot(securityContext, create);
 | 
			
		||||
    Response response = createOrUpdate(uriInfo, securityContext, bot, false);
 | 
			
		||||
    // ensures the secrets' manager store the credentials even when the botUser does not change
 | 
			
		||||
    bot = (Bot) response.getEntity();
 | 
			
		||||
    if (!BotType.BOT.equals(bot.getBotType())) {
 | 
			
		||||
      secretsManager.encryptOrDecryptBotCredentials(bot.getBotType().value(), bot.getBotUser().getName(), true);
 | 
			
		||||
    }
 | 
			
		||||
    return response;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @PATCH
 | 
			
		||||
@ -305,10 +320,66 @@ public class BotResource extends EntityResource<Bot, BotRepository> {
 | 
			
		||||
          boolean hardDelete,
 | 
			
		||||
      @Parameter(description = "Id of the Bot", schema = @Schema(type = "UUID")) @PathParam("id") UUID id)
 | 
			
		||||
      throws IOException {
 | 
			
		||||
    BotType botType = dao.get(null, id, EntityUtil.Fields.EMPTY_FIELDS).getBotType();
 | 
			
		||||
    if (!BotType.BOT.equals(botType)) {
 | 
			
		||||
      throw new IllegalArgumentException(String.format("[%s] can not be deleted.", botType.value()));
 | 
			
		||||
    }
 | 
			
		||||
    return delete(uriInfo, securityContext, id, true, hardDelete, false);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private Bot getBot(CreateBot create, String user) throws IOException {
 | 
			
		||||
    return copy(new Bot(), create, user).withBotUser(create.getBotUser());
 | 
			
		||||
    return copy(new Bot(), create, user)
 | 
			
		||||
        .withBotUser(create.getBotUser())
 | 
			
		||||
        .withBotType(BotType.BOT)
 | 
			
		||||
        .withFullyQualifiedName(create.getName());
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private boolean userHasRelationshipWithAnyBot(User user, Bot botUser) {
 | 
			
		||||
    if (user == null) {
 | 
			
		||||
      return false;
 | 
			
		||||
    }
 | 
			
		||||
    List<CollectionDAO.EntityRelationshipRecord> userBotRelationship = retrieveBotRelationshipsFor(user);
 | 
			
		||||
    return !userBotRelationship.isEmpty()
 | 
			
		||||
        && (botUser == null
 | 
			
		||||
            || userBotRelationship.stream().anyMatch(relationship -> !relationship.getId().equals(botUser.getId())));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private List<CollectionDAO.EntityRelationshipRecord> retrieveBotRelationshipsFor(User user) {
 | 
			
		||||
    return dao.findFrom(user.getId(), Entity.USER, Relationship.CONTAINS, Entity.BOT);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private Bot getBot(SecurityContext securityContext, CreateBot create) throws IOException {
 | 
			
		||||
    Bot bot = getBot(create, securityContext.getUserPrincipal().getName());
 | 
			
		||||
    Bot originalBot = retrieveBot(bot.getName());
 | 
			
		||||
    User botUser = retrieveUser(bot);
 | 
			
		||||
    if (botUser != null && !Boolean.TRUE.equals(botUser.getIsBot())) {
 | 
			
		||||
      throw new IllegalArgumentException(String.format("User [%s] is not a bot user", botUser.getName()));
 | 
			
		||||
    }
 | 
			
		||||
    if (userHasRelationshipWithAnyBot(botUser, originalBot)) {
 | 
			
		||||
      List<CollectionDAO.EntityRelationshipRecord> userBotRelationship = retrieveBotRelationshipsFor(botUser);
 | 
			
		||||
      bot =
 | 
			
		||||
          dao.get(null, userBotRelationship.stream().findFirst().orElseThrow().getId(), EntityUtil.Fields.EMPTY_FIELDS);
 | 
			
		||||
      throw new IllegalArgumentException(
 | 
			
		||||
          String.format("Bot user [%s] is already used by [%s] bot", botUser.getName(), bot.getName()));
 | 
			
		||||
    }
 | 
			
		||||
    return bot;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private User retrieveUser(Bot bot) {
 | 
			
		||||
    try {
 | 
			
		||||
      return UserRepository.class
 | 
			
		||||
          .cast(Entity.getEntityRepository(Entity.USER))
 | 
			
		||||
          .get(null, bot.getBotUser().getId(), EntityUtil.Fields.EMPTY_FIELDS);
 | 
			
		||||
    } catch (Exception exception) {
 | 
			
		||||
      return null;
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private Bot retrieveBot(String botName) {
 | 
			
		||||
    try {
 | 
			
		||||
      return dao.getByName(null, botName, EntityUtil.Fields.EMPTY_FIELDS);
 | 
			
		||||
    } catch (Exception e) {
 | 
			
		||||
      return null;
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -88,7 +88,6 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
 | 
			
		||||
  private PipelineServiceClient pipelineServiceClient;
 | 
			
		||||
  private OpenMetadataApplicationConfig openMetadataApplicationConfig;
 | 
			
		||||
  private final SecretsManager secretsManager;
 | 
			
		||||
  private CollectionDAO collectionDAO;
 | 
			
		||||
 | 
			
		||||
  @Getter private final IngestionPipelineRepository ingestionPipelineRepository;
 | 
			
		||||
 | 
			
		||||
@ -101,7 +100,6 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
 | 
			
		||||
 | 
			
		||||
  public IngestionPipelineResource(CollectionDAO dao, Authorizer authorizer, SecretsManager secretsManager) {
 | 
			
		||||
    super(IngestionPipeline.class, new IngestionPipelineRepository(dao, secretsManager), authorizer);
 | 
			
		||||
    this.collectionDAO = dao;
 | 
			
		||||
    this.secretsManager = secretsManager;
 | 
			
		||||
    this.ingestionPipelineRepository = new IngestionPipelineRepository(dao, secretsManager);
 | 
			
		||||
  }
 | 
			
		||||
@ -404,6 +402,8 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
 | 
			
		||||
      @Context UriInfo uriInfo, @PathParam("id") UUID id, @Context SecurityContext securityContext) throws IOException {
 | 
			
		||||
    Fields fields = getFields(FIELD_OWNER);
 | 
			
		||||
    IngestionPipeline ingestionPipeline = dao.get(uriInfo, id, fields);
 | 
			
		||||
    ingestionPipeline.setOpenMetadataServerConnection(
 | 
			
		||||
        new OpenMetadataServerConnectionBuilder(secretsManager, openMetadataApplicationConfig).build());
 | 
			
		||||
    pipelineServiceClient.deployPipeline(ingestionPipeline);
 | 
			
		||||
    decryptOrNullify(securityContext, ingestionPipeline);
 | 
			
		||||
    return addHref(uriInfo, ingestionPipeline);
 | 
			
		||||
@ -589,7 +589,7 @@ public class IngestionPipelineResource extends EntityResource<IngestionPipeline,
 | 
			
		||||
 | 
			
		||||
  private IngestionPipeline getIngestionPipeline(CreateIngestionPipeline create, String user) throws IOException {
 | 
			
		||||
    OpenMetadataServerConnection openMetadataServerConnection =
 | 
			
		||||
        new OpenMetadataServerConnectionBuilder(secretsManager, openMetadataApplicationConfig, collectionDAO).build();
 | 
			
		||||
        new OpenMetadataServerConnectionBuilder(secretsManager, openMetadataApplicationConfig).build();
 | 
			
		||||
    return copy(new IngestionPipeline(), create, user)
 | 
			
		||||
        .withPipelineType(create.getPipelineType())
 | 
			
		||||
        .withAirflowConfig(create.getAirflowConfig())
 | 
			
		||||
 | 
			
		||||
@ -75,6 +75,8 @@ import javax.ws.rs.core.UriInfo;
 | 
			
		||||
import lombok.Getter;
 | 
			
		||||
import lombok.extern.slf4j.Slf4j;
 | 
			
		||||
import org.apache.commons.lang3.StringUtils;
 | 
			
		||||
import org.jetbrains.annotations.Nullable;
 | 
			
		||||
import org.openmetadata.schema.EntityInterface;
 | 
			
		||||
import org.openmetadata.schema.api.security.AuthenticationConfiguration;
 | 
			
		||||
import org.openmetadata.schema.api.security.AuthorizerConfiguration;
 | 
			
		||||
import org.openmetadata.schema.api.teams.CreateUser;
 | 
			
		||||
@ -99,6 +101,7 @@ import org.openmetadata.schema.type.EntityHistory;
 | 
			
		||||
import org.openmetadata.schema.type.EntityReference;
 | 
			
		||||
import org.openmetadata.schema.type.Include;
 | 
			
		||||
import org.openmetadata.schema.type.MetadataOperation;
 | 
			
		||||
import org.openmetadata.schema.type.Relationship;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.auth.JwtResponse;
 | 
			
		||||
import org.openmetadata.service.exception.CatalogExceptionMessage;
 | 
			
		||||
@ -235,8 +238,6 @@ public class UserResource extends EntityResource<User, UserRepository> {
 | 
			
		||||
          @DefaultValue("non-deleted")
 | 
			
		||||
          Include include)
 | 
			
		||||
      throws IOException {
 | 
			
		||||
    // remove USER_PROTECTED_FIELDS from fieldsParam
 | 
			
		||||
    fieldsParam = removeUserProtectedFields(fieldsParam);
 | 
			
		||||
    ListFilter filter = new ListFilter(include).addQueryParam("team", teamParam);
 | 
			
		||||
    if (isAdmin != null) {
 | 
			
		||||
      filter.addQueryParam("isAdmin", String.valueOf(isAdmin));
 | 
			
		||||
@ -314,8 +315,6 @@ public class UserResource extends EntityResource<User, UserRepository> {
 | 
			
		||||
          @DefaultValue("non-deleted")
 | 
			
		||||
          Include include)
 | 
			
		||||
      throws IOException {
 | 
			
		||||
    // remove USER_PROTECTED_FIELDS from fieldsParam
 | 
			
		||||
    fieldsParam = removeUserProtectedFields(fieldsParam);
 | 
			
		||||
    return decryptOrNullify(securityContext, getInternal(uriInfo, securityContext, id, fieldsParam, include));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -350,8 +349,6 @@ public class UserResource extends EntityResource<User, UserRepository> {
 | 
			
		||||
          @DefaultValue("non-deleted")
 | 
			
		||||
          Include include)
 | 
			
		||||
      throws IOException {
 | 
			
		||||
    // remove USER_PROTECTED_FIELDS from fieldsParam
 | 
			
		||||
    fieldsParam = removeUserProtectedFields(fieldsParam);
 | 
			
		||||
    return decryptOrNullify(securityContext, getByNameInternal(uriInfo, securityContext, name, fieldsParam, include));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -556,11 +553,10 @@ public class UserResource extends EntityResource<User, UserRepository> {
 | 
			
		||||
      authorizer.authorize(securityContext, createOperationContext, resourceContext, true);
 | 
			
		||||
    }
 | 
			
		||||
    if (Boolean.TRUE.equals(create.getIsBot())) {
 | 
			
		||||
      addAuthMechanismToBot(user, create, uriInfo);
 | 
			
		||||
      return createOrUpdateBot(user, create, uriInfo, securityContext);
 | 
			
		||||
    }
 | 
			
		||||
    RestUtil.PutResponse<User> response = dao.createOrUpdate(uriInfo, user);
 | 
			
		||||
    addHref(uriInfo, response.getEntity());
 | 
			
		||||
    decryptOrNullify(securityContext, response.getEntity());
 | 
			
		||||
    return response.toResponse();
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -1439,51 +1435,110 @@ public class UserResource extends EntityResource<User, UserRepository> {
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private Response createOrUpdateBot(User user, CreateUser create, UriInfo uriInfo, SecurityContext securityContext)
 | 
			
		||||
      throws IOException {
 | 
			
		||||
    User original = retrieveBotUser(user, uriInfo);
 | 
			
		||||
    String botName = create.getBotName();
 | 
			
		||||
    EntityInterface bot = retrieveBot(botName);
 | 
			
		||||
    // check if the bot user exists
 | 
			
		||||
    if (!botHasRelationshipWithUser(bot, original)) {
 | 
			
		||||
      // throw an exception if user already has a relationship with a bot
 | 
			
		||||
      if (original != null && userHasRelationshipWithAnyBot(original, bot)) {
 | 
			
		||||
        List<CollectionDAO.EntityRelationshipRecord> userBotRelationship = retrieveBotRelationshipsFor(original);
 | 
			
		||||
        bot =
 | 
			
		||||
            Entity.getEntityRepository(Entity.BOT)
 | 
			
		||||
                .get(null, userBotRelationship.stream().findFirst().orElseThrow().getId(), Fields.EMPTY_FIELDS);
 | 
			
		||||
        throw new IllegalArgumentException(
 | 
			
		||||
            String.format("Bot user [%s] is already used by [%s] bot.", user.getName(), bot.getName()));
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
    addAuthMechanismToBot(user, create, uriInfo);
 | 
			
		||||
    RestUtil.PutResponse<User> response = dao.createOrUpdate(uriInfo, user);
 | 
			
		||||
    decryptOrNullify(securityContext, response.getEntity());
 | 
			
		||||
    return response.toResponse();
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private EntityInterface retrieveBot(String botName) {
 | 
			
		||||
    try {
 | 
			
		||||
      return Entity.getEntityRepository(Entity.BOT).getByName(null, botName, Fields.EMPTY_FIELDS);
 | 
			
		||||
    } catch (Exception e) {
 | 
			
		||||
      return null;
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private boolean userHasRelationshipWithAnyBot(User user, EntityInterface botUser) {
 | 
			
		||||
    List<CollectionDAO.EntityRelationshipRecord> userBotRelationship = retrieveBotRelationshipsFor(user);
 | 
			
		||||
    return !userBotRelationship.isEmpty()
 | 
			
		||||
        && (botUser == null
 | 
			
		||||
            || (userBotRelationship.stream().anyMatch(relationship -> !relationship.getId().equals(botUser.getId()))));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private List<CollectionDAO.EntityRelationshipRecord> retrieveBotRelationshipsFor(User user) {
 | 
			
		||||
    return dao.findFrom(user.getId(), Entity.USER, Relationship.CONTAINS, Entity.BOT);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private boolean botHasRelationshipWithUser(EntityInterface bot, User user) {
 | 
			
		||||
    if (bot == null || user == null) {
 | 
			
		||||
      return false;
 | 
			
		||||
    }
 | 
			
		||||
    List<CollectionDAO.EntityRelationshipRecord> botUserRelationships = retrieveBotRelationshipsFor(bot);
 | 
			
		||||
    return !botUserRelationships.isEmpty() && botUserRelationships.get(0).getId().equals(user.getId());
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private List<CollectionDAO.EntityRelationshipRecord> retrieveBotRelationshipsFor(EntityInterface bot) {
 | 
			
		||||
    return dao.findTo(bot.getId(), Entity.BOT, Relationship.CONTAINS, Entity.USER);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private void addAuthMechanismToBot(User user, @Valid CreateUser create, UriInfo uriInfo) {
 | 
			
		||||
    if (!Boolean.TRUE.equals(user.getIsBot())) {
 | 
			
		||||
      throw new IllegalArgumentException("Authentication mechanism change is only supported for bot users");
 | 
			
		||||
    }
 | 
			
		||||
    if (!isValidAuthenticationMechanism(create)) {
 | 
			
		||||
    if (isValidAuthenticationMechanism(create)) {
 | 
			
		||||
      AuthenticationMechanism authMechanism = create.getAuthenticationMechanism();
 | 
			
		||||
      AuthenticationMechanism.AuthType authType = authMechanism.getAuthType();
 | 
			
		||||
      switch (authType) {
 | 
			
		||||
        case JWT:
 | 
			
		||||
          User original = retrieveBotUser(user, uriInfo);
 | 
			
		||||
          if (original != null && !secretsManager.isLocal() && authMechanism.getConfig() != null) {
 | 
			
		||||
            original
 | 
			
		||||
                .getAuthenticationMechanism()
 | 
			
		||||
                .setConfig(
 | 
			
		||||
                    secretsManager.encryptOrDecryptBotUserCredentials(
 | 
			
		||||
                        user.getName(), authMechanism.getConfig(), false));
 | 
			
		||||
          }
 | 
			
		||||
          if (original == null || !hasAJWTAuthMechanism(original.getAuthenticationMechanism())) {
 | 
			
		||||
            JWTAuthMechanism jwtAuthMechanism =
 | 
			
		||||
                JsonUtils.convertValue(authMechanism.getConfig(), JWTAuthMechanism.class);
 | 
			
		||||
            authMechanism.setConfig(jwtTokenGenerator.generateJWTToken(user, jwtAuthMechanism.getJWTTokenExpiry()));
 | 
			
		||||
          } else {
 | 
			
		||||
            authMechanism = original.getAuthenticationMechanism();
 | 
			
		||||
          }
 | 
			
		||||
          break;
 | 
			
		||||
        case SSO:
 | 
			
		||||
          SSOAuthMechanism ssoAuthMechanism = JsonUtils.convertValue(authMechanism.getConfig(), SSOAuthMechanism.class);
 | 
			
		||||
          authMechanism.setConfig(ssoAuthMechanism);
 | 
			
		||||
          break;
 | 
			
		||||
        default:
 | 
			
		||||
          throw new IllegalArgumentException(
 | 
			
		||||
              String.format("Not supported authentication mechanism type: [%s]", authType.value()));
 | 
			
		||||
      }
 | 
			
		||||
      user.setAuthenticationMechanism(authMechanism);
 | 
			
		||||
    } else {
 | 
			
		||||
      throw new IllegalArgumentException(
 | 
			
		||||
          String.format("Authentication mechanism is empty bot user: [%s]", user.getName()));
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
    AuthenticationMechanism authMechanism = create.getAuthenticationMechanism();
 | 
			
		||||
    AuthenticationMechanism.AuthType authType = authMechanism.getAuthType();
 | 
			
		||||
    switch (authType) {
 | 
			
		||||
      case JWT:
 | 
			
		||||
        User original;
 | 
			
		||||
        try {
 | 
			
		||||
          original =
 | 
			
		||||
              dao.getByName(
 | 
			
		||||
                  uriInfo, user.getFullyQualifiedName(), new EntityUtil.Fields(List.of("authenticationMechanism")));
 | 
			
		||||
        } catch (EntityNotFoundException | IOException exc) {
 | 
			
		||||
          LOG.debug(String.format("User not found when adding auth mechanism for: [%s]", user.getName()));
 | 
			
		||||
          original = null;
 | 
			
		||||
        }
 | 
			
		||||
        if (original != null && !secretsManager.isLocal()) {
 | 
			
		||||
          original
 | 
			
		||||
              .getAuthenticationMechanism()
 | 
			
		||||
              .setConfig(
 | 
			
		||||
                  secretsManager.encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
                      user.getName(), authMechanism.getConfig(), false));
 | 
			
		||||
        }
 | 
			
		||||
        if (original == null || !hasAJWTAuthMechanism(original.getAuthenticationMechanism())) {
 | 
			
		||||
          JWTAuthMechanism jwtAuthMechanism = JsonUtils.convertValue(authMechanism.getConfig(), JWTAuthMechanism.class);
 | 
			
		||||
          authMechanism.setConfig(jwtTokenGenerator.generateJWTToken(user, jwtAuthMechanism.getJWTTokenExpiry()));
 | 
			
		||||
        } else {
 | 
			
		||||
          authMechanism = original.getAuthenticationMechanism();
 | 
			
		||||
        }
 | 
			
		||||
        break;
 | 
			
		||||
      case SSO:
 | 
			
		||||
        SSOAuthMechanism ssoAuthMechanism = JsonUtils.convertValue(authMechanism.getConfig(), SSOAuthMechanism.class);
 | 
			
		||||
        authMechanism.setConfig(ssoAuthMechanism);
 | 
			
		||||
        break;
 | 
			
		||||
      default:
 | 
			
		||||
        throw new IllegalArgumentException(
 | 
			
		||||
            String.format("Not supported authentication mechanism type: [%s]", authType.value()));
 | 
			
		||||
  @Nullable
 | 
			
		||||
  private User retrieveBotUser(User user, UriInfo uriInfo) {
 | 
			
		||||
    User original;
 | 
			
		||||
    try {
 | 
			
		||||
      original = dao.getByName(uriInfo, user.getFullyQualifiedName(), new Fields(List.of("authenticationMechanism")));
 | 
			
		||||
    } catch (EntityNotFoundException | IOException exc) {
 | 
			
		||||
      LOG.debug(String.format("User not found when adding auth mechanism for: [%s]", user.getName()));
 | 
			
		||||
      original = null;
 | 
			
		||||
    }
 | 
			
		||||
    user.setAuthenticationMechanism(authMechanism);
 | 
			
		||||
    return original;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private void addAuthMechanismToUser(User user, @Valid CreateUser create) {
 | 
			
		||||
@ -1532,14 +1587,10 @@ public class UserResource extends EntityResource<User, UserRepository> {
 | 
			
		||||
      }
 | 
			
		||||
      user.getAuthenticationMechanism()
 | 
			
		||||
          .setConfig(
 | 
			
		||||
              secretsManager.encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
              secretsManager.encryptOrDecryptBotUserCredentials(
 | 
			
		||||
                  user.getName(), user.getAuthenticationMechanism().getConfig(), false));
 | 
			
		||||
      return user;
 | 
			
		||||
    }
 | 
			
		||||
    return user;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private String removeUserProtectedFields(String fieldsParam) {
 | 
			
		||||
    return fieldsParam != null ? fieldsParam.replace("," + USER_PROTECTED_FIELDS, "") : null;
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -13,13 +13,11 @@
 | 
			
		||||
 | 
			
		||||
package org.openmetadata.service.secrets;
 | 
			
		||||
 | 
			
		||||
import java.io.IOException;
 | 
			
		||||
import java.util.HashMap;
 | 
			
		||||
import java.util.Map;
 | 
			
		||||
import lombok.Getter;
 | 
			
		||||
import org.openmetadata.schema.services.connections.metadata.SecretsManagerProvider;
 | 
			
		||||
import org.openmetadata.service.exception.SecretsManagerException;
 | 
			
		||||
import org.openmetadata.service.util.JsonUtils;
 | 
			
		||||
 | 
			
		||||
/** Secret Manager used for testing */
 | 
			
		||||
public class InMemorySecretsManager extends ThirdPartySecretsManager {
 | 
			
		||||
@ -56,12 +54,4 @@ public class InMemorySecretsManager extends ThirdPartySecretsManager {
 | 
			
		||||
    }
 | 
			
		||||
    return value;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public Object getBotConfig(String botName) {
 | 
			
		||||
    try {
 | 
			
		||||
      return JsonUtils.readValue(getSecret(buildSecretId(BOT_PREFIX, botName)), Object.class);
 | 
			
		||||
    } catch (IOException e) {
 | 
			
		||||
      throw new RuntimeException(e);
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -65,10 +65,15 @@ public class NoopSecretsManager extends SecretsManager {
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  public Object encryptOrDecryptIngestionBotCredentials(String botName, Object securityConfig, boolean encrypt) {
 | 
			
		||||
  public Object encryptOrDecryptBotUserCredentials(String botUserName, Object securityConfig, boolean encrypt) {
 | 
			
		||||
    return securityConfig;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  public Object encryptOrDecryptBotCredentials(String botName, String botUserName, boolean encrypt) {
 | 
			
		||||
    return null;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private void encryptOrDecryptField(Object connConfig, String field, Class<?> clazz, boolean encrypt)
 | 
			
		||||
      throws InvocationTargetException, IllegalAccessException {
 | 
			
		||||
    try {
 | 
			
		||||
 | 
			
		||||
@ -105,8 +105,9 @@ public abstract class SecretsManager {
 | 
			
		||||
 | 
			
		||||
  public abstract Object storeTestConnectionObject(TestServiceConnection testServiceConnection);
 | 
			
		||||
 | 
			
		||||
  public abstract Object encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
      String botName, Object securityConfig, boolean encrypt);
 | 
			
		||||
  public abstract Object encryptOrDecryptBotUserCredentials(String botUserName, Object securityConfig, boolean encrypt);
 | 
			
		||||
 | 
			
		||||
  public abstract Object encryptOrDecryptBotCredentials(String botName, String botUserName, boolean encrypt);
 | 
			
		||||
 | 
			
		||||
  public void validateServiceConnection(Object connectionConfig, String connectionType, ServiceType serviceType) {
 | 
			
		||||
    try {
 | 
			
		||||
 | 
			
		||||
@ -237,9 +237,9 @@ public class SecretsManagerMigrationService {
 | 
			
		||||
      User user = userRepository.dao.findEntityById(botUser.getId());
 | 
			
		||||
 | 
			
		||||
      Object authConfig =
 | 
			
		||||
          oldSecretManager.encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
          oldSecretManager.encryptOrDecryptBotUserCredentials(
 | 
			
		||||
              botUser.getName(), user.getAuthenticationMechanism().getConfig(), false);
 | 
			
		||||
      authConfig = newSecretManager.encryptOrDecryptIngestionBotCredentials(botUser.getName(), authConfig, true);
 | 
			
		||||
      authConfig = newSecretManager.encryptOrDecryptBotUserCredentials(botUser.getName(), authConfig, true);
 | 
			
		||||
 | 
			
		||||
      user.getAuthenticationMechanism().setConfig(authConfig);
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@ -13,19 +13,34 @@
 | 
			
		||||
 | 
			
		||||
package org.openmetadata.service.secrets;
 | 
			
		||||
 | 
			
		||||
import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.JWT;
 | 
			
		||||
import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.SSO;
 | 
			
		||||
 | 
			
		||||
import com.fasterxml.jackson.core.JsonProcessingException;
 | 
			
		||||
import java.util.List;
 | 
			
		||||
import org.jetbrains.annotations.Nullable;
 | 
			
		||||
import org.openmetadata.schema.api.services.ingestionPipelines.TestServiceConnection;
 | 
			
		||||
import org.openmetadata.schema.entity.services.ServiceType;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.User;
 | 
			
		||||
import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
 | 
			
		||||
import org.openmetadata.schema.services.connections.metadata.OpenMetadataServerConnection;
 | 
			
		||||
import org.openmetadata.schema.services.connections.metadata.SecretsManagerProvider;
 | 
			
		||||
import org.openmetadata.schema.teams.authn.JWTAuthMechanism;
 | 
			
		||||
import org.openmetadata.schema.teams.authn.SSOAuthMechanism;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.exception.InvalidServiceConnectionException;
 | 
			
		||||
import org.openmetadata.service.exception.SecretsManagerException;
 | 
			
		||||
import org.openmetadata.service.jdbi3.UserRepository;
 | 
			
		||||
import org.openmetadata.service.util.EntityUtil;
 | 
			
		||||
import org.openmetadata.service.util.JsonUtils;
 | 
			
		||||
 | 
			
		||||
public abstract class ThirdPartySecretsManager extends SecretsManager {
 | 
			
		||||
  public static final String DATABASE_METADATA_PIPELINE_SECRET_ID_PREFIX = "database-metadata-pipeline";
 | 
			
		||||
  public static final String TEST_CONNECTION_TEMP_SECRET_ID_PREFIX = "test-connection-temp";
 | 
			
		||||
  public static final String BOT_USER_PREFIX = "bot-user";
 | 
			
		||||
  public static final String BOT_PREFIX = "bot";
 | 
			
		||||
  public static final String AUTH_PROVIDER = "auth-provider";
 | 
			
		||||
  public static final String NULL_SECRET_STRING = "null";
 | 
			
		||||
 | 
			
		||||
  protected ThirdPartySecretsManager(SecretsManagerProvider secretsManagerProvider, String clusterPrefix) {
 | 
			
		||||
@ -70,11 +85,51 @@ public abstract class ThirdPartySecretsManager extends SecretsManager {
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  public Object encryptOrDecryptIngestionBotCredentials(String botName, Object securityConfig, boolean encrypt) {
 | 
			
		||||
    String secretName = buildSecretId(BOT_PREFIX, botName);
 | 
			
		||||
  public Object encryptOrDecryptBotUserCredentials(String botUserName, Object securityConfig, boolean encrypt) {
 | 
			
		||||
    String secretName = buildSecretId(BOT_USER_PREFIX, botUserName);
 | 
			
		||||
    return encryptOrDecryptObject(securityConfig, encrypt, secretName);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  // TODO: move this logic outside secrets manager
 | 
			
		||||
  public Object encryptOrDecryptBotCredentials(String botName, String botUserName, boolean encrypt) {
 | 
			
		||||
    String secretName = buildSecretId(BOT_PREFIX, botName);
 | 
			
		||||
    if (encrypt) {
 | 
			
		||||
      try {
 | 
			
		||||
        // save bot user auth config
 | 
			
		||||
        Object authConfig = encryptOrDecryptBotUserCredentials(botUserName, null, false);
 | 
			
		||||
        // save bot user auth provider
 | 
			
		||||
        User botUser =
 | 
			
		||||
            UserRepository.class
 | 
			
		||||
                .cast(Entity.getEntityRepository(Entity.USER))
 | 
			
		||||
                .getByName(null, botUserName, new EntityUtil.Fields(List.of("authenticationMechanism")));
 | 
			
		||||
        AuthenticationMechanism authMechanism = botUser.getAuthenticationMechanism();
 | 
			
		||||
        if (authMechanism != null) {
 | 
			
		||||
          String authProviderSecretName = buildSecretId(BOT_PREFIX, botName, AUTH_PROVIDER);
 | 
			
		||||
          String authProvider = null;
 | 
			
		||||
          if (JWT.equals(authMechanism.getAuthType())) {
 | 
			
		||||
            JWTAuthMechanism jwtAuthMechanism = JsonUtils.convertValue(authConfig, JWTAuthMechanism.class);
 | 
			
		||||
            encryptOrDecryptObject(
 | 
			
		||||
                new OpenMetadataJWTClientConfig().withJwtToken(jwtAuthMechanism.getJWTToken()), true, secretName);
 | 
			
		||||
            authProvider = OpenMetadataServerConnection.AuthProvider.OPENMETADATA.value();
 | 
			
		||||
          } else if (authConfig != null && SSO.equals(authMechanism.getAuthType())) {
 | 
			
		||||
            encryptOrDecryptObject(
 | 
			
		||||
                JsonUtils.convertValue(authConfig, SSOAuthMechanism.class).getAuthConfig(), true, secretName);
 | 
			
		||||
            authProvider =
 | 
			
		||||
                OpenMetadataServerConnection.AuthProvider.fromValue(
 | 
			
		||||
                        (String) JsonUtils.getMap(authConfig).get("ssoServiceType"))
 | 
			
		||||
                    .value();
 | 
			
		||||
          }
 | 
			
		||||
          encryptOrDecryptObject(authProvider, true, authProviderSecretName);
 | 
			
		||||
        }
 | 
			
		||||
      } catch (Exception e) {
 | 
			
		||||
        throw SecretsManagerException.byMessage(getClass().getSimpleName(), secretName, e.getMessage());
 | 
			
		||||
      }
 | 
			
		||||
    } else {
 | 
			
		||||
      return encryptOrDecryptObject(null, false, secretName);
 | 
			
		||||
    }
 | 
			
		||||
    return null;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  public Object encryptOrDecryptDbtConfigSource(Object dbtConfigSource, String serviceName, boolean encrypt) {
 | 
			
		||||
    String secretName = buildSecretId(DATABASE_METADATA_PIPELINE_SECRET_ID_PREFIX, serviceName);
 | 
			
		||||
@ -82,10 +137,10 @@ public abstract class ThirdPartySecretsManager extends SecretsManager {
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Nullable
 | 
			
		||||
  private Object encryptOrDecryptObject(Object securityConfig, boolean encrypt, String secretName) {
 | 
			
		||||
  private Object encryptOrDecryptObject(Object objectValue, boolean encrypt, String secretName) {
 | 
			
		||||
    try {
 | 
			
		||||
      if (encrypt) {
 | 
			
		||||
        String securityConfigJson = JsonUtils.pojoToJson(securityConfig);
 | 
			
		||||
        String securityConfigJson = JsonUtils.pojoToJson(objectValue);
 | 
			
		||||
        upsertSecret(secretName, securityConfigJson);
 | 
			
		||||
        return null;
 | 
			
		||||
      } else {
 | 
			
		||||
 | 
			
		||||
@ -27,12 +27,14 @@ import static org.openmetadata.service.security.SecurityUtil.DEFAULT_PRINCIPAL_D
 | 
			
		||||
 | 
			
		||||
import at.favre.lib.crypto.bcrypt.BCrypt;
 | 
			
		||||
import java.io.IOException;
 | 
			
		||||
import java.util.Arrays;
 | 
			
		||||
import java.util.HashMap;
 | 
			
		||||
import java.util.HashSet;
 | 
			
		||||
import java.util.List;
 | 
			
		||||
import java.util.Map;
 | 
			
		||||
import java.util.Set;
 | 
			
		||||
import java.util.UUID;
 | 
			
		||||
import java.util.stream.Collectors;
 | 
			
		||||
import javax.ws.rs.core.SecurityContext;
 | 
			
		||||
import lombok.extern.slf4j.Slf4j;
 | 
			
		||||
import org.apache.commons.lang3.exception.ExceptionUtils;
 | 
			
		||||
@ -40,6 +42,7 @@ import org.jdbi.v3.core.Jdbi;
 | 
			
		||||
import org.openmetadata.schema.api.configuration.airflow.AirflowConfiguration;
 | 
			
		||||
import org.openmetadata.schema.api.security.AuthenticationConfiguration;
 | 
			
		||||
import org.openmetadata.schema.entity.Bot;
 | 
			
		||||
import org.openmetadata.schema.entity.BotType;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.User;
 | 
			
		||||
import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
 | 
			
		||||
@ -53,7 +56,9 @@ import org.openmetadata.schema.type.ResourcePermission;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.OpenMetadataApplicationConfig;
 | 
			
		||||
import org.openmetadata.service.exception.EntityNotFoundException;
 | 
			
		||||
import org.openmetadata.service.jdbi3.BotRepository;
 | 
			
		||||
import org.openmetadata.service.jdbi3.EntityRepository;
 | 
			
		||||
import org.openmetadata.service.jdbi3.UserRepository;
 | 
			
		||||
import org.openmetadata.service.secrets.SecretsManager;
 | 
			
		||||
import org.openmetadata.service.secrets.SecretsManagerFactory;
 | 
			
		||||
import org.openmetadata.service.security.jwt.JWTTokenGenerator;
 | 
			
		||||
@ -75,7 +80,7 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
  private final String COLONDELIMETER = ":";
 | 
			
		||||
  private final String DEFAULT_ADMIN = ADMIN_USER_NAME;
 | 
			
		||||
  private Set<String> adminUsers;
 | 
			
		||||
  private Set<String> botUsers;
 | 
			
		||||
  private Set<String> botPrincipalUsers;
 | 
			
		||||
  private Set<String> testUsers;
 | 
			
		||||
  private String principalDomain;
 | 
			
		||||
 | 
			
		||||
@ -86,7 +91,8 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
    LOG.info(
 | 
			
		||||
        "Initializing DefaultAuthorizer with config {}", openMetadataApplicationConfig.getAuthorizerConfiguration());
 | 
			
		||||
    this.adminUsers = new HashSet<>(openMetadataApplicationConfig.getAuthorizerConfiguration().getAdminPrincipals());
 | 
			
		||||
    this.botUsers = new HashSet<>(openMetadataApplicationConfig.getAuthorizerConfiguration().getBotPrincipals());
 | 
			
		||||
    this.botPrincipalUsers =
 | 
			
		||||
        new HashSet<>(openMetadataApplicationConfig.getAuthorizerConfiguration().getBotPrincipals());
 | 
			
		||||
    this.testUsers = new HashSet<>(openMetadataApplicationConfig.getAuthorizerConfiguration().getTestPrincipals());
 | 
			
		||||
    this.principalDomain = openMetadataApplicationConfig.getAuthorizerConfiguration().getPrincipalDomain();
 | 
			
		||||
    this.secretsManager =
 | 
			
		||||
@ -121,11 +127,20 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
    }
 | 
			
		||||
 | 
			
		||||
    LOG.debug("Checking user entries for bot users");
 | 
			
		||||
    Set<String> botUsers = Arrays.stream(BotType.values()).map(BotType::value).collect(Collectors.toSet());
 | 
			
		||||
    botUsers.remove(BotType.BOT.value());
 | 
			
		||||
    botUsers.addAll(botPrincipalUsers);
 | 
			
		||||
    for (String botUser : botUsers) {
 | 
			
		||||
      User user = user(botUser, domain, botUser).withIsBot(true);
 | 
			
		||||
      User user = user(botUser, domain, botUser).withIsBot(true).withIsAdmin(false);
 | 
			
		||||
      user = addOrUpdateBotUser(user, openMetadataApplicationConfig);
 | 
			
		||||
      if (user != null) {
 | 
			
		||||
        Bot bot = bot(user).withBotUser(user.getEntityReference());
 | 
			
		||||
        BotType botType;
 | 
			
		||||
        try {
 | 
			
		||||
          botType = BotType.fromValue(botUser);
 | 
			
		||||
        } catch (IllegalArgumentException e) {
 | 
			
		||||
          botType = BotType.BOT;
 | 
			
		||||
        }
 | 
			
		||||
        Bot bot = bot(user).withBotUser(user.getEntityReference()).withBotType(botType);
 | 
			
		||||
        addOrUpdateBot(bot);
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
@ -326,13 +341,14 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
   *       </ul>
 | 
			
		||||
   * </ul>
 | 
			
		||||
   *
 | 
			
		||||
   * @param user
 | 
			
		||||
   * @param openMetadataApplicationConfig
 | 
			
		||||
   * @return
 | 
			
		||||
   * @param user the user
 | 
			
		||||
   * @param openMetadataApplicationConfig the OM config
 | 
			
		||||
   * @return enriched user
 | 
			
		||||
   */
 | 
			
		||||
  private User addOrUpdateBotUser(User user, OpenMetadataApplicationConfig openMetadataApplicationConfig) {
 | 
			
		||||
    AuthenticationMechanism authMechanism = retrieveAuthMechanism(user);
 | 
			
		||||
    User originalUser = retrieveAuthMechanism(user);
 | 
			
		||||
    // the user did not have an auth mechanism
 | 
			
		||||
    AuthenticationMechanism authMechanism = originalUser != null ? originalUser.getAuthenticationMechanism() : null;
 | 
			
		||||
    if (authMechanism == null) {
 | 
			
		||||
      AuthenticationConfiguration authConfig = openMetadataApplicationConfig.getAuthenticationConfiguration();
 | 
			
		||||
      AirflowConfiguration airflowConfig = openMetadataApplicationConfig.getAirflowConfiguration();
 | 
			
		||||
@ -385,6 +401,8 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
    user.setAuthenticationMechanism(authMechanism);
 | 
			
		||||
    user.setDescription(user.getDescription());
 | 
			
		||||
    user.setDisplayName(user.getDisplayName());
 | 
			
		||||
    return addOrUpdateUser(user);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -396,19 +414,18 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
    return new AuthenticationMechanism().withAuthType(authType).withConfig(config);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private AuthenticationMechanism retrieveAuthMechanism(User user) {
 | 
			
		||||
    EntityRepository<User> userRepository = Entity.getEntityRepository(Entity.USER);
 | 
			
		||||
  private User retrieveAuthMechanism(User user) {
 | 
			
		||||
    EntityRepository<User> userRepository = UserRepository.class.cast(Entity.getEntityRepository(Entity.USER));
 | 
			
		||||
    try {
 | 
			
		||||
      User originalUser =
 | 
			
		||||
          userRepository.getByName(
 | 
			
		||||
              null, user.getFullyQualifiedName(), new EntityUtil.Fields(List.of("authenticationMechanism")));
 | 
			
		||||
      AuthenticationMechanism authMechanism = user.getAuthenticationMechanism();
 | 
			
		||||
          userRepository.getByName(null, user.getName(), new EntityUtil.Fields(List.of("authenticationMechanism")));
 | 
			
		||||
      AuthenticationMechanism authMechanism = originalUser.getAuthenticationMechanism();
 | 
			
		||||
      if (authMechanism != null) {
 | 
			
		||||
        Object config =
 | 
			
		||||
            secretsManager.encryptOrDecryptIngestionBotCredentials(user.getName(), authMechanism.getConfig(), false);
 | 
			
		||||
            secretsManager.encryptOrDecryptBotUserCredentials(user.getName(), authMechanism.getConfig(), false);
 | 
			
		||||
        authMechanism.setConfig(config != null ? config : authMechanism.getConfig());
 | 
			
		||||
      }
 | 
			
		||||
      return originalUser.getAuthenticationMechanism();
 | 
			
		||||
      return originalUser;
 | 
			
		||||
    } catch (IOException | EntityNotFoundException e) {
 | 
			
		||||
      LOG.debug("Bot entity: {} does not exists.", user);
 | 
			
		||||
      return null;
 | 
			
		||||
@ -416,7 +433,13 @@ public class DefaultAuthorizer implements Authorizer {
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private void addOrUpdateBot(Bot bot) {
 | 
			
		||||
    EntityRepository<Bot> botRepository = Entity.getEntityRepository(Entity.BOT);
 | 
			
		||||
    EntityRepository<Bot> botRepository = BotRepository.class.cast(Entity.getEntityRepository(Entity.BOT));
 | 
			
		||||
    Bot originalBot;
 | 
			
		||||
    try {
 | 
			
		||||
      originalBot = botRepository.getByName(null, bot.getName(), EntityUtil.Fields.EMPTY_FIELDS);
 | 
			
		||||
      bot.setBotUser(originalBot.getBotUser());
 | 
			
		||||
    } catch (Exception e) {
 | 
			
		||||
    }
 | 
			
		||||
    try {
 | 
			
		||||
      RestUtil.PutResponse<Bot> addedBot = botRepository.createOrUpdate(null, bot);
 | 
			
		||||
      LOG.debug("Added bot entry: {}", addedBot.getEntity().getName());
 | 
			
		||||
 | 
			
		||||
@ -1,14 +1,12 @@
 | 
			
		||||
package org.openmetadata.service.util;
 | 
			
		||||
 | 
			
		||||
import java.io.IOException;
 | 
			
		||||
import java.util.ArrayList;
 | 
			
		||||
import java.util.Iterator;
 | 
			
		||||
import java.util.List;
 | 
			
		||||
import java.util.Set;
 | 
			
		||||
import lombok.extern.slf4j.Slf4j;
 | 
			
		||||
import org.openmetadata.api.configuration.airflow.SSLConfig;
 | 
			
		||||
import org.openmetadata.schema.api.configuration.airflow.AirflowConfiguration;
 | 
			
		||||
import org.openmetadata.schema.entity.Bot;
 | 
			
		||||
import org.openmetadata.schema.entity.BotType;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.User;
 | 
			
		||||
import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
 | 
			
		||||
@ -16,18 +14,16 @@ import org.openmetadata.schema.services.connections.metadata.OpenMetadataServerC
 | 
			
		||||
import org.openmetadata.schema.services.connections.metadata.SecretsManagerProvider;
 | 
			
		||||
import org.openmetadata.schema.teams.authn.JWTAuthMechanism;
 | 
			
		||||
import org.openmetadata.schema.teams.authn.SSOAuthMechanism;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.OpenMetadataApplicationConfig;
 | 
			
		||||
import org.openmetadata.service.exception.EntityNotFoundException;
 | 
			
		||||
import org.openmetadata.service.jdbi3.BotRepository;
 | 
			
		||||
import org.openmetadata.service.jdbi3.CollectionDAO;
 | 
			
		||||
import org.openmetadata.service.jdbi3.UserRepository;
 | 
			
		||||
import org.openmetadata.service.secrets.SecretsManager;
 | 
			
		||||
 | 
			
		||||
@Slf4j
 | 
			
		||||
public class OpenMetadataServerConnectionBuilder {
 | 
			
		||||
 | 
			
		||||
  private static final String INGESTION_BOT = "ingestion-bot";
 | 
			
		||||
 | 
			
		||||
  OpenMetadataServerConnection.AuthProvider authProvider;
 | 
			
		||||
  String bot;
 | 
			
		||||
  Object securityConfig;
 | 
			
		||||
@ -41,9 +37,7 @@ public class OpenMetadataServerConnectionBuilder {
 | 
			
		||||
  SecretsManager secretsManager;
 | 
			
		||||
 | 
			
		||||
  public OpenMetadataServerConnectionBuilder(
 | 
			
		||||
      SecretsManager secretsManager,
 | 
			
		||||
      OpenMetadataApplicationConfig openMetadataApplicationConfig,
 | 
			
		||||
      CollectionDAO collectionDAO) {
 | 
			
		||||
      SecretsManager secretsManager, OpenMetadataApplicationConfig openMetadataApplicationConfig) {
 | 
			
		||||
    this.secretsManager = secretsManager;
 | 
			
		||||
    // TODO: https://github.com/open-metadata/OpenMetadata/issues/7712
 | 
			
		||||
    authProvider =
 | 
			
		||||
@ -53,9 +47,9 @@ public class OpenMetadataServerConnectionBuilder {
 | 
			
		||||
                openMetadataApplicationConfig.getAuthenticationConfiguration().getProvider());
 | 
			
		||||
 | 
			
		||||
    if (!OpenMetadataServerConnection.AuthProvider.NO_AUTH.equals(authProvider)) {
 | 
			
		||||
      botRepository = new BotRepository(collectionDAO);
 | 
			
		||||
      userRepository = new UserRepository(collectionDAO, secretsManager);
 | 
			
		||||
      User botUser = retrieveBotUser(openMetadataApplicationConfig);
 | 
			
		||||
      botRepository = BotRepository.class.cast(Entity.getEntityRepository(Entity.BOT));
 | 
			
		||||
      userRepository = UserRepository.class.cast(Entity.getEntityRepository(Entity.USER));
 | 
			
		||||
      User botUser = retrieveBotUser();
 | 
			
		||||
      if (secretsManager.isLocal()) {
 | 
			
		||||
        securityConfig = extractSecurityConfig(botUser);
 | 
			
		||||
      }
 | 
			
		||||
@ -111,44 +105,20 @@ public class OpenMetadataServerConnectionBuilder {
 | 
			
		||||
        .withVerifySSL(verifySSL)
 | 
			
		||||
        .withClusterName(clusterName)
 | 
			
		||||
        .withSecretsManagerProvider(secretsManagerProvider)
 | 
			
		||||
        .withWorkflowBot(bot)
 | 
			
		||||
        .withSslConfig(airflowSSLConfig);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private User retrieveBotUser(OpenMetadataApplicationConfig openMetadataApplicationConfig) {
 | 
			
		||||
    Set<String> botPrincipals = openMetadataApplicationConfig.getAuthorizerConfiguration().getBotPrincipals();
 | 
			
		||||
    if (botPrincipals == null || botPrincipals.isEmpty()) {
 | 
			
		||||
      throw new IllegalArgumentException(
 | 
			
		||||
          "Please, add at least one bot to the 'authorizerConfiguration.botPrincipals' in the OpenMetadata configuration");
 | 
			
		||||
    }
 | 
			
		||||
    User botUser = null;
 | 
			
		||||
    // try to find ingestion-bot user
 | 
			
		||||
    if (botPrincipals.contains(INGESTION_BOT)) {
 | 
			
		||||
      botUser = retrieveIngestionBotUser(INGESTION_BOT);
 | 
			
		||||
      addBotNameIfUserExists(botUser, INGESTION_BOT);
 | 
			
		||||
    }
 | 
			
		||||
    // sort botPrincipals in order to use always the same alternate bot
 | 
			
		||||
    List<String> sortedBotPrincipals = new ArrayList<>(botPrincipals);
 | 
			
		||||
    sortedBotPrincipals.sort(String::compareTo);
 | 
			
		||||
    Iterator<String> it = sortedBotPrincipals.iterator();
 | 
			
		||||
    while (botUser == null && it.hasNext()) {
 | 
			
		||||
      String botName = it.next();
 | 
			
		||||
      botUser = retrieveIngestionBotUser(botName);
 | 
			
		||||
      if (botUser != null && botUser.getAuthenticationMechanism() == null) {
 | 
			
		||||
        botUser = null;
 | 
			
		||||
      }
 | 
			
		||||
      addBotNameIfUserExists(botUser, botName);
 | 
			
		||||
    }
 | 
			
		||||
  private User retrieveBotUser() {
 | 
			
		||||
    User botUser = retrieveIngestionBotUser(BotType.INGESTION_BOT.value());
 | 
			
		||||
    if (botUser == null) {
 | 
			
		||||
      throw new IllegalArgumentException(
 | 
			
		||||
          "Please, create at least one bot with valid authentication mechanism that matches any of the names from 'authorizerConfiguration.botPrincipals' in the OpenMetadata configuration");
 | 
			
		||||
      throw new IllegalArgumentException("Please, verify that the ingestion-bot is present.");
 | 
			
		||||
    }
 | 
			
		||||
    return botUser;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private User retrieveIngestionBotUser(String botName) {
 | 
			
		||||
    try {
 | 
			
		||||
      Bot bot = botRepository.getByName(null, botName, new EntityUtil.Fields(List.of("botUser")));
 | 
			
		||||
      Bot bot = botRepository.getByName(null, botName, EntityUtil.Fields.EMPTY_FIELDS);
 | 
			
		||||
      if (bot.getBotUser() == null) {
 | 
			
		||||
        return null;
 | 
			
		||||
      }
 | 
			
		||||
@ -160,8 +130,8 @@ public class OpenMetadataServerConnectionBuilder {
 | 
			
		||||
      if (user.getAuthenticationMechanism() != null) {
 | 
			
		||||
        user.getAuthenticationMechanism()
 | 
			
		||||
            .setConfig(
 | 
			
		||||
                secretsManager.encryptOrDecryptIngestionBotCredentials(
 | 
			
		||||
                    botName, user.getAuthenticationMechanism().getConfig(), false));
 | 
			
		||||
                secretsManager.encryptOrDecryptBotUserCredentials(
 | 
			
		||||
                    user.getName(), user.getAuthenticationMechanism().getConfig(), false));
 | 
			
		||||
      }
 | 
			
		||||
      return user;
 | 
			
		||||
    } catch (IOException | EntityNotFoundException ex) {
 | 
			
		||||
 | 
			
		||||
@ -1,17 +0,0 @@
 | 
			
		||||
package org.openmetadata.service.validators;
 | 
			
		||||
 | 
			
		||||
import java.lang.annotation.ElementType;
 | 
			
		||||
import java.lang.annotation.Retention;
 | 
			
		||||
import java.lang.annotation.RetentionPolicy;
 | 
			
		||||
import java.lang.annotation.Target;
 | 
			
		||||
import javax.validation.Payload;
 | 
			
		||||
 | 
			
		||||
@Target({ElementType.FIELD})
 | 
			
		||||
@Retention(RetentionPolicy.RUNTIME)
 | 
			
		||||
public @interface AirflowConfigValidation {
 | 
			
		||||
  String message() default "This will be replaced by the validation";
 | 
			
		||||
 | 
			
		||||
  Class<?>[] groups() default {};
 | 
			
		||||
 | 
			
		||||
  Class<? extends Payload>[] payload() default {};
 | 
			
		||||
}
 | 
			
		||||
@ -1,88 +0,0 @@
 | 
			
		||||
package org.openmetadata.service.fixtures;
 | 
			
		||||
 | 
			
		||||
import io.dropwizard.db.DataSourceFactory;
 | 
			
		||||
import java.util.List;
 | 
			
		||||
import org.openmetadata.api.configuration.airflow.SSLConfig;
 | 
			
		||||
import org.openmetadata.schema.api.configuration.airflow.AirflowConfiguration;
 | 
			
		||||
import org.openmetadata.schema.security.client.Auth0SSOClientConfig;
 | 
			
		||||
import org.openmetadata.schema.security.client.AzureSSOClientConfig;
 | 
			
		||||
import org.openmetadata.schema.security.client.CustomOIDCSSOClientConfig;
 | 
			
		||||
import org.openmetadata.schema.security.client.OktaSSOClientConfig;
 | 
			
		||||
import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
 | 
			
		||||
import org.openmetadata.schema.security.ssl.ValidateSSLClientConfig;
 | 
			
		||||
import org.openmetadata.schema.services.connections.metadata.OpenMetadataServerConnection;
 | 
			
		||||
import org.openmetadata.service.OpenMetadataApplicationConfig;
 | 
			
		||||
import org.openmetadata.service.migration.MigrationConfiguration;
 | 
			
		||||
 | 
			
		||||
public class ConfigurationFixtures {
 | 
			
		||||
 | 
			
		||||
  public static OpenMetadataApplicationConfig buildOpenMetadataApplicationConfig(
 | 
			
		||||
      OpenMetadataServerConnection.AuthProvider authProvider) {
 | 
			
		||||
    OpenMetadataApplicationConfig openMetadataApplicationConfig = new OpenMetadataApplicationConfig();
 | 
			
		||||
    DataSourceFactory dataSourceFactory = new DataSourceFactory();
 | 
			
		||||
    dataSourceFactory.setDriverClass("driverClass");
 | 
			
		||||
    dataSourceFactory.setUrl("http://localhost");
 | 
			
		||||
    MigrationConfiguration migrationConfiguration = new MigrationConfiguration();
 | 
			
		||||
    migrationConfiguration.setPath("/fake/path");
 | 
			
		||||
    openMetadataApplicationConfig.setDataSourceFactory(dataSourceFactory);
 | 
			
		||||
    openMetadataApplicationConfig.setMigrationConfiguration(migrationConfiguration);
 | 
			
		||||
    openMetadataApplicationConfig.setAirflowConfiguration(buildAirflowConfig(authProvider));
 | 
			
		||||
    return openMetadataApplicationConfig;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static AirflowConfiguration buildAirflowConfig(OpenMetadataServerConnection.AuthProvider authProvider) {
 | 
			
		||||
    AirflowConfiguration airflowConfiguration = new AirflowConfiguration();
 | 
			
		||||
    airflowConfiguration.setUsername("admin");
 | 
			
		||||
    airflowConfiguration.setPassword("admin");
 | 
			
		||||
    airflowConfiguration.setApiEndpoint("http://localhost:8080/api");
 | 
			
		||||
    airflowConfiguration.setMetadataApiEndpoint("http://localhost:8585/api");
 | 
			
		||||
    return airflowConfiguration;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static AirflowConfiguration buildAirflowSSLConfig(OpenMetadataServerConnection.AuthProvider authProvider) {
 | 
			
		||||
    AirflowConfiguration airflowConfiguration = new AirflowConfiguration();
 | 
			
		||||
    airflowConfiguration.setUsername("admin");
 | 
			
		||||
    airflowConfiguration.setPassword("admin");
 | 
			
		||||
    airflowConfiguration.setApiEndpoint("http://localhost:8080/api");
 | 
			
		||||
    airflowConfiguration.setMetadataApiEndpoint("http://localhost:8585/api");
 | 
			
		||||
    airflowConfiguration.setVerifySSL(String.valueOf(OpenMetadataServerConnection.VerifySSL.VALIDATE));
 | 
			
		||||
 | 
			
		||||
    ValidateSSLClientConfig validateSSLClientConfig = new ValidateSSLClientConfig().withCertificatePath("/public.cert");
 | 
			
		||||
    SSLConfig sslConfig = new SSLConfig().withValidate(validateSSLClientConfig);
 | 
			
		||||
 | 
			
		||||
    airflowConfiguration.setSslConfig(sslConfig);
 | 
			
		||||
    return airflowConfiguration;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static OktaSSOClientConfig buildOktaSSOClientConfig() {
 | 
			
		||||
    return new OktaSSOClientConfig()
 | 
			
		||||
        .withClientId("1234")
 | 
			
		||||
        .withEmail("test@test.com")
 | 
			
		||||
        .withOrgURL("https://okta.domain.com")
 | 
			
		||||
        .withPrivateKey("34123")
 | 
			
		||||
        .withScopes(List.of("local", "prod", "test"));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static Auth0SSOClientConfig buildAuth0SSOClientConfig() {
 | 
			
		||||
    return new Auth0SSOClientConfig().withClientId("1234").withDomain("local").withSecretKey("34123");
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static AzureSSOClientConfig buildAzureClientConfig() {
 | 
			
		||||
    return new AzureSSOClientConfig()
 | 
			
		||||
        .withClientId("1234")
 | 
			
		||||
        .withClientSecret("34123")
 | 
			
		||||
        .withAuthority("local")
 | 
			
		||||
        .withScopes(List.of("local", "prod", "test"));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static OpenMetadataJWTClientConfig buildOpenMetadataJWTClientConfig() {
 | 
			
		||||
    return new OpenMetadataJWTClientConfig().withJwtToken("fakeToken");
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public static CustomOIDCSSOClientConfig buildCustomOIDCSSOClientConfig() {
 | 
			
		||||
    return new CustomOIDCSSOClientConfig()
 | 
			
		||||
        .withClientId("1234")
 | 
			
		||||
        .withSecretKey("34123")
 | 
			
		||||
        .withTokenEndpoint("https://localhost/");
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
@ -435,7 +435,7 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
 | 
			
		||||
 | 
			
		||||
    T entity = createEntity(createRequest(test, 0), ADMIN_AUTH_HEADERS);
 | 
			
		||||
 | 
			
		||||
    String allFields = String.join(",", Entity.getAllowedFields(entityClass));
 | 
			
		||||
    String allFields = getAllowedFields();
 | 
			
		||||
 | 
			
		||||
    // GET an entity by ID with all the field names of an entity should be successful
 | 
			
		||||
    getEntity(entity.getId(), allFields, ADMIN_AUTH_HEADERS);
 | 
			
		||||
@ -2218,4 +2218,8 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
 | 
			
		||||
    // In requests send minimum entity reference information to ensure the server fills rest of the details
 | 
			
		||||
    return ref != null ? new EntityReference().withType(ref.getType()).withId(ref.getId()) : null;
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  protected String getAllowedFields() {
 | 
			
		||||
    return String.join(",", Entity.getAllowedFields(entityClass));
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -1,26 +1,31 @@
 | 
			
		||||
package org.openmetadata.service.resources.bots;
 | 
			
		||||
 | 
			
		||||
import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
 | 
			
		||||
import static org.openmetadata.service.util.TestUtils.ADMIN_AUTH_HEADERS;
 | 
			
		||||
import static org.openmetadata.service.util.TestUtils.assertResponse;
 | 
			
		||||
 | 
			
		||||
import java.io.IOException;
 | 
			
		||||
import java.net.URISyntaxException;
 | 
			
		||||
import java.util.Map;
 | 
			
		||||
import java.util.Objects;
 | 
			
		||||
import lombok.SneakyThrows;
 | 
			
		||||
import org.apache.http.client.HttpResponseException;
 | 
			
		||||
import org.junit.jupiter.api.BeforeAll;
 | 
			
		||||
import org.junit.jupiter.api.BeforeEach;
 | 
			
		||||
import org.junit.jupiter.api.Test;
 | 
			
		||||
import org.junit.jupiter.api.TestInfo;
 | 
			
		||||
import org.openmetadata.schema.api.CreateBot;
 | 
			
		||||
import org.openmetadata.schema.api.teams.CreateUser;
 | 
			
		||||
import org.openmetadata.schema.entity.Bot;
 | 
			
		||||
import org.openmetadata.schema.entity.BotType;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.User;
 | 
			
		||||
import org.openmetadata.schema.type.EntityReference;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.resources.EntityResourceTest;
 | 
			
		||||
import org.openmetadata.service.resources.bots.BotResource.BotList;
 | 
			
		||||
import org.openmetadata.service.resources.teams.UserResourceTest;
 | 
			
		||||
import org.openmetadata.service.util.ResultList;
 | 
			
		||||
 | 
			
		||||
class BotResourceTest extends EntityResourceTest<Bot, CreateBot> {
 | 
			
		||||
public class BotResourceTest extends EntityResourceTest<Bot, CreateBot> {
 | 
			
		||||
  public static User botUser;
 | 
			
		||||
  public static EntityReference botUserRef;
 | 
			
		||||
 | 
			
		||||
@ -32,10 +37,19 @@ class BotResourceTest extends EntityResourceTest<Bot, CreateBot> {
 | 
			
		||||
  @BeforeAll
 | 
			
		||||
  public void setup(TestInfo test) throws URISyntaxException, IOException {
 | 
			
		||||
    super.setup(test);
 | 
			
		||||
    UserResourceTest userResourceTest = new UserResourceTest();
 | 
			
		||||
    CreateUser createUser = userResourceTest.createRequest("botUser", "", "", null);
 | 
			
		||||
    botUser = new UserResourceTest().createEntity(createUser, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    botUserRef = botUser.getEntityReference();
 | 
			
		||||
    createUser();
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @BeforeEach
 | 
			
		||||
  public void beforeEach() throws HttpResponseException {
 | 
			
		||||
    ResultList<Bot> bots = listEntities(null, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    for (Bot bot : bots.getData()) {
 | 
			
		||||
      try {
 | 
			
		||||
        deleteEntity(bot.getId(), true, true, ADMIN_AUTH_HEADERS);
 | 
			
		||||
        createUser();
 | 
			
		||||
      } catch (Exception ignored) {
 | 
			
		||||
      }
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
@ -47,9 +61,7 @@ class BotResourceTest extends EntityResourceTest<Bot, CreateBot> {
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
  void delete_ensureBotUserDelete(TestInfo test) throws IOException {
 | 
			
		||||
    UserResourceTest userResourceTest = new UserResourceTest();
 | 
			
		||||
    CreateUser createUser = userResourceTest.createRequest(test);
 | 
			
		||||
    User testUser = new UserResourceTest().createEntity(createUser, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    User testUser = new UserResourceTest().createUser("test-deleter", true);
 | 
			
		||||
    EntityReference testUserRef = testUser.getEntityReference();
 | 
			
		||||
 | 
			
		||||
    CreateBot create = createRequest(test).withBotUser(testUserRef);
 | 
			
		||||
@ -61,8 +73,53 @@ class BotResourceTest extends EntityResourceTest<Bot, CreateBot> {
 | 
			
		||||
    assertEntityDeleted(testUser.getId(), true);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
  void put_failIfUserIsAlreadyUsedByAnotherBot(TestInfo test) throws IOException {
 | 
			
		||||
    // create a bot user
 | 
			
		||||
    User testUser = new UserResourceTest().createUser("bot-test-user", true);
 | 
			
		||||
    EntityReference botUserRef = Objects.requireNonNull(testUser).getEntityReference();
 | 
			
		||||
    // create a bot
 | 
			
		||||
    CreateBot create = createRequest(test).withBotUser(botUserRef);
 | 
			
		||||
    createEntity(create, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    // create another bot with the same bot user
 | 
			
		||||
    CreateBot failCreateRequest = createRequest(test).withName("wrong-bot").withBotUser(botUserRef);
 | 
			
		||||
    assertResponse(
 | 
			
		||||
        () -> createEntity(failCreateRequest, ADMIN_AUTH_HEADERS),
 | 
			
		||||
        BAD_REQUEST,
 | 
			
		||||
        "Bot user [bot-test-user] is already used by [bot_put_failIfUserIsAlreadyUsedByAnotherBot] bot");
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
  void put_failIfUserIsNotBot(TestInfo test) throws IOException {
 | 
			
		||||
    // create a non bot user
 | 
			
		||||
    User testUser = new UserResourceTest().createUser("bot-test-user", false);
 | 
			
		||||
    EntityReference userRef = Objects.requireNonNull(testUser).getEntityReference();
 | 
			
		||||
    CreateBot failCreateRequest = createRequest(test).withBotUser(userRef);
 | 
			
		||||
    // fail because it is not a bot
 | 
			
		||||
    assertResponse(
 | 
			
		||||
        () -> createEntity(failCreateRequest, ADMIN_AUTH_HEADERS),
 | 
			
		||||
        BAD_REQUEST,
 | 
			
		||||
        "User [bot-test-user] is not a bot user");
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
  void delete_failIfUserIsIngestionBot(TestInfo test) throws IOException {
 | 
			
		||||
    // get ingestion bot
 | 
			
		||||
    Bot ingestionBot = getEntityByName(BotType.INGESTION_BOT.value(), "", ADMIN_AUTH_HEADERS);
 | 
			
		||||
    // fail because it is a system bot
 | 
			
		||||
    assertResponse(
 | 
			
		||||
        () -> deleteEntity(ingestionBot.getId(), true, true, ADMIN_AUTH_HEADERS),
 | 
			
		||||
        BAD_REQUEST,
 | 
			
		||||
        "[ingestion-bot] can not be deleted.");
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  public CreateBot createRequest(String name) {
 | 
			
		||||
    if (name != null && name.contains("entityListWithPagination_200")) {
 | 
			
		||||
      return new CreateBot()
 | 
			
		||||
          .withName(name)
 | 
			
		||||
          .withBotUser(Objects.requireNonNull(new UserResourceTest().createUser(name, true)).getEntityReference());
 | 
			
		||||
    }
 | 
			
		||||
    return new CreateBot().withName(name).withBotUser(botUserRef);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
@ -85,4 +142,11 @@ class BotResourceTest extends EntityResourceTest<Bot, CreateBot> {
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  public void assertFieldChange(String fieldName, Object expected, Object actual) throws IOException {}
 | 
			
		||||
 | 
			
		||||
  private void createUser() {
 | 
			
		||||
    botUser = new UserResourceTest().createUser("botUser", true);
 | 
			
		||||
    if (botUser != null) {
 | 
			
		||||
      botUserRef = botUser.getEntityReference();
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -31,6 +31,7 @@ import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
 | 
			
		||||
import static org.openmetadata.service.exception.CatalogExceptionMessage.entityNotFound;
 | 
			
		||||
import static org.openmetadata.service.exception.CatalogExceptionMessage.notAdmin;
 | 
			
		||||
import static org.openmetadata.service.exception.CatalogExceptionMessage.permissionNotAllowed;
 | 
			
		||||
import static org.openmetadata.service.resources.teams.UserResource.USER_PROTECTED_FIELDS;
 | 
			
		||||
import static org.openmetadata.service.security.SecurityUtil.authHeaders;
 | 
			
		||||
import static org.openmetadata.service.util.EntityUtil.fieldAdded;
 | 
			
		||||
import static org.openmetadata.service.util.EntityUtil.fieldDeleted;
 | 
			
		||||
@ -61,6 +62,7 @@ import java.util.Date;
 | 
			
		||||
import java.util.HashMap;
 | 
			
		||||
import java.util.List;
 | 
			
		||||
import java.util.Map;
 | 
			
		||||
import java.util.Objects;
 | 
			
		||||
import java.util.TimeZone;
 | 
			
		||||
import java.util.UUID;
 | 
			
		||||
import java.util.function.Predicate;
 | 
			
		||||
@ -72,6 +74,7 @@ import org.junit.jupiter.api.MethodOrderer;
 | 
			
		||||
import org.junit.jupiter.api.Test;
 | 
			
		||||
import org.junit.jupiter.api.TestInfo;
 | 
			
		||||
import org.junit.jupiter.api.TestMethodOrder;
 | 
			
		||||
import org.openmetadata.schema.api.CreateBot;
 | 
			
		||||
import org.openmetadata.schema.api.teams.CreateUser;
 | 
			
		||||
import org.openmetadata.schema.entity.data.Table;
 | 
			
		||||
import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
 | 
			
		||||
@ -90,6 +93,7 @@ import org.openmetadata.schema.type.MetadataOperation;
 | 
			
		||||
import org.openmetadata.schema.type.Profile;
 | 
			
		||||
import org.openmetadata.service.Entity;
 | 
			
		||||
import org.openmetadata.service.resources.EntityResourceTest;
 | 
			
		||||
import org.openmetadata.service.resources.bots.BotResourceTest;
 | 
			
		||||
import org.openmetadata.service.resources.databases.TableResourceTest;
 | 
			
		||||
import org.openmetadata.service.resources.locations.LocationResourceTest;
 | 
			
		||||
import org.openmetadata.service.resources.teams.UserResource.UserList;
 | 
			
		||||
@ -721,6 +725,41 @@ public class UserResourceTest extends EntityResourceTest<User, CreateUser> {
 | 
			
		||||
    assertEntityReferences(List.of(DATA_CONSUMER_ROLE_REF, DATA_STEWARD_ROLE_REF), user_team21.getInheritedRoles());
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
  void put_failIfBotUserIsAlreadyAssignedToAnotherBot(TestInfo test) throws HttpResponseException {
 | 
			
		||||
    BotResourceTest botResourceTest = new BotResourceTest();
 | 
			
		||||
    String botName = "test-bot-user-fail";
 | 
			
		||||
    // create bot user
 | 
			
		||||
    CreateUser createBotUser = creatBotUserRequest("test-bot-user", true).withBotName(botName);
 | 
			
		||||
    User botUser = updateEntity(createBotUser, CREATED, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    EntityReference botUserRef = Objects.requireNonNull(botUser).getEntityReference();
 | 
			
		||||
    // assign bot user to a bot
 | 
			
		||||
    CreateBot create = botResourceTest.createRequest(test).withBotUser(botUserRef).withName(botName);
 | 
			
		||||
    botResourceTest.createEntity(create, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    // put user with a different bot name
 | 
			
		||||
    CreateUser createWrongBotUser = creatBotUserRequest("test-bot-user", true).withBotName("test-bot-user-fail-2");
 | 
			
		||||
    assertResponse(
 | 
			
		||||
        () -> updateEntity(createWrongBotUser, BAD_REQUEST, ADMIN_AUTH_HEADERS),
 | 
			
		||||
        BAD_REQUEST,
 | 
			
		||||
        String.format("Bot user [test-bot-user] is already used by [%s] bot.", botName));
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Test
 | 
			
		||||
  void put_ok_ifBotUserIsBotUserOfBot(TestInfo test) throws HttpResponseException {
 | 
			
		||||
    BotResourceTest botResourceTest = new BotResourceTest();
 | 
			
		||||
    String botName = "test-bot-ok";
 | 
			
		||||
    // create bot user
 | 
			
		||||
    CreateUser createBotUser = creatBotUserRequest("test-bot-user-ok", true).withBotName(botName);
 | 
			
		||||
    User botUser = updateEntity(createBotUser, CREATED, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    EntityReference botUserRef = Objects.requireNonNull(botUser).getEntityReference();
 | 
			
		||||
    // assign bot user to a bot
 | 
			
		||||
    CreateBot create = botResourceTest.createRequest(test).withBotUser(botUserRef).withName(botName);
 | 
			
		||||
    botResourceTest.createEntity(create, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    // put again user with same bot name
 | 
			
		||||
    CreateUser createDifferentBotUser = creatBotUserRequest("test-bot-user-ok", true).withBotName(botName);
 | 
			
		||||
    updateEntity(createDifferentBotUser, OK, ADMIN_AUTH_HEADERS);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private DecodedJWT decodedJWT(String token) {
 | 
			
		||||
    DecodedJWT jwt;
 | 
			
		||||
    try {
 | 
			
		||||
@ -862,4 +901,30 @@ public class UserResourceTest extends EntityResourceTest<User, CreateUser> {
 | 
			
		||||
      assertCommonFieldChange(fieldName, expected, actual);
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  @Override
 | 
			
		||||
  protected String getAllowedFields() {
 | 
			
		||||
    List<String> allowedFields = Entity.getAllowedFields(entityClass);
 | 
			
		||||
    allowedFields.removeAll(of(USER_PROTECTED_FIELDS.split(",")));
 | 
			
		||||
    return String.join(",", allowedFields);
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  public User createUser(String botName, boolean isBot) {
 | 
			
		||||
    try {
 | 
			
		||||
      CreateUser createUser = creatBotUserRequest(botName, isBot);
 | 
			
		||||
      return createEntity(createUser, ADMIN_AUTH_HEADERS);
 | 
			
		||||
    } catch (Exception ignore) {
 | 
			
		||||
      return null;
 | 
			
		||||
    }
 | 
			
		||||
  }
 | 
			
		||||
 | 
			
		||||
  private CreateUser creatBotUserRequest(String botUserName, boolean isBot) {
 | 
			
		||||
    return createRequest(botUserName, "", "", null)
 | 
			
		||||
        .withIsBot(isBot)
 | 
			
		||||
        .withIsAdmin(false)
 | 
			
		||||
        .withAuthenticationMechanism(
 | 
			
		||||
            new AuthenticationMechanism()
 | 
			
		||||
                .withAuthType(AuthenticationMechanism.AuthType.JWT)
 | 
			
		||||
                .withConfig(new JWTAuthMechanism().withJWTTokenExpiry(JWTTokenExpiry.Unlimited)));
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
@ -30,6 +30,10 @@
 | 
			
		||||
      "description": "When true indicates user is a bot with appropriate privileges",
 | 
			
		||||
      "type": "boolean"
 | 
			
		||||
    },
 | 
			
		||||
    "botName": {
 | 
			
		||||
      "description": "User bot name if we want to associate this bot with an specific bot",
 | 
			
		||||
      "type": "string"
 | 
			
		||||
    },
 | 
			
		||||
    "isAdmin": {
 | 
			
		||||
      "description": "When true indicates user is an administrator for the system with superuser privileges",
 | 
			
		||||
      "type": "boolean",
 | 
			
		||||
 | 
			
		||||
@ -6,6 +6,15 @@
 | 
			
		||||
  "type": "object",
 | 
			
		||||
  "javaType": "org.openmetadata.schema.entity.Bot",
 | 
			
		||||
  "javaInterfaces": ["org.openmetadata.schema.EntityInterface"],
 | 
			
		||||
  "definitions": {
 | 
			
		||||
    "botType": {
 | 
			
		||||
      "javaType": "org.openmetadata.schema.entity.BotType",
 | 
			
		||||
      "description": "Type of bot",
 | 
			
		||||
      "type": "string",
 | 
			
		||||
      "enum": ["bot", "ingestion-bot"],
 | 
			
		||||
      "default": "bot"
 | 
			
		||||
    }
 | 
			
		||||
  },
 | 
			
		||||
  "properties": {
 | 
			
		||||
    "id": {
 | 
			
		||||
      "description": "Unique identifier of a bot instance.",
 | 
			
		||||
@ -31,6 +40,10 @@
 | 
			
		||||
      "description": "Bot user created for this bot on behalf of which the bot performs all the operations, such as updating description, responding on the conversation threads, etc.",
 | 
			
		||||
      "$ref" : "../type/entityReference.json"
 | 
			
		||||
    },
 | 
			
		||||
    "botType" : {
 | 
			
		||||
      "$ref": "#/definitions/botType",
 | 
			
		||||
      "default": "bot"
 | 
			
		||||
    },
 | 
			
		||||
    "version": {
 | 
			
		||||
      "description": "Metadata version of the entity.",
 | 
			
		||||
      "$ref": "../type/entityHistory.json#/definitions/entityVersion"
 | 
			
		||||
 | 
			
		||||
@ -96,11 +96,6 @@
 | 
			
		||||
        }
 | 
			
		||||
      ]
 | 
			
		||||
    },
 | 
			
		||||
    "workflowBot": {
 | 
			
		||||
      "description": "OpenMetadata bot used for the ingestion",
 | 
			
		||||
      "type": "string",
 | 
			
		||||
      "default": "ingestion-bot"
 | 
			
		||||
    },
 | 
			
		||||
    "apiVersion": {
 | 
			
		||||
      "description": "OpenMetadata server API version to use.",
 | 
			
		||||
      "type": "string",
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user