diff --git a/datadog_sync/commands/shared/options.py b/datadog_sync/commands/shared/options.py index 3904f35b..68869aa4 100644 --- a/datadog_sync/commands/shared/options.py +++ b/datadog_sync/commands/shared/options.py @@ -326,6 +326,86 @@ def click_config_file_provider(ctx: Context, opts: CustomOptionClass, value: Non help=f"AWS session token, only used if --storage-type is '{constants.S3_STORAGE_TYPE}'", cls=CustomOptionClass, ), + # GCS options + option( + "--gcs-bucket-name", + envvar=constants.GCS_BUCKET_NAME, + required=False, + help=f"GCS bucket name, only used if --storage-type is '{constants.GCS_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--gcs-bucket-key-prefix-source", + default=constants.SOURCE_PATH_DEFAULT, + show_default=True, + envvar=constants.GCS_BUCKET_KEY_PREFIX_SOURCE, + required=False, + help=f"GCS bucket source key prefix, only used if --storage-type is '{constants.GCS_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--gcs-bucket-key-prefix-destination", + default=constants.DESTINATION_PATH_DEFAULT, + show_default=True, + envvar=constants.GCS_BUCKET_KEY_PREFIX_DESTINATION, + required=False, + help=f"GCS bucket destination key prefix, only used if --storage-type is '{constants.GCS_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--gcs-service-account-key-file", + envvar=constants.GCS_SERVICE_ACCOUNT_KEY_FILE, + required=False, + help=f"Path to GCS service account key file, only used if --storage-type is '{constants.GCS_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + # Azure options + option( + "--azure-container-name", + envvar=constants.AZURE_CONTAINER_NAME, + required=False, + help=f"Azure container name, only used if --storage-type is '{constants.AZURE_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--azure-container-key-prefix-source", + default=constants.SOURCE_PATH_DEFAULT, + show_default=True, + envvar=constants.AZURE_CONTAINER_KEY_PREFIX_SOURCE, + required=False, + help=f"Azure container source key prefix, only used if --storage-type is '{constants.AZURE_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--azure-container-key-prefix-destination", + default=constants.DESTINATION_PATH_DEFAULT, + show_default=True, + envvar=constants.AZURE_CONTAINER_KEY_PREFIX_DESTINATION, + required=False, + help=f"Azure container destination key prefix, only used if --storage-type is '{constants.AZURE_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--azure-storage-account-name", + envvar=constants.AZURE_STORAGE_ACCOUNT_NAME, + required=False, + help=f"Azure storage account name, only used if --storage-type is '{constants.AZURE_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--azure-storage-account-key", + envvar=constants.AZURE_STORAGE_ACCOUNT_KEY, + required=False, + help=f"Azure storage account key, only used if --storage-type is '{constants.AZURE_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), + option( + "--azure-storage-connection-string", + envvar=constants.AZURE_STORAGE_CONNECTION_STRING, + required=False, + help=f"Azure storage connection string, only used if --storage-type is '{constants.AZURE_STORAGE_TYPE}'", + cls=CustomOptionClass, + ), ] diff --git a/datadog_sync/constants.py b/datadog_sync/constants.py index cdd127f8..0ca54ec0 100644 --- a/datadog_sync/constants.py +++ b/datadog_sync/constants.py @@ -29,9 +29,13 @@ LOCAL_STORAGE_TYPE = "local" S3_STORAGE_TYPE = "s3" +GCS_STORAGE_TYPE = "gcs" +AZURE_STORAGE_TYPE = "azure" STORAGE_TYPES = [ LOCAL_STORAGE_TYPE, S3_STORAGE_TYPE, + GCS_STORAGE_TYPE, + AZURE_STORAGE_TYPE, ] DD_DESTINATION_RESOURCES_PATH = "DD_DESTINATION_RESOURCES_PATH" @@ -53,6 +57,30 @@ "aws_session_token", ] +# GCS env parameter names +GCS_BUCKET_NAME = "GCS_BUCKET_NAME" +GCS_BUCKET_KEY_PREFIX_SOURCE = "GCS_BUCKET_KEY_PREFIX_SOURCE" +GCS_BUCKET_KEY_PREFIX_DESTINATION = "GCS_BUCKET_KEY_PREFIX_DESTINATION" +GCS_SERVICE_ACCOUNT_KEY_FILE = "GCS_SERVICE_ACCOUNT_KEY_FILE" +GCS_CONFIG_PROPERTIES = [ + "gcs_bucket_name", + "gcs_service_account_key_file", +] + +# Azure env parameter names +AZURE_CONTAINER_NAME = "AZURE_CONTAINER_NAME" +AZURE_CONTAINER_KEY_PREFIX_SOURCE = "AZURE_CONTAINER_KEY_PREFIX_SOURCE" +AZURE_CONTAINER_KEY_PREFIX_DESTINATION = "AZURE_CONTAINER_KEY_PREFIX_DESTINATION" +AZURE_STORAGE_ACCOUNT_NAME = "AZURE_STORAGE_ACCOUNT_NAME" +AZURE_STORAGE_ACCOUNT_KEY = "AZURE_STORAGE_ACCOUNT_KEY" +AZURE_STORAGE_CONNECTION_STRING = "AZURE_STORAGE_CONNECTION_STRING" +AZURE_CONFIG_PROPERTIES = [ + "azure_container_name", + "azure_storage_account_name", + "azure_storage_account_key", + "azure_storage_connection_string", +] + # Default variables DEFAULT_API_URL = "https://api.datadoghq.com" diff --git a/datadog_sync/utils/configuration.py b/datadog_sync/utils/configuration.py index 9ac71d2c..b8f87260 100644 --- a/datadog_sync/utils/configuration.py +++ b/datadog_sync/utils/configuration.py @@ -13,10 +13,14 @@ from datadog_sync.constants import ( Command, AWS_CONFIG_PROPERTIES, + AZURE_CONFIG_PROPERTIES, + AZURE_STORAGE_TYPE, DESTINATION_PATH_DEFAULT, DESTINATION_PATH_PARAM, FALSE, FORCE, + GCS_CONFIG_PROPERTIES, + GCS_STORAGE_TYPE, LOCAL_STORAGE_TYPE, LOGGER_NAME, RESOURCE_PER_FILE, @@ -230,6 +234,38 @@ def build_config(cmd: Command, **kwargs: Optional[Any]) -> Configuration: if not property_value: logger.warning(f"Missing AWS configuration parameter: {aws_config_property}") config[aws_config_property] = property_value + elif storage_type == GCS_STORAGE_TYPE: + logger.info("Using GCS to store state files") + storage_type = StorageType.GCS_BUCKET + + local_source_resources_path = kwargs.get(SOURCE_PATH_PARAM, SOURCE_PATH_DEFAULT) + source_resources_path = kwargs.get("gcs_bucket_key_prefix_source", local_source_resources_path) + + local_destination_resources_path = kwargs.get(DESTINATION_PATH_PARAM, DESTINATION_PATH_DEFAULT) + destination_resources_path = kwargs.get("gcs_bucket_key_prefix_destination", local_destination_resources_path) + + for gcs_config_property in GCS_CONFIG_PROPERTIES: + property_value = kwargs.get(gcs_config_property, None) + if not property_value: + logger.warning(f"Missing GCS configuration parameter: {gcs_config_property}") + config[gcs_config_property] = property_value + elif storage_type == AZURE_STORAGE_TYPE: + logger.info("Using Azure Blob Storage to store state files") + storage_type = StorageType.AZURE_BLOB_CONTAINER + + local_source_resources_path = kwargs.get(SOURCE_PATH_PARAM, SOURCE_PATH_DEFAULT) + source_resources_path = kwargs.get("azure_container_key_prefix_source", local_source_resources_path) + + local_destination_resources_path = kwargs.get(DESTINATION_PATH_PARAM, DESTINATION_PATH_DEFAULT) + destination_resources_path = kwargs.get( + "azure_container_key_prefix_destination", local_destination_resources_path + ) + + for azure_config_property in AZURE_CONFIG_PROPERTIES: + property_value = kwargs.get(azure_config_property, None) + if not property_value: + logger.warning(f"Missing Azure configuration parameter: {azure_config_property}") + config[azure_config_property] = property_value elif storage_type == LOCAL_STORAGE_TYPE: logger.info("Using local filesystem to store state files") storage_type = StorageType.LOCAL_FILE diff --git a/datadog_sync/utils/state.py b/datadog_sync/utils/state.py index 77ee38e6..85a9dd79 100644 --- a/datadog_sync/utils/state.py +++ b/datadog_sync/utils/state.py @@ -14,6 +14,8 @@ ) from datadog_sync.utils.storage._base_storage import BaseStorage, StorageData from datadog_sync.utils.storage.aws_s3_bucket import AWSS3Bucket +from datadog_sync.utils.storage.azure_blob_container import AzureBlobContainer +from datadog_sync.utils.storage.gcs_bucket import GCSBucket from datadog_sync.utils.storage.local_file import LocalFile from datadog_sync.utils.storage.storage_types import StorageType @@ -39,6 +41,26 @@ def __init__(self, type_: StorageType = StorageType.LOCAL_FILE, **kwargs: object config=config, resource_per_file=resource_per_file, ) + elif type_ == StorageType.GCS_BUCKET: + config = kwargs.get("config", {}) + if not config: + raise ValueError("GCS configuration not found") + self._storage: BaseStorage = GCSBucket( + source_resources_path=source_resources_path, + destination_resources_path=destination_resources_path, + config=config, + resource_per_file=resource_per_file, + ) + elif type_ == StorageType.AZURE_BLOB_CONTAINER: + config = kwargs.get("config", {}) + if not config: + raise ValueError("Azure configuration not found") + self._storage: BaseStorage = AzureBlobContainer( + source_resources_path=source_resources_path, + destination_resources_path=destination_resources_path, + config=config, + resource_per_file=resource_per_file, + ) else: raise NotImplementedError(f"Storage type {type_} not implemented") diff --git a/datadog_sync/utils/storage/aws_s3_bucket.py b/datadog_sync/utils/storage/aws_s3_bucket.py index d79ead59..a1d050dd 100644 --- a/datadog_sync/utils/storage/aws_s3_bucket.py +++ b/datadog_sync/utils/storage/aws_s3_bucket.py @@ -54,6 +54,8 @@ def __init__( self.client = boto3.client("s3") self.bucket_name = config.get("aws_bucket_name", "") + if not self.bucket_name: + raise ValueError("AWS S3 bucket name is required") def get(self, origin: Origin) -> StorageData: log.info("AWS S3 get called") diff --git a/datadog_sync/utils/storage/azure_blob_container.py b/datadog_sync/utils/storage/azure_blob_container.py new file mode 100644 index 00000000..0b3b8350 --- /dev/null +++ b/datadog_sync/utils/storage/azure_blob_container.py @@ -0,0 +1,115 @@ +# Unless explicitly stated otherwise all files in this repository are licensed +# under the 3-clause BSD style license (see LICENSE). +# This product includes software developed at Datadog (https://www.datadoghq.com/). +# Copyright 2019 Datadog, Inc. + +import json +import logging + +from azure.storage.blob import BlobServiceClient, ContainerClient +from azure.identity import DefaultAzureCredential + +from datadog_sync.constants import ( + Origin, + DESTINATION_PATH_DEFAULT, + LOGGER_NAME, + SOURCE_PATH_DEFAULT, +) +from datadog_sync.utils.storage._base_storage import BaseStorage, StorageData + + +log = logging.getLogger(LOGGER_NAME) + + +class AzureBlobContainer(BaseStorage): + + def __init__( + self, + source_resources_path=SOURCE_PATH_DEFAULT, + destination_resources_path=DESTINATION_PATH_DEFAULT, + resource_per_file=False, + config=None, + ) -> None: + log.info("Azure Blob Storage init called") + super().__init__() + self.source_resources_path = source_resources_path + self.destination_resources_path = destination_resources_path + self.resource_per_file = resource_per_file + if not config: + raise ValueError("No Azure configuration passed in") + + container_name = config.get("azure_container_name", "") + if not container_name: + raise ValueError("Azure container name is required") + connection_string = config.get("azure_storage_connection_string", None) + account_name = config.get("azure_storage_account_name", None) + account_key = config.get("azure_storage_account_key", None) + + if connection_string: + log.info("Azure Blob Storage configured with connection string") + self.container_client = ContainerClient.from_connection_string( + conn_str=connection_string, + container_name=container_name, + ) + elif account_name and account_key: + log.info("Azure Blob Storage configured with account name and key") + account_url = f"https://{account_name}.blob.core.windows.net" + blob_service_client = BlobServiceClient(account_url=account_url, credential=account_key) + self.container_client = blob_service_client.get_container_client(container_name) + elif account_name: + log.info("Azure Blob Storage configured with default credentials") + account_url = f"https://{account_name}.blob.core.windows.net" + blob_service_client = BlobServiceClient(account_url=account_url, credential=DefaultAzureCredential()) + self.container_client = blob_service_client.get_container_client(container_name) + else: + raise ValueError("Azure storage requires at least a connection string or storage account name") + + def get(self, origin: Origin) -> StorageData: + log.info("Azure Blob Storage get called") + data = StorageData() + + if origin in [Origin.SOURCE, Origin.ALL]: + for blob in self.container_client.list_blobs(name_starts_with=self.source_resources_path): + if blob.name.endswith(".json"): + resource_type = blob.name.split(".")[0].split("/")[-1] + try: + content = self.container_client.download_blob(blob.name).readall().decode("utf-8") + data.source[resource_type].update(json.loads(content)) + except json.decoder.JSONDecodeError: + log.warning(f"invalid json in azure source resource file: {resource_type}") + + if origin in [Origin.DESTINATION, Origin.ALL]: + for blob in self.container_client.list_blobs(name_starts_with=self.destination_resources_path): + if blob.name.endswith(".json"): + resource_type = blob.name.split(".")[0].split("/")[-1] + try: + content = self.container_client.download_blob(blob.name).readall().decode("utf-8") + data.destination[resource_type].update(json.loads(content)) + except json.decoder.JSONDecodeError: + log.warning(f"invalid json in azure destination resource file: {resource_type}") + + return data + + def put(self, origin: Origin, data: StorageData) -> None: + log.info("Azure Blob Storage put called") + if origin in [Origin.SOURCE, Origin.ALL]: + for resource_type, resource_data in data.source.items(): + base_key = f"{self.source_resources_path}/{resource_type}" + if self.resource_per_file: + for _id, resource in resource_data.items(): + key = f"{base_key}.{_id}.json" + self.container_client.upload_blob(name=key, data=json.dumps({_id: resource}), overwrite=True) + else: + key = f"{base_key}.json" + self.container_client.upload_blob(name=key, data=json.dumps(resource_data), overwrite=True) + + if origin in [Origin.DESTINATION, Origin.ALL]: + for resource_type, resource_data in data.destination.items(): + base_key = f"{self.destination_resources_path}/{resource_type}" + if self.resource_per_file: + for _id, resource in resource_data.items(): + key = f"{base_key}.{_id}.json" + self.container_client.upload_blob(name=key, data=json.dumps({_id: resource}), overwrite=True) + else: + key = f"{base_key}.json" + self.container_client.upload_blob(name=key, data=json.dumps(resource_data), overwrite=True) diff --git a/datadog_sync/utils/storage/gcs_bucket.py b/datadog_sync/utils/storage/gcs_bucket.py new file mode 100644 index 00000000..fbf98bdb --- /dev/null +++ b/datadog_sync/utils/storage/gcs_bucket.py @@ -0,0 +1,105 @@ +# Unless explicitly stated otherwise all files in this repository are licensed +# under the 3-clause BSD style license (see LICENSE). +# This product includes software developed at Datadog (https://www.datadoghq.com/). +# Copyright 2019 Datadog, Inc. + +import json +import logging + +from google.cloud import storage as gcs_storage + +from datadog_sync.constants import ( + Origin, + DESTINATION_PATH_DEFAULT, + LOGGER_NAME, + SOURCE_PATH_DEFAULT, +) +from datadog_sync.utils.storage._base_storage import BaseStorage, StorageData + + +log = logging.getLogger(LOGGER_NAME) + + +class GCSBucket(BaseStorage): + + def __init__( + self, + source_resources_path=SOURCE_PATH_DEFAULT, + destination_resources_path=DESTINATION_PATH_DEFAULT, + resource_per_file=False, + config=None, + ) -> None: + log.info("GCS init called") + super().__init__() + self.source_resources_path = source_resources_path + self.destination_resources_path = destination_resources_path + self.resource_per_file = resource_per_file + if not config: + raise ValueError("No GCS configuration passed in") + + key_file = config.get("gcs_service_account_key_file", None) + if key_file: + log.info("GCS configured with service account key file") + self.client = gcs_storage.Client.from_service_account_json(key_file) + else: + log.info("GCS configured with application default credentials") + self.client = gcs_storage.Client() + + bucket_name = config.get("gcs_bucket_name", "") + if not bucket_name: + raise ValueError("GCS bucket name is required") + self.bucket = self.client.bucket(bucket_name) + + def get(self, origin: Origin) -> StorageData: + log.info("GCS get called") + data = StorageData() + + if origin in [Origin.SOURCE, Origin.ALL]: + for blob in self.bucket.list_blobs(prefix=self.source_resources_path): + if blob.name.endswith(".json"): + resource_type = blob.name.split(".")[0].split("/")[-1] + try: + content = blob.download_as_text() + data.source[resource_type].update(json.loads(content)) + except json.decoder.JSONDecodeError: + log.warning(f"invalid json in gcs source resource file: {resource_type}") + + if origin in [Origin.DESTINATION, Origin.ALL]: + for blob in self.bucket.list_blobs(prefix=self.destination_resources_path): + if blob.name.endswith(".json"): + resource_type = blob.name.split(".")[0].split("/")[-1] + try: + content = blob.download_as_text() + data.destination[resource_type].update(json.loads(content)) + except json.decoder.JSONDecodeError: + log.warning(f"invalid json in gcs destination resource file: {resource_type}") + + return data + + def put(self, origin: Origin, data: StorageData) -> None: + log.info("GCS put called") + if origin in [Origin.SOURCE, Origin.ALL]: + for resource_type, resource_data in data.source.items(): + base_key = f"{self.source_resources_path}/{resource_type}" + if self.resource_per_file: + for _id, resource in resource_data.items(): + key = f"{base_key}.{_id}.json" + self.bucket.blob(key).upload_from_string( + json.dumps({_id: resource}), content_type="application/json" + ) + else: + key = f"{base_key}.json" + self.bucket.blob(key).upload_from_string(json.dumps(resource_data), content_type="application/json") + + if origin in [Origin.DESTINATION, Origin.ALL]: + for resource_type, resource_data in data.destination.items(): + base_key = f"{self.destination_resources_path}/{resource_type}" + if self.resource_per_file: + for _id, resource in resource_data.items(): + key = f"{base_key}.{_id}.json" + self.bucket.blob(key).upload_from_string( + json.dumps({_id: resource}), content_type="application/json" + ) + else: + key = f"{base_key}.json" + self.bucket.blob(key).upload_from_string(json.dumps(resource_data), content_type="application/json") diff --git a/datadog_sync/utils/storage/storage_types.py b/datadog_sync/utils/storage/storage_types.py index 5b779027..77a9ca4f 100644 --- a/datadog_sync/utils/storage/storage_types.py +++ b/datadog_sync/utils/storage/storage_types.py @@ -9,3 +9,5 @@ class StorageType(Enum): LOCAL_FILE = 1 AWS_S3_BUCKET = 2 + GCS_BUCKET = 3 + AZURE_BLOB_CONTAINER = 4 diff --git a/setup.cfg b/setup.cfg index c8ca534d..0b06fe83 100644 --- a/setup.cfg +++ b/setup.cfg @@ -39,6 +39,9 @@ install_requires = tqdm==4.66.3 certifi>=2022.12.7 python-dateutil + google-cloud-storage>=2.14.0 + azure-storage-blob>=12.19.0 + azure-identity>=1.15.0 setup_requires = setuptools>=67.6.0 setuptools_scm diff --git a/tests/unit/test_aws_s3_bucket.py b/tests/unit/test_aws_s3_bucket.py new file mode 100644 index 00000000..077840ce --- /dev/null +++ b/tests/unit/test_aws_s3_bucket.py @@ -0,0 +1,238 @@ +# Unless explicitly stated otherwise all files in this repository are licensed +# under the 3-clause BSD style license (see LICENSE). +# This product includes software developed at Datadog (https://www.datadoghq.com/). +# Copyright 2019 Datadog, Inc. + +import io +import json +from unittest.mock import MagicMock, patch + +import pytest + +from datadog_sync.constants import Origin +from datadog_sync.utils.storage.aws_s3_bucket import AWSS3Bucket + + +@pytest.fixture +def mock_s3_client(): + with patch("datadog_sync.utils.storage.aws_s3_bucket.boto3") as mock_boto3: + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + yield mock_boto3, mock_client + + +class TestAWSS3Bucket: + def test_init_with_explicit_credentials(self, mock_s3_client): + mock_boto3, mock_client = mock_s3_client + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + + mock_boto3.client.assert_called_once_with( + "s3", + region_name="us-east-1", + aws_access_key_id="AKID", + aws_secret_access_key="SECRET", + aws_session_token="", + ) + assert bucket.bucket_name == "test-bucket" + + def test_init_with_default_credentials(self, mock_s3_client): + mock_boto3, mock_client = mock_s3_client + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": None, + "aws_access_key_id": None, + "aws_secret_access_key": None, + "aws_session_token": None, + } + ) + + mock_boto3.client.assert_called_once_with("s3") + assert bucket.bucket_name == "test-bucket" + + def test_init_no_config(self): + with pytest.raises(ValueError, match="No S3 configuration passed in"): + AWSS3Bucket(config=None) + + def test_init_missing_bucket_name(self, mock_s3_client): + with pytest.raises(ValueError, match="AWS S3 bucket name is required"): + AWSS3Bucket( + config={ + "aws_bucket_name": "", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + + def test_get_source(self, mock_s3_client): + _, mock_client = mock_s3_client + + mock_client.list_objects_v2.return_value = { + "Contents": [{"Key": "resources/source/monitors.json"}], + "IsTruncated": False, + } + mock_client.get_object.return_value = { + "Body": io.BytesIO(json.dumps({"id1": {"name": "monitor1"}}).encode("utf-8")) + } + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + data = bucket.get(Origin.SOURCE) + + assert dict(data.source["monitors"]) == {"id1": {"name": "monitor1"}} + assert len(data.destination) == 0 + + def test_get_destination(self, mock_s3_client): + _, mock_client = mock_s3_client + + mock_client.list_objects_v2.return_value = { + "Contents": [{"Key": "resources/destination/dashboards.json"}], + "IsTruncated": False, + } + mock_client.get_object.return_value = { + "Body": io.BytesIO(json.dumps({"id2": {"title": "dash1"}}).encode("utf-8")) + } + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + data = bucket.get(Origin.DESTINATION) + + assert len(data.source) == 0 + assert dict(data.destination["dashboards"]) == {"id2": {"title": "dash1"}} + + def test_get_all(self, mock_s3_client): + _, mock_client = mock_s3_client + + mock_client.list_objects_v2.side_effect = [ + { + "Contents": [{"Key": "resources/source/monitors.json"}], + "IsTruncated": False, + }, + { + "Contents": [{"Key": "resources/destination/dashboards.json"}], + "IsTruncated": False, + }, + ] + mock_client.get_object.side_effect = [ + {"Body": io.BytesIO(json.dumps({"id1": {"name": "monitor1"}}).encode("utf-8"))}, + {"Body": io.BytesIO(json.dumps({"id2": {"title": "dash1"}}).encode("utf-8"))}, + ] + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + data = bucket.get(Origin.ALL) + + assert dict(data.source["monitors"]) == {"id1": {"name": "monitor1"}} + assert dict(data.destination["dashboards"]) == {"id2": {"title": "dash1"}} + + def test_get_skips_non_json(self, mock_s3_client): + _, mock_client = mock_s3_client + + mock_client.list_objects_v2.return_value = { + "Contents": [{"Key": "resources/source/readme.txt"}], + "IsTruncated": False, + } + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + data = bucket.get(Origin.SOURCE) + + assert len(data.source) == 0 + mock_client.get_object.assert_not_called() + + def test_put_single_file(self, mock_s3_client): + _, mock_client = mock_s3_client + + mock_client.list_objects_v2.return_value = {"IsTruncated": False} + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + } + ) + data = bucket.get(Origin.SOURCE) + data.source["monitors"] = {"id1": {"name": "monitor1"}} + + bucket.put(Origin.SOURCE, data) + + mock_client.put_object.assert_called_once_with( + Body=bytes(json.dumps({"id1": {"name": "monitor1"}}), "UTF-8"), + Bucket="test-bucket", + Key="resources/source/monitors.json", + ) + + def test_put_resource_per_file(self, mock_s3_client): + _, mock_client = mock_s3_client + + mock_client.list_objects_v2.return_value = {"IsTruncated": False} + + bucket = AWSS3Bucket( + config={ + "aws_bucket_name": "test-bucket", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKID", + "aws_secret_access_key": "SECRET", + "aws_session_token": "", + }, + resource_per_file=True, + ) + data = bucket.get(Origin.SOURCE) + data.source["monitors"] = {"id1": {"name": "mon1"}, "id2": {"name": "mon2"}} + + bucket.put(Origin.SOURCE, data) + + assert mock_client.put_object.call_count == 2 + mock_client.put_object.assert_any_call( + Body=bytes(json.dumps({"id1": {"name": "mon1"}}), "UTF-8"), + Bucket="test-bucket", + Key="resources/source/monitors.id1.json", + ) + mock_client.put_object.assert_any_call( + Body=bytes(json.dumps({"id2": {"name": "mon2"}}), "UTF-8"), + Bucket="test-bucket", + Key="resources/source/monitors.id2.json", + ) diff --git a/tests/unit/test_azure_blob_container.py b/tests/unit/test_azure_blob_container.py new file mode 100644 index 00000000..bb8d5035 --- /dev/null +++ b/tests/unit/test_azure_blob_container.py @@ -0,0 +1,268 @@ +# Unless explicitly stated otherwise all files in this repository are licensed +# under the 3-clause BSD style license (see LICENSE). +# This product includes software developed at Datadog (https://www.datadoghq.com/). +# Copyright 2019 Datadog, Inc. + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from datadog_sync.constants import Origin +from datadog_sync.utils.storage.azure_blob_container import AzureBlobContainer + + +@pytest.fixture +def mock_azure_connection_string(): + with patch("datadog_sync.utils.storage.azure_blob_container.ContainerClient") as mock_container_cls: + mock_container = MagicMock() + mock_container_cls.from_connection_string.return_value = mock_container + yield mock_container_cls, mock_container + + +@pytest.fixture +def mock_azure_account_key(): + with patch("datadog_sync.utils.storage.azure_blob_container.BlobServiceClient") as mock_bsc_cls: + mock_bsc = MagicMock() + mock_bsc_cls.return_value = mock_bsc + mock_container = MagicMock() + mock_bsc.get_container_client.return_value = mock_container + yield mock_bsc_cls, mock_bsc, mock_container + + +class TestAzureBlobContainer: + def test_init_with_connection_string(self, mock_azure_connection_string): + mock_container_cls, mock_container = mock_azure_connection_string + + AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "DefaultEndpointsProtocol=https;AccountName=test", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + + mock_container_cls.from_connection_string.assert_called_once_with( + conn_str="DefaultEndpointsProtocol=https;AccountName=test", + container_name="test-container", + ) + + def test_init_with_account_key(self, mock_azure_account_key): + mock_bsc_cls, mock_bsc, mock_container = mock_azure_account_key + + AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": None, + "azure_storage_account_name": "myaccount", + "azure_storage_account_key": "mykey", + } + ) + + mock_bsc_cls.assert_called_once_with( + account_url="https://myaccount.blob.core.windows.net", + credential="mykey", + ) + mock_bsc.get_container_client.assert_called_once_with("test-container") + + def test_init_with_default_credentials(self): + with patch("datadog_sync.utils.storage.azure_blob_container.BlobServiceClient") as mock_bsc_cls, patch( + "datadog_sync.utils.storage.azure_blob_container.DefaultAzureCredential" + ) as mock_cred_cls: + mock_bsc = MagicMock() + mock_bsc_cls.return_value = mock_bsc + mock_container = MagicMock() + mock_bsc.get_container_client.return_value = mock_container + mock_cred = MagicMock() + mock_cred_cls.return_value = mock_cred + + AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": None, + "azure_storage_account_name": "myaccount", + "azure_storage_account_key": None, + } + ) + + mock_bsc_cls.assert_called_once_with( + account_url="https://myaccount.blob.core.windows.net", + credential=mock_cred, + ) + + def test_init_no_config(self): + with pytest.raises(ValueError, match="No Azure configuration passed in"): + AzureBlobContainer(config=None) + + def test_init_missing_container_name(self, mock_azure_connection_string): + with pytest.raises(ValueError, match="Azure container name is required"): + AzureBlobContainer( + config={ + "azure_container_name": "", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + + def test_init_missing_account_info(self): + with pytest.raises(ValueError, match="Azure storage requires"): + AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": None, + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + + def test_get_source(self, mock_azure_connection_string): + _, mock_container = mock_azure_connection_string + + blob1 = MagicMock() + blob1.name = "resources/source/monitors.json" + mock_container.list_blobs.return_value = [blob1] + + download_mock = MagicMock() + download_mock.readall.return_value = json.dumps({"id1": {"name": "monitor1"}}).encode("utf-8") + mock_container.download_blob.return_value = download_mock + + container = AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + data = container.get(Origin.SOURCE) + + assert dict(data.source["monitors"]) == {"id1": {"name": "monitor1"}} + assert len(data.destination) == 0 + + def test_get_destination(self, mock_azure_connection_string): + _, mock_container = mock_azure_connection_string + + blob1 = MagicMock() + blob1.name = "resources/destination/dashboards.json" + mock_container.list_blobs.return_value = [blob1] + + download_mock = MagicMock() + download_mock.readall.return_value = json.dumps({"id2": {"title": "dash1"}}).encode("utf-8") + mock_container.download_blob.return_value = download_mock + + container = AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + data = container.get(Origin.DESTINATION) + + assert len(data.source) == 0 + assert dict(data.destination["dashboards"]) == {"id2": {"title": "dash1"}} + + def test_get_all(self, mock_azure_connection_string): + _, mock_container = mock_azure_connection_string + + source_blob = MagicMock() + source_blob.name = "resources/source/monitors.json" + + dest_blob = MagicMock() + dest_blob.name = "resources/destination/dashboards.json" + + mock_container.list_blobs.side_effect = [[source_blob], [dest_blob]] + + source_download = MagicMock() + source_download.readall.return_value = json.dumps({"id1": {"name": "monitor1"}}).encode("utf-8") + dest_download = MagicMock() + dest_download.readall.return_value = json.dumps({"id2": {"title": "dash1"}}).encode("utf-8") + mock_container.download_blob.side_effect = [source_download, dest_download] + + container = AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + data = container.get(Origin.ALL) + + assert dict(data.source["monitors"]) == {"id1": {"name": "monitor1"}} + assert dict(data.destination["dashboards"]) == {"id2": {"title": "dash1"}} + + def test_get_skips_non_json(self, mock_azure_connection_string): + _, mock_container = mock_azure_connection_string + + blob1 = MagicMock() + blob1.name = "resources/source/readme.txt" + mock_container.list_blobs.return_value = [blob1] + + container = AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + data = container.get(Origin.SOURCE) + + assert len(data.source) == 0 + + def test_put_single_file(self, mock_azure_connection_string): + _, mock_container = mock_azure_connection_string + mock_container.list_blobs.return_value = [] + + container = AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + } + ) + data = container.get(Origin.SOURCE) + data.source["monitors"] = {"id1": {"name": "monitor1"}} + + container.put(Origin.SOURCE, data) + + mock_container.upload_blob.assert_called_once_with( + name="resources/source/monitors.json", + data=json.dumps({"id1": {"name": "monitor1"}}), + overwrite=True, + ) + + def test_put_resource_per_file(self, mock_azure_connection_string): + _, mock_container = mock_azure_connection_string + mock_container.list_blobs.return_value = [] + + container = AzureBlobContainer( + config={ + "azure_container_name": "test-container", + "azure_storage_connection_string": "connstr", + "azure_storage_account_name": None, + "azure_storage_account_key": None, + }, + resource_per_file=True, + ) + data = container.get(Origin.SOURCE) + data.source["monitors"] = {"id1": {"name": "mon1"}, "id2": {"name": "mon2"}} + + container.put(Origin.SOURCE, data) + + assert mock_container.upload_blob.call_count == 2 + mock_container.upload_blob.assert_any_call( + name="resources/source/monitors.id1.json", + data=json.dumps({"id1": {"name": "mon1"}}), + overwrite=True, + ) + mock_container.upload_blob.assert_any_call( + name="resources/source/monitors.id2.json", + data=json.dumps({"id2": {"name": "mon2"}}), + overwrite=True, + ) diff --git a/tests/unit/test_gcs_bucket.py b/tests/unit/test_gcs_bucket.py new file mode 100644 index 00000000..9a705535 --- /dev/null +++ b/tests/unit/test_gcs_bucket.py @@ -0,0 +1,154 @@ +# Unless explicitly stated otherwise all files in this repository are licensed +# under the 3-clause BSD style license (see LICENSE). +# This product includes software developed at Datadog (https://www.datadoghq.com/). +# Copyright 2019 Datadog, Inc. + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from datadog_sync.constants import Origin +from datadog_sync.utils.storage.gcs_bucket import GCSBucket + + +@pytest.fixture +def mock_gcs_client(): + with patch("datadog_sync.utils.storage.gcs_bucket.gcs_storage") as mock_storage: + mock_client = MagicMock() + mock_storage.Client.return_value = mock_client + mock_bucket = MagicMock() + mock_client.bucket.return_value = mock_bucket + yield mock_storage, mock_client, mock_bucket + + +class TestGCSBucket: + def test_init_with_service_account_key(self): + with patch("datadog_sync.utils.storage.gcs_bucket.gcs_storage") as mock_storage: + mock_client = MagicMock() + mock_storage.Client.from_service_account_json.return_value = mock_client + mock_bucket = MagicMock() + mock_client.bucket.return_value = mock_bucket + + GCSBucket( + config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": "/path/to/key.json"} + ) + + mock_storage.Client.from_service_account_json.assert_called_once_with("/path/to/key.json") + mock_client.bucket.assert_called_once_with("test-bucket") + + def test_init_with_default_credentials(self, mock_gcs_client): + mock_storage, mock_client, mock_bucket = mock_gcs_client + + GCSBucket(config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}) + + mock_storage.Client.assert_called_once() + mock_client.bucket.assert_called_once_with("test-bucket") + + def test_init_no_config(self): + with pytest.raises(ValueError, match="No GCS configuration passed in"): + GCSBucket(config=None) + + def test_init_missing_bucket_name(self, mock_gcs_client): + with pytest.raises(ValueError, match="GCS bucket name is required"): + GCSBucket(config={"gcs_bucket_name": "", "gcs_service_account_key_file": None}) + + def test_get_source(self, mock_gcs_client): + _, _, mock_bucket = mock_gcs_client + + blob1 = MagicMock() + blob1.name = "resources/source/monitors.json" + blob1.download_as_text.return_value = json.dumps({"id1": {"name": "monitor1"}}) + + mock_bucket.list_blobs.return_value = [blob1] + + bucket = GCSBucket(config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}) + data = bucket.get(Origin.SOURCE) + + assert dict(data.source["monitors"]) == {"id1": {"name": "monitor1"}} + assert len(data.destination) == 0 + + def test_get_destination(self, mock_gcs_client): + _, _, mock_bucket = mock_gcs_client + + blob1 = MagicMock() + blob1.name = "resources/destination/dashboards.json" + blob1.download_as_text.return_value = json.dumps({"id2": {"title": "dash1"}}) + + mock_bucket.list_blobs.return_value = [blob1] + + bucket = GCSBucket(config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}) + data = bucket.get(Origin.DESTINATION) + + assert len(data.source) == 0 + assert dict(data.destination["dashboards"]) == {"id2": {"title": "dash1"}} + + def test_get_all(self, mock_gcs_client): + _, _, mock_bucket = mock_gcs_client + + source_blob = MagicMock() + source_blob.name = "resources/source/monitors.json" + source_blob.download_as_text.return_value = json.dumps({"id1": {"name": "monitor1"}}) + + dest_blob = MagicMock() + dest_blob.name = "resources/destination/dashboards.json" + dest_blob.download_as_text.return_value = json.dumps({"id2": {"title": "dash1"}}) + + mock_bucket.list_blobs.side_effect = [[source_blob], [dest_blob]] + + bucket = GCSBucket(config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}) + data = bucket.get(Origin.ALL) + + assert dict(data.source["monitors"]) == {"id1": {"name": "monitor1"}} + assert dict(data.destination["dashboards"]) == {"id2": {"title": "dash1"}} + + def test_get_skips_non_json(self, mock_gcs_client): + _, _, mock_bucket = mock_gcs_client + + blob1 = MagicMock() + blob1.name = "resources/source/readme.txt" + + mock_bucket.list_blobs.return_value = [blob1] + + bucket = GCSBucket(config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}) + data = bucket.get(Origin.SOURCE) + + assert len(data.source) == 0 + + def test_put_single_file(self, mock_gcs_client): + _, _, mock_bucket = mock_gcs_client + + mock_blob = MagicMock() + mock_bucket.blob.return_value = mock_blob + mock_bucket.list_blobs.return_value = [] + + bucket = GCSBucket(config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}) + data = bucket.get(Origin.SOURCE) + data.source["monitors"] = {"id1": {"name": "monitor1"}} + + bucket.put(Origin.SOURCE, data) + + mock_bucket.blob.assert_called_with("resources/source/monitors.json") + mock_blob.upload_from_string.assert_called_once_with( + json.dumps({"id1": {"name": "monitor1"}}), content_type="application/json" + ) + + def test_put_resource_per_file(self, mock_gcs_client): + _, _, mock_bucket = mock_gcs_client + + mock_blob = MagicMock() + mock_bucket.blob.return_value = mock_blob + mock_bucket.list_blobs.return_value = [] + + bucket = GCSBucket( + config={"gcs_bucket_name": "test-bucket", "gcs_service_account_key_file": None}, + resource_per_file=True, + ) + data = bucket.get(Origin.SOURCE) + data.source["monitors"] = {"id1": {"name": "mon1"}, "id2": {"name": "mon2"}} + + bucket.put(Origin.SOURCE, data) + + assert mock_bucket.blob.call_count == 2 + mock_bucket.blob.assert_any_call("resources/source/monitors.id1.json") + mock_bucket.blob.assert_any_call("resources/source/monitors.id2.json")