Merge pull request #91 from whdalsrnt/master

refactor: add provider_filter to BudgetUsage Model
cloudforet-io · Sep 12, 2023 · 475516a · 475516a
2 parents 7929d82 + 771144b
commit 475516a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git 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
@@ -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
@@ -3,13 +3,22 @@
 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='')
     date = StringField(max_length=7, required=True)
     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)