[TTDB-394] Added multiprocessing for dump mode (#25)

* [TTDB-394] Added multiprocessing for dump mode
- Added `server_settings` for connections
- Added docs for `db_utils`
- Renamed CLI parameter `threads` to `db_connections_per_process`
- Updated tests run arguments

* [TTDB-394] Updated parameter `db_connections_per_process in `README.md`

* [TTDB-394] Added check exists schema with utils

* [TTDB-394] Fixed check exists schema with utils

* [TTDB-394] Fixed json output for "view" modes

* [TTDB-394] Fixed escaping quotes

---------

Co-authored-by: Max Ibragimov <maxim.ibragimov@tantorlabs.ru>

Author: max-ibragimow

README.md

@@ -347,12 +347,12 @@ python pg_anon.py --help
 
 Common pg_anon options:
 
-| Option        | Description                                                    |
-|---------------|----------------------------------------------------------------|
-| `--debug`     | Enable debug mode (default false)                              |
-| `--verbose`   | Configure verbose mode: [info, debug, error] (default info)    |
-| `--threads`   | Amount of threads for IO operations (default 4)                |
-| `--processes` | Amount of processes for multiprocessing operations (default 4) |
+| Option                          | Description                                                          |
+|---------------------------------|----------------------------------------------------------------------|
+| `--debug`                       | Enable debug mode (default false)                                    |
+| `--verbose`                     | Configure verbose mode: [info, debug, error] (default info)          |
+| `--db-connections-per-process`  | Amount of connections for IO operations for each process (default 4) |
+| `--processes`                   | Amount of processes for multiprocessing operations (default 4)       |
 
 Database configuration options:
 

pg_anon/common/constants.py

@@ -0,0 +1 @@
+ANON_UTILS_DB_SCHEMA_NAME = 'anon_funcs'

pg_anon/common/db_queries.py

@@ -1,5 +1,6 @@
 import re
 
+from pg_anon.common.constants import ANON_UTILS_DB_SCHEMA_NAME
 from pg_anon.common.dto import FieldInfo
 
 
@@ -9,15 +10,15 @@ def get_query_limit(limit: int) -> str:
 
 def get_query_get_scan_fields(limit: int = None, count_only: bool = False):
     if not count_only:
-        fields = """
+        fields = f"""
             SELECT DISTINCT
             n.nspname,
             c.relname,
             a.attname AS column_name,
             format_type(a.atttypid, a.atttypmod) as type,
             c.oid, a.attnum,
-            anon_funcs.digest(n.nspname || '.' || c.relname || '.' || a.attname, '', 'md5') as obj_id,
-            anon_funcs.digest(n.nspname || '.' || c.relname, '', 'md5') as tbl_id
+            {ANON_UTILS_DB_SCHEMA_NAME}.digest(n.nspname || '.' || c.relname || '.' || a.attname, '', 'md5') as obj_id,
+            {ANON_UTILS_DB_SCHEMA_NAME}.digest(n.nspname || '.' || c.relname, '', 'md5') as tbl_id
         """
         order_by = 'ORDER BY 1, 2, a.attnum' if count_only else ''
     else:

pg_anon/common/db_utils.py

@@ -3,28 +3,69 @@
 import asyncpg
 from asyncpg import Connection
 
+from pg_anon.common.constants import ANON_UTILS_DB_SCHEMA_NAME
 from pg_anon.common.db_queries import get_query_get_scan_fields
 from pg_anon.common.dto import FieldInfo
 
 
-async def get_scan_fields_list(connection_params: Dict, limit: int = None) -> List:
-    db_conn = await asyncpg.connect(**connection_params)
+async def check_anon_utils_db_schema_exists(connection_params: Dict, server_settings: Dict = None) -> bool:
+    """
+    Checks exists db schema what consists predefined anonymization utils
+    :param connection_params: Required connection parameters such as host, login, password and etc.
+    :param server_settings: Optional server settings for new connection. Can consists of timeout settings, application name and etc.
+    :return: Exists schema or not
+    """
+    query = f"""
+    select exists (select schema_name FROM information_schema.schemata where "schema_name" = '{ANON_UTILS_DB_SCHEMA_NAME}');
+    """
+
+    db_conn = await asyncpg.connect(**connection_params, server_settings=server_settings)
+    exists = await db_conn.fetchval(query)
+    await db_conn.close()
+    return exists
+
+
+async def get_scan_fields_list(connection_params: Dict, server_settings: Dict = None, limit: int = None) -> List:
+    """
+    Get fields list for scan sensitive data
+    :param connection_params: Required connection parameters such as host, login, password and etc.
+    :param server_settings: Optional server settings for new connection. Can consists of timeout settings, application name and etc.
+    :param limit: Limit the number of results to return.
+    :return: resulted fields list for processing
+    """
     query = get_query_get_scan_fields(limit=limit)
+
+    db_conn = await asyncpg.connect(**connection_params, server_settings=server_settings)
     fields_list = await db_conn.fetch(query)
     await db_conn.close()
     return fields_list
 
 
-async def get_scan_fields_count(connection_params: Dict) -> int:
-    db_conn = await asyncpg.connect(**connection_params)
+async def get_scan_fields_count(connection_params: Dict, server_settings: Dict = None) -> int:
+    """
+    Get count of fields for scan sensitive data
+    :param connection_params: Required connection parameters such as host, login, password and etc.
+    :param server_settings: Optional server settings for new connection. Can consists of timeout settings, application name and etc.
+    :return: count of resulted fields list for processing
+    """
     query = get_query_get_scan_fields(count_only=True)
+
+    db_conn = await asyncpg.connect(**connection_params, server_settings=server_settings)
     count = await db_conn.fetchval(query)
     await db_conn.close()
     return count
 
 
-async def get_fields_list(connection_params: Dict, table_schema: str, table_name: str) -> List:
-    db_conn = await asyncpg.connect(**connection_params)
+async def get_fields_list(connection_params: Dict, table_schema: str, table_name: str, server_settings: Dict = None) -> List:
+    """
+    Get fields list for dump
+    :param connection_params: Required connection parameters such as host, login, password and etc.
+    :param table_schema: Table schema name
+    :param table_name: Table name
+    :param server_settings: Optional server settings for new connection. Can consists of timeout settings, application name and etc.
+    :return: fields list for dump
+    """
+    db_conn = await asyncpg.connect(**connection_params, server_settings=server_settings)
     fields_list = await db_conn.fetch(
         """
             SELECT column_name, udt_name FROM information_schema.columns

pg_anon/common/multiprocessing_utils.py

@@ -0,0 +1,34 @@
+import time
+from typing import List, Callable
+
+import aioprocessing
+
+
+async def init_process(name: str, ctx, target_func: Callable, tasks: List, *args, **kwargs):
+    from pg_anon.context import Context
+
+    ctx: Context
+    start_t = time.time()
+    ctx.logger.info(f"================> Process [{name}] started. Input items: {len(tasks)}")
+    queue = aioprocessing.AioQueue()
+
+    p = aioprocessing.AioProcess(
+        target=target_func,
+        args=(name, ctx, queue, tasks, *args),
+        kwargs=kwargs,
+    )
+    p.start()
+    res = None
+    while True:
+        result = await queue.coro_get()
+        if result is None:
+            break
+        res = result
+    await p.coro_join()
+    end_t = time.time()
+    elapsed = round(end_t - start_t, 2)
+    result_item_log = str(len(res)) if res is not None else "0"
+    ctx.logger.info(
+        f"<================ Process [{name}] finished, elapsed: {elapsed} sec. Result {result_item_log} item(s)"
+    )
+    return res

pg_anon/common/utils.py

@@ -206,7 +206,7 @@ async def get_dump_query(ctx, table_schema: str, table_name: str, table_rule,
         (table_schema + "_" + table_name).encode()
     ).hexdigest()
 
-    files["%s.bin.gz" % hashed_name] = {"schema": table_schema, "table": table_name}
+    files[f"{hashed_name}.bin.gz"] = {"schema": table_schema, "table": table_name}
 
     if not found_white_list:
         included_objs.append(
@@ -217,7 +217,6 @@ async def get_dump_query(ctx, table_schema: str, table_name: str, table_rule,
                 or ctx.args.dbg_stage_2_validate_data
                 or ctx.args.dbg_stage_3_validate_full):
             query = "SELECT * FROM %s %s" % (table_name_full, ctx.validate_limit)
-            ctx.logger.info(str(query))
             return query
         else:
             query = f"SELECT * FROM {table_name_full}"
@@ -242,52 +241,42 @@ async def get_dump_query(ctx, table_schema: str, table_name: str, table_rule,
             # the table is transferred with the specific fields for anonymization
             fields_list = await get_fields_list(
                 connection_params=ctx.conn_params,
+                server_settings=ctx.server_settings,
                 table_schema=table_schema,
                 table_name=table_name
             )
 
             sql_expr = ""
 
-            def check_fld(fld_name):
-                if fld_name in table_rule["fields"]:
-                    return fld_name, table_rule["fields"][fld_name]
+            def check_field(field_name: str):
+                if field_name in table_rule["fields"]:
+                    return field_name, table_rule["fields"][field_name]
                 return None, None
 
             for cnt, column_info in enumerate(fields_list):
                 column_name = column_info["column_name"]
                 udt_name = column_info["udt_name"]
-                fld_name, fld_val = check_fld(column_name)
-                if fld_name:
-                    if fld_val.find("SQL:") == 0:
-                        sql_expr += f'({fld_val[4:]}) as "{fld_name}"'
+                field_name, field_value = check_field(column_name)
+
+                if field_name:
+                    if field_value.find("SQL:") == 0:
+                        sql_expr += f'({field_value[4:]}) as "{field_name}"'
                     else:
-                        sql_expr += f'{fld_val}::{udt_name} as "{fld_name}"'
+                        sql_expr += f'{field_value}::{udt_name} as "{field_name}"'
                 else:
                     # field "as is"
-                    if (
-                            not column_name.islower() and not column_name.isupper()
-                    ) or column_name.isupper():
-                        sql_expr += f'"{column_name}" as "{column_name}"'
-                    else:
-                        sql_expr += f'"{column_name}" as "{column_name}"'
+                    sql_expr += f'"{column_name}" as "{column_name}"'
+
                 if cnt != len(fields_list) - 1:
                     sql_expr += ",\n"
 
+            query = f"SELECT {sql_expr} FROM {table_name_full}"
             if (ctx.args.dbg_stage_1_validate_dict
                     or ctx.args.dbg_stage_2_validate_data
                     or ctx.args.dbg_stage_3_validate_full):
-                query = "SELECT %s FROM %s %s" % (
-                    sql_expr,
-                    table_name_full,
-                    ctx.validate_limit,
-                )
-                return query
-            else:
-                query = "SELECT %s FROM %s" % (
-                    sql_expr,
-                    table_name_full,
-                )
-                return query
+                query += f" {ctx.validate_limit}"
+
+            return query
 
 
 def get_file_name_from_path(path: str) -> str:

pg_anon/context.py

@@ -2,6 +2,7 @@
 import os
 from typing import Dict, Optional
 
+from pg_anon.common.constants import ANON_UTILS_DB_SCHEMA_NAME
 from pg_anon.common.enums import VerboseOptions, AnonMode, ScanMode
 from pg_anon.common.utils import (
     exception_handler,
@@ -24,12 +25,20 @@ def __init__(self, args):
         self.total_rows = 0
         self.create_dict_sens_matches = {}  # for create-dict mode
         self.create_dict_no_sens_matches = {}  # for create-dict mode
-        self.exclude_schemas = ["anon_funcs", "columnar_internal"]
+        self.exclude_schemas = [ANON_UTILS_DB_SCHEMA_NAME, "columnar_internal"]
         self.logger = None
 
         if args.db_user_password == "" and os.environ.get("PGPASSWORD") is not None:
             args.db_user_password = os.environ["PGPASSWORD"]
 
+        self.server_settings = {
+            "application_name": "pg_anon",
+            "statement_timeout": "0",
+            "lock_timeout": "0",
+            "idle_in_transaction_session_timeout": "0",
+            "idle_session_timeout": "0",
+        }
+
         self.conn_params = {
             "host": args.db_host,
             "database": args.db_name,
@@ -208,8 +217,8 @@ def read_prepared_dict(self, save_dict_file_name_for_each_rule: bool = False):
                         dictionary_rule['dict_file_name'] = dict_file
                 self.prepared_dictionary_obj["dictionary"].extend(dictionary_rules)
 
-            if dictionary := dict_data.get("dictionary_exclude", []):
-                self.prepared_dictionary_obj["dictionary_exclude"].extend(dictionary)
+            if dictionary_exclude_rules := dict_data.get("dictionary_exclude", []):
+                self.prepared_dictionary_obj["dictionary_exclude"].extend(dictionary_exclude_rules)
 
             if validate_tables := dict_data.get("validate_tables", []):
                 self.prepared_dictionary_obj["validate_tables"].extend(validate_tables)
@@ -260,10 +269,10 @@ def get_arg_parser():
             help="In 'create-dict' mode input file or file list with scan rules of sensitive and not sensitive fields"
         )
         parser.add_argument(
-            "--threads",
+            "--db-connections-per-process",
             type=int,
             default=4,
-            help="Amount of threads for IO operations.",
+            help="Amount of db connections for each process for IO operations.",
         )
         parser.add_argument(
             "--processes",