put logic to common function

oamg · Oct 30, 2024 · 8abbc1b · 8abbc1b
1 parent 470367d
commit 8abbc1b
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 32 deletions.
diff --git a/commands/command_utils.py b/commands/command_utils.py
@@ -1,6 +1,7 @@
 import json
 import os
 import re
+import resource
 
 from leapp.exceptions import CommandError
 from leapp.utils import path
@@ -140,3 +141,27 @@ def vet_upgrade_path(args):
                     flavor=flavor,
                     choices=','.join(supported_target_versions)))
     return (target_release, flavor)
+
+
+def set_resource_limits():
+    """
+    Set resource limits for the maximum number of open file descriptors and the maximum writable file size.
+
+    :raises: `CommandError` if the resource limits cannot be set
+    """
+
+    soft_nofile, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
+    soft_fsize, _ = resource.getrlimit(resource.RLIMIT_FSIZE)
+    nofile_limit = 1024*16
+
+    if soft_nofile < nofile_limit:
+        try:
+            resource.setrlimit(resource.RLIMIT_NOFILE, (nofile_limit, nofile_limit))
+        except OSError as err:
+            raise CommandError('Failed to set limit for maximum number of open file descriptors: {}'.format(err))
+
+    if soft_fsize != resource.RLIM_INFINITY:
+        try:
+            resource.setrlimit(resource.RLIMIT_FSIZE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
+        except OSError as err:
+            raise CommandError('Failed to set limit for maximum writeable file size: {}'.format(err))
diff --git a/commands/preupgrade/__init__.py b/commands/preupgrade/__init__.py
@@ -1,5 +1,4 @@
 import os
-import resource
 import sys
 import uuid
 
@@ -60,21 +59,7 @@ def preupgrade(args, breadcrumbs):
     except LeappError as exc:
         raise CommandError(exc.message)
 
-    soft_nofile, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
-    soft_fsize, _ = resource.getrlimit(resource.RLIMIT_FSIZE)
-    nofile_limit = 1024*16
-
-    if soft_nofile < nofile_limit:
-        try:
-            resource.setrlimit(resource.RLIMIT_NOFILE, (nofile_limit, nofile_limit))
-        except OSError as err:
-            raise CommandError('Failed to set limit for maximum number of open file descriptors: {}'.format(err))
-
-    if soft_fsize != resource.RLIM_INFINITY:
-        try:
-            resource.setrlimit(resource.RLIMIT_FSIZE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
-        except OSError as err:
-            raise CommandError('Failed to set limit for maximum writeable file size: {}'.format(err))
+    command_utils.set_resource_limits()
 
     workflow = repositories.lookup_workflow('IPUWorkflow')()
     util.warn_if_unsupported(configuration)

diff --git a/commands/upgrade/__init__.py b/commands/upgrade/__init__.py
@@ -1,5 +1,4 @@
 import os
-import resource
 import sys
 import uuid
 
@@ -91,21 +90,7 @@ def upgrade(args, breadcrumbs):
     except LeappError as exc:
         raise CommandError(exc.message)
 
-    soft_nofile, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
-    soft_fsize, _ = resource.getrlimit(resource.RLIMIT_FSIZE)
-    nofile_limit = 1024*16
-
-    if soft_nofile < nofile_limit:
-        try:
-            resource.setrlimit(resource.RLIMIT_NOFILE, (nofile_limit, nofile_limit))
-        except OSError as err:
-            raise CommandError('Failed to set limit for maximum number of open file descriptors: {}'.format(err))
-
-    if soft_fsize != resource.RLIM_INFINITY:
-        try:
-            resource.setrlimit(resource.RLIMIT_FSIZE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
-        except OSError as err:
-            raise CommandError('Failed to set limit for maximum writeable file size: {}'.format(err))
+    command_utils.set_resource_limits()
 
     workflow = repositories.lookup_workflow('IPUWorkflow')(auto_reboot=args.reboot)
     util.process_whitelist_experimental(repositories, workflow, configuration, logger)