feat(cost-report): adjust thread pool size (#4804)

Signed-off-by: Di Wang <[email protected]>
app-sre · Jan 7, 2025 · c6b6d8b · c6b6d8b
1 parent 2dec96a
commit c6b6d8b
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 15 deletions.
diff --git a/tools/cli_commands/cost_report/aws.py b/tools/cli_commands/cost_report/aws.py
@@ -125,13 +125,13 @@ def _build_report(
         )
 
     @classmethod
-    def create(cls) -> Self:
+    def create(cls, thread_pool_size: int = THREAD_POOL_SIZE) -> Self:
         gql_api = gql.get_api()
         secret = fetch_cost_report_secret(gql_api)
         cost_management_api = CostManagementApi.create_from_secret(secret)
         return cls(
             gql_api=gql_api,
             cost_management_api=cost_management_api,
             cost_management_console_base_url=secret["console_base_url"],
-            thread_pool_size=THREAD_POOL_SIZE,
+            thread_pool_size=thread_pool_size,
         )
diff --git a/tools/cli_commands/cost_report/openshift.py b/tools/cli_commands/cost_report/openshift.py
@@ -155,12 +155,12 @@ def _build_report(
         )
 
     @classmethod
-    def create(cls) -> Self:
+    def create(cls, thread_pool_size: int = THREAD_POOL_SIZE) -> Self:
         gql_api = gql.get_api()
         secret = fetch_cost_report_secret(gql_api)
         cost_management_api = CostManagementApi.create_from_secret(secret)
         return cls(
             gql_api=gql_api,
             cost_management_api=cost_management_api,
-            thread_pool_size=THREAD_POOL_SIZE,
+            thread_pool_size=thread_pool_size,
         )
diff --git a/tools/cli_commands/cost_report/openshift_cost_optimization.py b/tools/cli_commands/cost_report/openshift_cost_optimization.py
@@ -176,12 +176,12 @@ def _build_resource_config(response: ResourceConfigResponse) -> str | None:
         return f"{round(response.amount)}{response.format}"
 
     @classmethod
-    def create(cls) -> Self:
+    def create(cls, thread_pool_size: int = THREAD_POOL_SIZE) -> Self:
         gql_api = gql.get_api()
         secret = fetch_cost_report_secret(gql_api)
         cost_management_api = CostManagementApi.create_from_secret(secret)
         return cls(
             gql_api=gql_api,
             cost_management_api=cost_management_api,
-            thread_pool_size=THREAD_POOL_SIZE,
+            thread_pool_size=thread_pool_size,
         )
diff --git a/tools/qontract_cli.py b/tools/qontract_cli.py
@@ -2771,22 +2771,27 @@ def sort_by_severity(item: dict[str, str]) -> int:
 
 @get.command()
 @click.pass_context
-def aws_cost_report(ctx):
-    command = AwsCostReportCommand.create()
+@thread_pool_size(default=5)
+def aws_cost_report(ctx, thread_pool_size):
+    command = AwsCostReportCommand.create(thread_pool_size=thread_pool_size)
     print(command.execute())
 
 
 @get.command()
 @click.pass_context
-def openshift_cost_report(ctx):
-    command = OpenShiftCostReportCommand.create()
+@thread_pool_size(default=5)
+def openshift_cost_report(ctx, thread_pool_size):
+    command = OpenShiftCostReportCommand.create(thread_pool_size=thread_pool_size)
     print(command.execute())
 
 
 @get.command()
 @click.pass_context
-def openshift_cost_optimization_report(ctx):
-    command = OpenShiftCostOptimizationReportCommand.create()
+@thread_pool_size(default=5)
+def openshift_cost_optimization_report(ctx, thread_pool_size):
+    command = OpenShiftCostOptimizationReportCommand.create(
+        thread_pool_size=thread_pool_size
+    )
     print(command.execute())
 
 

diff --git a/tools/test/test_qontract_cli.py b/tools/test/test_qontract_cli.py
@@ -145,7 +145,7 @@ def test_get_aws_cost_report(env_vars, mock_queries, mock_aws_cost_report_comman
 
     assert result.exit_code == 0
     assert result.output == "some report\n"
-    mock_aws_cost_report_command.create.assert_called_once_with()
+    mock_aws_cost_report_command.create.assert_called_once_with(thread_pool_size=5)
     mock_aws_cost_report_command.create.return_value.execute.assert_called_once_with()
 
 
@@ -169,7 +169,9 @@ def test_get_openshift_cost_report(
 
     assert result.exit_code == 0
     assert result.output == "some report\n"
-    mock_openshift_cost_report_command.create.assert_called_once_with()
+    mock_openshift_cost_report_command.create.assert_called_once_with(
+        thread_pool_size=5
+    )
     mock_openshift_cost_report_command.create.return_value.execute.assert_called_once_with()
 
 
@@ -193,5 +195,7 @@ def test_get_openshift_cost_optimization_report(
 
     assert result.exit_code == 0
     assert result.output == "some report\n"
-    mock_openshift_cost_optimization_report_command.create.assert_called_once_with()
+    mock_openshift_cost_optimization_report_command.create.assert_called_once_with(
+        thread_pool_size=5
+    )
     mock_openshift_cost_optimization_report_command.create.return_value.execute.assert_called_once_with()