Adding option for limiting the number of entries to migrate in simulation mode.

diff --git a/lib/ldap_migration/__init__.py b/lib/ldap_migration/__init__.py
index 702f4b4d587095ee0b6f5a66c1346c65950aaab2..9478511500f69b917b6a5cb124903a2b6bc36c18 100644 (file)
--- a/lib/ldap_migration/__init__.py
+++ b/lib/ldap_migration/__init__.py
@@ -15,6 +15,7 @@ import sys
 import os
 import time
 import re
+import argparse
 
 # 3rd party modules
 
@@ -35,7 +36,7 @@ from fb_tools.errors import FbAppError
 from .config import LDAPMigrationConfiguration
 from .idict import CaseInsensitiveDict
 
-__version__ = '0.5.4'
+__version__ = '0.6.0'
 
 LOG = logging.getLogger(__name__)
 CFG_BASENAME = 'ldap-migration.ini'
@@ -52,6 +53,57 @@ class CommonLDAPMigrationError(FbAppError, LDAPException):
     pass
 
 
+# =============================================================================
+class NonNegativeItegerOptionAction(argparse.Action):
+
+    # -------------------------------------------------------------------------
+    def __call__(self, parser, namespace, value, option_string=None):
+
+        if value < 0:
+            msg = "The option must not be negative (given: {}).".format(value)
+            raise argparse.ArgumentError(self, msg)
+
+        setattr(namespace, self.dest, value)
+
+
+# =============================================================================
+class LimitedItegerOptionAction(argparse.Action):
+
+    # -------------------------------------------------------------------------
+    def __init__(self, option_strings, min_val=None, max_val=None, *args, **kwargs):
+
+        self._min_val = min_val
+        self._max_val = max_val
+
+        super(LimitedItegerOptionAction, self).__init__(
+            option_strings=option_strings, *args, **kwargs)
+
+    # -------------------------------------------------------------------------
+    def __call__(self, parser, namespace, value, option_string=None):
+
+        val = 0
+        try:
+            val = int(value)
+        except Exception as e:
+            msg = "Got a {c} for converting {v!r} into an integer value: {e}".format(
+                    c=e.__class__.__name__, v=value, e=e)
+            raise argparse.ArgumentError(self, msg)
+
+        if self._min_val is not None:
+            if val < self._min_val:
+                msg = "The option must be greater or equal to {m} (given: {v}).".format(
+                    m=self._min_val, v=val)
+                raise argparse.ArgumentError(self, msg)
+
+        if self._max_val is not None:
+            if val > self._max_val:
+                msg = "The option must be less or equal to {m} (given: {v}).".format(
+                    m=self._max_val, v=val)
+                raise argparse.ArgumentError(self, msg)
+
+        setattr(namespace, self.dest, val)
+
+
 # =============================================================================
 class LDAPMigrationApplication(BaseApplication):
     """
@@ -81,6 +133,8 @@ class LDAPMigrationApplication(BaseApplication):
         self.all_dns_file = None
         self.structural_dns_file = None
 
+        self.limit = 0
+
         self.object_classes = CaseInsensitiveDict()
         self.attribute_types = CaseInsensitiveDict()
         self.dns = CaseInsensitiveDict()
@@ -144,10 +198,20 @@ class LDAPMigrationApplication(BaseApplication):
 
         app_group.add_argument(
             '-T', '--timeout', dest='timeout', type=int, metavar='SECONDS',
+            action=LimitedItegerOptionAction,
+            min_val=1, max_val=LDAPMigrationConfiguration.max_timeout,
             help="The timeout in seconds for LDAP operations (default: {}).".format(
                 LDAPMigrationConfiguration.default_timeout),
         )
 
+        app_group.add_argument(
+            '-L', '--limit', dest='limit', type=int, metavar='NUMBER',
+            action=NonNegativeItegerOptionAction,
+            help=(
+                "Limiting the migration to the first NUMBER entries. "
+                "This option is valid only in simulation mode."),
+        )
+
         app_group.add_argument(
             '-c', '--config', '--config-file', dest='cfg_file', metavar='FILE',
             action=CfgFileOptionAction,
@@ -260,6 +324,17 @@ class LDAPMigrationApplication(BaseApplication):
                 self.arg_parser.print_usage(sys.stdout)
                 self.exit(1)
 
+        if self.args.limit:
+            if self.simulate:
+                self.limit = self.args.limit
+            else:
+                LOG.warn(
+                    "Limiting the number of entries to migrate is valid only in "
+                    "siulation mode.")
+                print()
+                self.arg_parser.print_usage(sys.stdout)
+                self.exit(1)
+
         self.initialized = True
 
     # -------------------------------------------------------------------------
@@ -619,8 +694,9 @@ class LDAPMigrationApplication(BaseApplication):
                     LOG.debug("Found DN {!r}.".format(old_dn))
                 self.register_dn_tokens(old_dn, entry['attributes']['objectClass'], self.dns)
                 fh.write("{old} => {new}\n".format(old=old_dn, new=new_dn))
-#                if item_count >= 100:
-#                    break
+                if self.limit:
+                    if item_count >= self.limit:
+                        break
 
         LOG.info("Found {nr} items in subtree of {sfx!r}.".format(
             nr=item_count, sfx=self.config.suffix))

diff --git a/lib/ldap_migration/config.py b/lib/ldap_migration/config.py
index 6f6c1fe8dedf2f8880d1a9724f9277b0b5155021..fbaf235b39a7b58f3b58a40ef0180a7cf1e1e639 100644 (file)
--- a/lib/ldap_migration/config.py
+++ b/lib/ldap_migration/config.py
@@ -21,7 +21,7 @@ from fb_tools.common import to_bool
 
 from fb_tools.config import ConfigError, BaseConfiguration
 
-__version__ = '0.2.1'
+__version__ = '0.2.2'
 
 LOG = logging.getLogger(__name__)
 
@@ -55,6 +55,7 @@ class LDAPMigrationConfiguration(BaseConfiguration):
     default_timeout = 20
     max_tcp_port = (2**16 - 1)
     default_xlations_definitions_base = 'xlations.yaml'
+    max_timeout = 3600
 
     re_bind_dn = re.compile(r'^\s*bind[_-]?dn\s*$', re.IGNORECASE)
     re_bind_pw = re.compile(r'^\s*bind[_-]?(?:pw|passwd|passsword)\s*$', re.IGNORECASE)
@@ -146,7 +147,7 @@ class LDAPMigrationConfiguration(BaseConfiguration):
             return
         val = int(value)
         err_msg = "Invalid timeout {!r} for LDAP operations, must be 0 < SECONDS < 3600."
-        if val <= 0 or val > 3600:
+        if val <= 0 or val > self.max_timeout:
             msg = err_msg.format(value)
             raise ValueError(msg)
         self._timeout = val