Repository: froOzzy/pg_stage
Pull request #64: Добавлена возможность обфускации uuid на основе значения другой колонки
Author: login: Makcxim
Created: 2025-03-05T10:39:23Z | Merged: 2025-03-09T14:14:44Z

Проблема:
два несвязанных сервера с базами  
в обоих базах лежит колонка, которая после обфускации должна быть одинаковой на обоих серверах  

те нужно из одних и тех же данных генерировать данные для другой колонки  

Решение:
генерировать uuid на основе другой колонки, чтобы на двух серверах uuid был одинаковым при одинаковом значение исходной колонки  

Что сделал:  
для комментов вида:
```SQL
COMMENT ON COLUMN uuid IS 'anon: [{"mutation_name": "uuid5_by_source_value", "mutation_kwargs": {"source_column": "username", "namespace": "73556080-0000-0000-0000-735560800000"}}]';
```
будет сгенерирован uuid на основе значения колонки `username` текущей строки + переданного `uuid namespace` и текущей даты (чтобы изо дня в день значения менялись) (одно из условий)  





Diff:
diff --git a/src/pg_stage/mutator.py b/src/pg_stage/mutator.py
index 8657238..9b67d7c 100644
--- a/src/pg_stage/mutator.py
+++ b/src/pg_stage/mutator.py
@@ -1,9 +1,9 @@
-import random
 import datetime
+import random
 import uuid
 from typing import Any, Callable, List
 
-from mimesis import Person, Address, Datetime, Internet, Numbers
+from mimesis import Address, Datetime, Internet, Numbers, Person
 from mimesis.builtins import RussiaSpecProvider
 
 
@@ -516,3 +516,34 @@ def mutation_uuid4(**kwargs: Any) -> str:
         :return: строка uuid4
         """
         return str(uuid.uuid4())
+
+    def mutation_uuid5_by_source_value(self, **kwargs: Any) -> str:
+        """
+        Метод для формирования uuid5 с использованием значения из указанной колонки,
+        текущей даты и переданного uuid namespace.
+
+        :param kwargs:
+            source_column: Название колонки, значение которой хотим использовать.
+            namespace: Uuid namespace.
+        :return: Строка uuid5.
+        """
+        source_column: str = kwargs.get('source_column')
+        if not source_column:
+            raise ValueError('Argument "source_column" not found')
+
+        obfuscated_values: dict[str, Any] = kwargs.get('obfuscated_values', {})
+        source_value: str = obfuscated_values.get(source_column)
+        if not source_value:
+            raise ValueError('Value of "source_column" not found')
+
+        uuid_namespace: str = kwargs.get('namespace')
+        if uuid_namespace is None:
+            raise ValueError('Argument "namespace" not found')
+
+        try:
+            uuid_namespace: uuid.UUID = uuid.UUID(uuid_namespace)
+        except Exception as e:
+            raise ValueError('Invalid uuid namespace given') from e
+
+        date_today: datetime.date = self._now.date()
+        return str(uuid.uuid5(uuid_namespace, f'{source_value}-{date_today}'))
diff --git a/src/pg_stage/obfuscator.py b/src/pg_stage/obfuscator.py
index 4b19227..b2fb252 100644
--- a/src/pg_stage/obfuscator.py
+++ b/src/pg_stage/obfuscator.py
@@ -1,8 +1,8 @@
+import json
 import re
 import sys
-import json
 from collections import defaultdict
-from typing import Optional, List, Set, Dict
+from typing import Dict, List, Optional, Set
 from uuid import uuid4
 
 from pg_stage.mutator import Mutator
@@ -152,10 +152,25 @@ def _parse_comment_table(self, *, line: str) -> str:
 
         return line
 
+    def _sort_columns_by_source_column_exists(self, table_mutations_by_column: dict) -> list:
+        """
+        Метод для сортировки столбцов на основе наличия параметра `source_column` в аргументах мутации.
+        :return: Список с верным порядком прохождения столбцов
+        """
+        independent_columns = []
+        dependent_columns = []
+        for column_name in self._enumerate_table_columns:
+            mutations_for_column = table_mutations_by_column.get(column_name)
+            if mutations_for_column and 'source_column' in mutations_for_column[0].get('mutation_kwargs', {}):
+                dependent_columns.append(column_name)
+                continue
+            independent_columns.append(column_name)
+        return independent_columns + dependent_columns
+
     def _prepared_data(self, *, line: str) -> Optional[str]:
         """
         Метод для обработки данных.
