diff --git a/src/spaceone/cost_analysis/info/budget_usage_info.py b/src/spaceone/cost_analysis/info/budget_usage_info.py index 66a34af5..01a40dce 100644 --- a/src/spaceone/cost_analysis/info/budget_usage_info.py +++ b/src/spaceone/cost_analysis/info/budget_usage_info.py @@ -2,11 +2,23 @@ from spaceone.api.cost_analysis.v1 import budget_usage_pb2 from spaceone.core.pygrpc.message_type import * from spaceone.core import utils -from spaceone.cost_analysis.model.budget_usage_model import BudgetUsage +from spaceone.cost_analysis.model.budget_usage_model import BudgetUsage, ProviderFilter __all__ = ['BudgetUsageInfo', 'BudgetUsagesInfo'] +def ProviderFilterInfo(provider_filter_vo: ProviderFilter): + if provider_filter_vo is None: + return None + + info = { + 'state': provider_filter_vo.state, + 'providers': list(provider_filter_vo.providers) + } + + return budget_usage_pb2.BudgetUsageProviderFilter(**info) + + def BudgetUsageInfo(budget_usage_vo: BudgetUsage, minimal=False): info = { 'budget_id': budget_usage_vo.budget_id, @@ -15,6 +27,7 @@ def BudgetUsageInfo(budget_usage_vo: BudgetUsage, minimal=False): 'cost': budget_usage_vo.cost, 'limit': budget_usage_vo.limit, 'currency': budget_usage_vo.currency, + 'provider_filter': ProviderFilterInfo(budget_usage_vo.provider_filter), 'project_id': budget_usage_vo.project_id, 'project_group_id': budget_usage_vo.project_group_id, 'data_source_id': budget_usage_vo.data_source_id, diff --git a/src/spaceone/cost_analysis/manager/budget_usage_manager.py b/src/spaceone/cost_analysis/manager/budget_usage_manager.py index 11de3340..8544e3b5 100644 --- a/src/spaceone/cost_analysis/manager/budget_usage_manager.py +++ b/src/spaceone/cost_analysis/manager/budget_usage_manager.py @@ -35,6 +35,7 @@ def create_budget_usages(self, budget_vo: Budget): 'cost': 0, 'limit': limit_per_month, 'currency': budget_vo.currency, + 'provider_filter': budget_vo.provider_filter.to_dict(), 'budget': budget_vo, 'project_id': budget_vo.project_id, 'project_group_id': budget_vo.project_group_id, @@ -53,6 +54,7 @@ def create_budget_usages(self, budget_vo: Budget): 'cost': 0, 'limit': planned_limit.limit, 'currency': budget_vo.currency, + 'provider_filter': budget_vo.provider_filter.to_dict(), 'budget': budget_vo, 'project_id': budget_vo.project_id, 'project_group_id': budget_vo.project_group_id, diff --git a/src/spaceone/cost_analysis/model/budget_usage_model.py b/src/spaceone/cost_analysis/model/budget_usage_model.py index 553c42a5..bbed8569 100644 --- a/src/spaceone/cost_analysis/model/budget_usage_model.py +++ b/src/spaceone/cost_analysis/model/budget_usage_model.py @@ -3,6 +3,14 @@ from spaceone.core.model.mongo_model import MongoModel +class ProviderFilter(EmbeddedDocument): + state = StringField(max_length=20, default='ENABLED', choices=('ENABLED', 'DISABLED')) + providers = ListField(StringField(), default=[]) + + def to_dict(self): + return dict(self.to_mongo()) + + class BudgetUsage(MongoModel): budget_id = StringField(max_length=40, required=True) name = StringField(max_length=255, default='') @@ -10,6 +18,7 @@ class BudgetUsage(MongoModel): cost = FloatField(required=True) limit = FloatField(required=True) currency = StringField(default=None, null=True) + provider_filter = EmbeddedDocumentField(ProviderFilter, required=True) budget = ReferenceField('Budget', reverse_delete_rule=CASCADE) project_id = StringField(max_length=40, default=None, null=True) project_group_id = StringField(max_length=40, default=None, null=True)