Update README.md and minor fixes (#9)

* docs: Update README.md
* fix: Remove deprecated CLI options
* chore: Minor fixes

Author: baevmo

README.md

@@ -3,5 +3,4 @@
 from pg_anon import MainRoutine
 
 if __name__ == "__main__":
-    loop = asyncio.get_event_loop()
-    loop.run_until_complete(MainRoutine().run())
+    asyncio.run(MainRoutine().run())

pg_anon.py

@@ -4,8 +4,7 @@
 
 
 def _run_pg_anon():
-    loop = asyncio.get_event_loop()
-    loop.run_until_complete(MainRoutine().run())
+    asyncio.run(MainRoutine().run())
 
 
 if __name__ == "__main__":

pg_anon/__main__.py

@@ -5,22 +5,12 @@
 import subprocess
 import sys
 import traceback
-from enum import Enum
+from enum import StrEnum
 
 from pkg_resources import parse_version as version
 
 
-class BasicEnum:
-    def __str__(self):
-        return self.value
-
-
-class OutputFormat(BasicEnum, Enum):
-    BINARY = "binary"
-    TEXT = "text"
-
-
-class ResultCode(BasicEnum, Enum):
+class ResultCode(StrEnum):
     DONE = "done"
     FAIL = "fail"
     UNKNOWN = "unknown"
@@ -32,13 +22,13 @@ class PgAnonResult:
     result_data = None
 
 
-class VerboseOptions(BasicEnum, Enum):
+class VerboseOptions(StrEnum):
     INFO = "info"
     DEBUG = "debug"
     ERROR = "error"
 
 
-class AnonMode(BasicEnum, Enum):
+class AnonMode(StrEnum):
     DUMP = "dump"  # dump table contents to files using dictionary
     RESTORE = "restore"  # create tables in target database and load data from files
     INIT = "init"  # create a schema with anonymization helper functions
@@ -51,7 +41,7 @@ class AnonMode(BasicEnum, Enum):
     CREATE_DICT = "create-dict"  # create dictionary
 
 
-class ScanMode(BasicEnum, Enum):
+class ScanMode(StrEnum):
     FULL = "full"
     PARTIAL = "partial"
 

pg_anon/common.py

@@ -4,7 +4,6 @@
 from pg_anon.common import (
     exception_handler,
     AnonMode,
-    OutputFormat,
     VerboseOptions,
     ScanMode,
 )
@@ -86,20 +85,6 @@ def get_arg_parser():
         parser.add_argument(
             "--mode", type=AnonMode, choices=list(AnonMode), default=AnonMode.INIT
         )
-        parser.add_argument(
-            "--copy-options",
-            dest="copy_options",
-            default="",
-            help='Options for COPY command like "with binary".',
-        )
-        parser.add_argument(
-            "--format",
-            dest="format",
-            type=OutputFormat,
-            choices=list(OutputFormat),
-            default=OutputFormat.BINARY.value,
-            help="COPY data format, can be overwritten by --copy-options. Selects the data format to be read or written: text, csv or binary.",
-        )
         parser.add_argument(
             "--verbose",
             dest="verbose",
@@ -109,9 +94,30 @@ def get_arg_parser():
             help="Enable verbose output",
         )
         parser.add_argument("--dict-file", type=str, default="")
-        parser.add_argument("--threads", type=int, default=4)
-        parser.add_argument("--pg-dump", type=str, default="/usr/bin/pg_dump")
-        parser.add_argument("--pg-restore", type=str, default="/usr/bin/pg_restore")
+        parser.add_argument(
+            "--threads",
+            type=int,
+            default=4,
+            help="Amount of threads for IO operations.",
+        )
+        parser.add_argument(
+            "--processes",
+            type=int,
+            default=4,
+            help="Amount of processes for multiprocessing operations.",
+        )
+        parser.add_argument(
+            "--pg-dump",
+            type=str,
+            default="/usr/bin/pg_dump",
+            help="Path to the `pg_dump` Postgres tool",
+        )
+        parser.add_argument(
+            "--pg-restore",
+            type=str,
+            default="/usr/bin/pg_restore",
+            help="Path to the `pg_dump` Postgres tool.",
+        )
         parser.add_argument("--output-dir", type=str, default="")
         parser.add_argument("--input-dir", type=str, default="")
         parser.add_argument(
@@ -128,26 +134,24 @@ def get_arg_parser():
         )
         parser.add_argument("--clear-output-dir", action="store_true", default=False)
         parser.add_argument(
-            "--drop-custom-check-constr", action="store_true", default=False
-        )
-        parser.add_argument(
-            "--seq-init-by-max-value",
+            "--drop-custom-check-constr",
             action="store_true",
             default=False,
-            help="""Initialize sequences based on maximum values. Otherwise, the sequences will be initialized
-                based on the values of the source database.""",
+            help="Drop all CHECK constrains containing user-defined procedures to avoid performance "
+            "degradation at the data loading stage.",
         )
         parser.add_argument(
-            "--disable-checks",
+            "--seq-init-by-max-value",
             action="store_true",
             default=False,
-            help="""Disable checks of disk space and PostgreSQL version""",
+            help="Initialize sequences based on maximum values. Otherwise, the sequences "
+            "will be initialized based on the values of the source database.",
         )
         parser.add_argument(
-            "--skip-data",
+            "--disable-checks",
             action="store_true",
             default=False,
-            help="""Don't copy data. Only the database structure will be dumped""",
+            help="Disable checks of disk space and PostgreSQL version.",
         )
         parser.add_argument(
             "--scan-mode",

pg_anon/context.py

@@ -475,7 +475,7 @@ async def init_process(name, ctx, fields_info_chunk: List[FieldInfo]):
 
     p = aioprocessing.AioProcess(
         target=process_impl,
-        args=(name, ctx, queue, fields_info_chunk, ctx.conn_params, ctx.args.threads),
+        args=(name, ctx, queue, fields_info_chunk, ctx.conn_params, ctx.args.processes),
     )
     p.start()
     res = None

pg_anon/create_dict.py

@@ -115,6 +115,8 @@ def setup_logger(self):
             self.logger.setLevel(log_level)
 
     def close_logger_handlers(self):
+        if not self.logger:  # FIXME: Return an exception for --help command
+            return
         for handler in self.logger.handlers[
             :
         ]:  # iterate over a copy of the handlers list

pg_anon/pg_anon.py

@@ -245,7 +245,7 @@ async def make_restore(ctx):
 
     if not db_is_empty and ctx.args.mode != AnonMode.SYNC_DATA_RESTORE:
         db_conn.close()
-        raise Exception("Target DB is not empty!")
+        raise Exception(f"Target DB {ctx.conn_params['database']} is not empty!")
 
     metadata_file = open(
         os.path.join(ctx.current_dir, "dict", ctx.args.input_dir, "metadata.json"), "r"