-        :return: новая строка с данными
+        :return: Новая строка с данными
         """
         if self._is_delete:
             return None
@@ -164,12 +179,15 @@ def _prepared_data(self, *, line: str) -> Optional[str]:
         if not table_mutations_by_column:
             return line
 
-        result = []
+        sorted_columns = self._sort_columns_by_source_column_exists(table_mutations_by_column)
+
+        obfuscated_values = {}
         table_values = line.split(self.delimiter)
-        for column_name, column_index in self._enumerate_table_columns.items():
+        for column_name in sorted_columns:
+            column_index = self._enumerate_table_columns[column_name]
             mutations_for_column = table_mutations_by_column.get(column_name)
             if not mutations_for_column:
-                result.append(table_values[column_index])
+                obfuscated_values[column_name] = table_values[column_index]
                 continue
 
             is_obfuscated = False
@@ -184,14 +202,17 @@ def _prepared_data(self, *, line: str) -> Optional[str]:
                 mutation_conditions = mutation_for_column['mutation_conditions']
                 if not self._checking_conditions(conditions=mutation_conditions, table_values=table_values):
                     if mutation_index + 1 == len_mutations_for_column:
-                        result.append(table_values[column_index])
+                        obfuscated_values[column_name] = table_values[column_index]
                         break
 
                     continue
 
+                if 'source_column' in mutation_kwargs:
+                    mutation_kwargs['obfuscated_values'] = obfuscated_values
+
                 if not mutation_relations:
                     is_obfuscated = True
-                    result.append(mutation_func(**mutation_kwargs))
+                    obfuscated_values[column_name] = mutation_func(**mutation_kwargs)
                     continue
 
                 new_value = None
@@ -222,10 +243,14 @@ def _prepared_data(self, *, line: str) -> Optional[str]:
 
                     self._relation_values[relation_fk] = new_value
 
-                result.append(new_value)
+                obfuscated_values[column_name] = new_value
                 is_obfuscated = True
 
-        return self.delimiter.join(result)
+        sorted_result: list[Optional[str]] = [None] * len(self._enumerate_table_columns)
+        for column_name, column_index in self._enumerate_table_columns.items():
+            sorted_result[column_index] = obfuscated_values[column_name]
+
+        return self.delimiter.join(sorted_result)
 
     def _parse_copy_values(self, *, line: str) -> Optional[str]:
         """
diff --git a/tests/sql/test_obfuscate_uuid_by_source_value.sql b/tests/sql/test_obfuscate_uuid_by_source_value.sql
new file mode 100644
index 0000000..58884e4
--- /dev/null
+++ b/tests/sql/test_obfuscate_uuid_by_source_value.sql
@@ -0,0 +1,7 @@
+COMMENT ON COLUMN table_1.username IS 'anon: [{"mutation_name": "phone_number", "mutation_kwargs": {"mask": "791########", "unique": true}, "conditions": [{"column_name": "username", "operation": "by_pattern", "value": "^7[0-9]{10}$"}]}]';
+COMMENT ON COLUMN table_1.uuid IS 'anon: [{"mutation_name": "uuid5_by_source_value", "mutation_kwargs": {"source_column": "username", "namespace": "6ba7b810-9dad-11d1-80b4-00c04fd430c8"}}]';
+COPY table_1 (uuid, username, last_name) FROM stdin;
+20f654fe-b27d-4051-9fd4-000000000001	79999999999	111n
+20f654fe-b27d-4051-9fd4-000000000002	79888888888	222n
+20f654fe-b27d-4051-9fd4-000000000003	79777777777	333n
+\.
\ No newline at end of file
diff --git a/tests/test_prepared_data.py b/tests/test_prepared_data.py
index 74bf61a..2928238 100644
--- a/tests/test_prepared_data.py
+++ b/tests/test_prepared_data.py
@@ -1,3 +1,6 @@
+import uuid
+from datetime import date
+
 from src.pg_stage.obfuscator import Obfuscator
 
 
@@ -151,3 +154,26 @@ def test_two_mutation_for_one_column(obfuscator_object: Obfuscator):
     assert prepared_result[0][0] != '79999999999'  # nosec
     assert prepared_result[1][0] != 'test@mail.ru'  # nosec
     assert prepared_result[2][0] != '89999999999'  # nosec
+
+
+def test_obfuscate_uuid_by_source_column(obfuscator_object: Obfuscator):
+    """
+    Arrange: Дамп таблицы, в которой определены мутации uuid на основе мутации другой колонки
+    Act: Вызов функции `_parse_line` класса Obfuscator
+    Assert: В stdout данные поля мутированы на основе мутированного значения другой колонки
+    """
+    with open('tests/sql/test_obfuscate_uuid_by_source_value.sql') as file:
+        dump_sql = file.read()
+
+    result = []
+    for line in dump_sql.splitlines():
+        new_line = obfuscator_object._parse_line(line=line)
+        if new_line is not None:
+            result.append(new_line)
+
+    prepared_result = [line.split('\t') for line in result if '111n' in line or '222n' in line or '333n' in line]
+
+    date_today = date.today()
+    uuid_namespace = uuid.UUID(str('6ba7b810-9dad-11d1-80b4-00c04fd430c8'))
+    for (column_uuid, phone, last_name) in prepared_result:
+        assert column_uuid == str(uuid.uuid5(uuid_namespace, f'{phone}-{date_today}'))  # nosec


Reviews:
user: login: ilyichumakov
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T11:01:12Z
user: login: ilyichumakov
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-07T09:06:14Z
user: login: froOzzy
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-07T17:44:31Z
user: login: froOzzy
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-09T05:42:10Z

Review comments:
diff_hunk: @@ -189,6 +189,13 @@ def _prepared_data(self, *, line: str) -> Optional[str]:
 
                     continue
 
+                if 'source_column' in mutation_kwargs:
+                    source_column = mutation_kwargs['source_column']
+                    source_index = self._enumerate_table_columns.get(source_column)
+                    if source_index is None:
+                        raise ValueError(f'Column {source_column} not found in table {self._table_name}')
+                    mutation_kwargs['source_value'] = table_values[source_index]
path: src/pg_stage/obfuscator.py
user: login: ilyichumakov
body: берётся исходное значение колонки, если на колонке также есть мутация, то возможен рассинхрон
created_at: 2025-03-05T11:01:12Z
author_association: COLLABORATOR
diff_hunk: @@ -222,10 +248,14 @@ def _prepared_data(self, *, line: str) -> Optional[str]:
 
                     self._relation_values[relation_fk] = new_value
 
-                result.append(new_value)
+                obfuscated_values[column_name] = new_value
                 is_obfuscated = True
 
-        return self.delimiter.join(result)
+        sorted_result: list[str | None] = [None] * len(self._enumerate_table_columns)
path: src/pg_stage/obfuscator.py
user: login: froOzzy
body: Такой тип типизации не поддерживается в python3.9, а библиотека пока что поддерживает эту версию
 `list[str | None]`
created_at: 2025-03-07T17:35:25Z
author_association: OWNER
diff_hunk: @@ -182,16 +201,23 @@ def _prepared_data(self, *, line: str) -> Optional[str]:
                 mutation_kwargs = mutation_for_column['mutation_kwargs']
                 mutation_relations = mutation_for_column['mutation_relations']
                 mutation_conditions = mutation_for_column['mutation_conditions']
+                mutation_source_column = mutation_for_column['mutation_source_column']
                 if not self._checking_conditions(conditions=mutation_conditions, table_values=table_values):
                     if mutation_index + 1 == len_mutations_for_column:
-                        result.append(table_values[column_index])
+                        obfuscated_values[column_name] = table_values[column_index]
                         break
 
                     continue
 
+                if mutation_source_column:
+                    source_index = self._enumerate_table_columns.get(mutation_source_column)
+                    if source_index is None:
+                        raise ValueError(f'Column {mutation_source_column} not found in table {self._table_name}')
+                    mutation_kwargs['source_value'] = obfuscated_values[mutation_source_column]
path: src/pg_stage/obfuscator.py
user: login: froOzzy
body: Т.к. все мутаторы на вход получают `kwargs`, то можно передать `obfuscated_values` в аргументы вызова мутатора - есть подозрение что дальше в работе это будет актуально многим мутаторам (связка из уже обфусцированных данных). После уже проверять внутри все необходимое. Так же это будет проще для понимания, т.к. это условие нужно сейчас лишь для одного мутатора
created_at: 2025-03-07T17:42:12Z
author_association: OWNER
diff_hunk: @@ -516,3 +516,36 @@ def mutation_uuid4(**kwargs: Any) -> str:
         :return: строка uuid4
         """
         return str(uuid.uuid4())
+
+    @staticmethod
+    def mutation_uuid5_by_source_value(**kwargs: Any) -> str:
+        """
+        Метод для формирования uuid5 с использованием значения из указанной колонки,
+        текущей даты и переданного uuid namespace.
+
+        :param kwargs:
+            source_column: Название колонки, значение которой хотим использовать.
+            namespace: Uuid namespace.
+        :return: Строка uuid5.
+        """
+        source_column: str = kwargs.get('source_column')
+        if not source_column:
+            raise ValueError('Argument "source_column" not found')
+
+        obfuscated_values: dict = kwargs.get('obfuscated_values')
+        source_value: str = obfuscated_values.get(source_column)
path: src/pg_stage/mutator.py
user: login: froOzzy
body: Можем упасть до проверки значения ниже, т.к. obfuscated_values может быть None
created_at: 2025-03-09T05:38:10Z
author_association: OWNER
diff_hunk: @@ -516,3 +516,36 @@ def mutation_uuid4(**kwargs: Any) -> str:
         :return: строка uuid4
         """
         return str(uuid.uuid4())
+
+    @staticmethod
+    def mutation_uuid5_by_source_value(**kwargs: Any) -> str:
+        """
+        Метод для формирования uuid5 с использованием значения из указанной колонки,
+        текущей даты и переданного uuid namespace.
+
+        :param kwargs:
+            source_column: Название колонки, значение которой хотим использовать.
+            namespace: Uuid namespace.
+        :return: Строка uuid5.
+        """
+        source_column: str = kwargs.get('source_column')
+        if not source_column:
+            raise ValueError('Argument "source_column" not found')
+
+        obfuscated_values: dict = kwargs.get('obfuscated_values')
path: src/pg_stage/mutator.py
user: login: froOzzy
body: dict[str, Any]
created_at: 2025-03-09T05:38:24Z
author_association: OWNER
diff_hunk: @@ -516,3 +516,36 @@ def mutation_uuid4(**kwargs: Any) -> str:
         :return: строка uuid4
         """
         return str(uuid.uuid4())
+
+    @staticmethod
+    def mutation_uuid5_by_source_value(**kwargs: Any) -> str:
+        """
+        Метод для формирования uuid5 с использованием значения из указанной колонки,
+        текущей даты и переданного uuid namespace.
+
+        :param kwargs:
+            source_column: Название колонки, значение которой хотим использовать.
+            namespace: Uuid namespace.
+        :return: Строка uuid5.
+        """
+        source_column: str = kwargs.get('source_column')
+        if not source_column:
+            raise ValueError('Argument "source_column" not found')
+
+        obfuscated_values: dict = kwargs.get('obfuscated_values')
+        source_value: str = obfuscated_values.get(source_column)
+        if not source_value:
+            raise ValueError('Value of "source_column" not found')
+
+        uuid_namespace = kwargs.get('namespace')
+        if uuid_namespace is None:
+            raise ValueError('Argument "namespace" not found')
+
+        if not isinstance(uuid_namespace, uuid.UUID):
path: src/pg_stage/mutator.py
user: login: froOzzy
body: uuid_namespace всегда будет строкой, по другому через комментарии к БД нельзя передать этот параметр
created_at: 2025-03-09T05:39:15Z
author_association: OWNER
diff_hunk: @@ -516,3 +516,36 @@ def mutation_uuid4(**kwargs: Any) -> str:
         :return: строка uuid4
         """
         return str(uuid.uuid4())
+
+    @staticmethod
+    def mutation_uuid5_by_source_value(**kwargs: Any) -> str:
+        """
+        Метод для формирования uuid5 с использованием значения из указанной колонки,
+        текущей даты и переданного uuid namespace.
+
+        :param kwargs:
+            source_column: Название колонки, значение которой хотим использовать.
+            namespace: Uuid namespace.
+        :return: Строка uuid5.
+        """
+        source_column: str = kwargs.get('source_column')
+        if not source_column:
+            raise ValueError('Argument "source_column" not found')
+
+        obfuscated_values: dict = kwargs.get('obfuscated_values')
+        source_value: str = obfuscated_values.get(source_column)
+        if not source_value:
+            raise ValueError('Value of "source_column" not found')
+
+        uuid_namespace = kwargs.get('namespace')
+        if uuid_namespace is None:
+            raise ValueError('Argument "namespace" not found')
+
+        if not isinstance(uuid_namespace, uuid.UUID):
+            try:
+                uuid_namespace = uuid.UUID(str(uuid_namespace))
+            except Exception as e:
+                raise ValueError('Invalid uuid namespace given') from e
+
+        date_today = datetime.date.today()
path: src/pg_stage/mutator.py
user: login: froOzzy
body: Можно (или уже там есть) вынести в __init__
created_at: 2025-03-09T05:39:53Z
author_association: OWNER
diff_hunk: @@ -126,6 +126,7 @@ def _parse_comment_column(self, *, line: str) -> str:
                     'mutation_kwargs': mutation_param.get('mutation_kwargs', {}),
                     'mutation_relations': mutation_param.get('relations', []),
                     'mutation_conditions': mutation_param.get('conditions', []),
+                    'mutation_source_column': mutation_param.get('source_column'),
path: src/pg_stage/obfuscator.py
user: login: froOzzy
body: Кажется это уже лишнее
created_at: 2025-03-09T05:41:19Z
author_association: OWNER

PR comments:
user: login: Makcxim
created_at: 2025-03-07T06:39:22Z
body: Теперь:  

для обфускации используется не исходный телефон, а уже измененный    
 
Что делаем:   
сортируем колонки на основе наличия `source_column`    
на данный момент нет возможности указать на колонку, у которой уже проставлен `source_column`    
(не гарантируется верная сортировка в таком случае)  

после проходимся по столбцам в нужной сортировке, на момент прохождения столбца с `source_column`  
в словарике уже будет находиться обфусцированное значение нужного столбца, берем его оттуда  

убрал `result`, тк он делал то же самое, но в другой структуре  

в конце значения возвращаются в исходную последовательность 


Добавил тест
author_association: CONTRIBUTOR
user: login: Makcxim
created_at: 2025-03-08T16:30:05Z
body: перенес `source_column `внутрь `mutation_kwargs`, тк значение нужно только внутри мутатора  

теперь во все колонки с `source_column` передаем все обфусцированные значения, чтобы не добавлять проверок в обфускатор  

внутри мутаторов будем проверять на нужные значения  


author_association: CONTRIBUTOR