From 61b22359638ebed40e3ff49aeb05d7531e05b5a3 Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Mon, 6 Oct 2025 18:27:38 +0300 Subject: [PATCH 1/7] Fix alembic.ini file not being found Error: No 'script_location' key found in configuration --- definitions/config.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/definitions/config.py b/definitions/config.py index 195ebd6..39361d9 100644 --- a/definitions/config.py +++ b/definitions/config.py @@ -31,6 +31,12 @@ def main( # behavior changed incompatibly in py3.3 command_line.parser.error("too few arguments") Config.get_template_directory = get_template_directory # type: ignore + if options.config is None: + if os.path.isfile("alembic.ini"): + options.config = "alembic.ini" + else: + raise EnvironmentError("File alembic.ini does not exist and was not provided in command line. See --help.") + config = Config( file_=options.config, ini_section=options.name, From 185e4df6092bf04c29d9b9b7e0ffe4586cddf4ef Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Mon, 3 Nov 2025 23:58:44 +0300 Subject: [PATCH 2/7] Finished schema comparator refractoring TODO: finish table comparator & encryption operations --- definitions/config.py | 4 +- definitions/custom_scripts/__init__.py | 16 +-- .../custom_scripts/operations_encrypt.py | 90 +++++++++++++ definitions/custom_scripts/schemas.py | 118 +++++++----------- definitions/custom_scripts/tables.py | 101 +++++++++------ 5 files changed, 197 insertions(+), 132 deletions(-) create mode 100644 definitions/custom_scripts/operations_encrypt.py diff --git a/definitions/config.py b/definitions/config.py index 39361d9..6b12df1 100644 --- a/definitions/config.py +++ b/definitions/config.py @@ -35,7 +35,9 @@ def main( if os.path.isfile("alembic.ini"): options.config = "alembic.ini" else: - raise EnvironmentError("File alembic.ini does not exist and was not provided in command line. See --help.") + raise EnvironmentError( + "File alembic.ini does not exist and was not provided in command line. See --help." + ) config = Config( file_=options.config, diff --git a/definitions/custom_scripts/__init__.py b/definitions/custom_scripts/__init__.py index fa343af..5d979b6 100644 --- a/definitions/custom_scripts/__init__.py +++ b/definitions/custom_scripts/__init__.py @@ -6,17 +6,5 @@ render_drop_group, render_drop_sequence, ) -from .schemas import add_table_schema_to_model, compare_schemas - -__all__ = [ - "create_table_schema", - "drop_table_schema", - "render_create_sequence", - "render_drop_sequence", - "add_table_schema_to_model", - "compare_schemas", - "create_group", - "delete_group", - "render_create_group", - "render_drop_group", -] +from .schemas import add_table_schema_to_model, compare_for_groups +from .tables import compare_for_encrypted, compare_for_sensitive diff --git a/definitions/custom_scripts/operations_encrypt.py b/definitions/custom_scripts/operations_encrypt.py new file mode 100644 index 0000000..198635c --- /dev/null +++ b/definitions/custom_scripts/operations_encrypt.py @@ -0,0 +1,90 @@ +from alembic.operations import MigrateOperation, Operations + + +@Operations.register_operation("encrypt_table") +class EncryptTableOp(MigrateOperation): + def __init__(self, table_name, key_table_name, id_column, columns): + self.table_name = table_name + self.key_table_name = key_table_name + self.id_column = id_column + self.columns = columns + + @classmethod + def encrypt_table( + cls, operations, table_name, key_table_name=None, id_column="id", columns=None + ): + op = EncryptTableOp(table_name, key_table_name, id_column) + return operations.invoke(op) + + def reverse(self): + return DecryptTableOp( + self.table_name, self.key_table_name, self.id_column, self.columns + ) + + +@Operations.register_operation("decrypt_table") +class DecryptTableOp(MigrateOperation): + def __init__(self, table_name, key_table_name, id_column, columns): + self.table_name = table_name + self.key_table_name = key_table_name + self.id_column = id_column + self.columns = columns + + @classmethod + def decrypt_table(cls, operations, table_name, key_table_name, id_column, columns): + op = DecryptTableOp(table_name, key_table_name, id_column, columns) + return operations.invoke(op) + + def reverse(self): + return EncryptTableOp( + self.table_name, self.key_table_name, self.id_column, self.columns + ) + + +@Operations.register_operation("encrypt_column") +class EncryptColumnOp(MigrateOperation): + def __init__(self, column_name, key_table_name, id_column): + self.column_name = column_name + self.key_table_name = key_table_name + self.id_column = id_column + + @classmethod + def encrypt_column(cls, operations, column_name, key_table_name, id_column): + op = EncryptColumnOp(column_name, key_table_name, id_column) + return operations.invoke(op) + + def reverse(self): + return DecryptColumnOp(self.column_name, self.key_table_name, self.id_column) + + +@Operations.register_operation("decrypt_column") +class DecryptColumnOp(MigrateOperation): + def __init__(self, column_name, key_table_name, id_column): + self.column_name = column_name + self.key_table_name = key_table_name + self.id_column = id_column + + @classmethod + def decrypt_column(cls, operations, column_name, key_table_name, id_column): + op = DecryptColumnOp(column_name, key_table_name, id_column) + return operations.invoke(op) + + def reverse(self): + return EncryptColumnOp(self.column_name, self.key_table_name, self.id_column) + + +@Operations.implementation_for(EncryptTableOp) +def encrypt_table(operations, operation): + print(type(operations), type(operation)) + + +@Operations.implementation_for(DecryptTableOp) +def decrypt_table(operations, operation): ... + + +@Operations.implementation_for(EncryptColumnOp) +def encrypt_column(operations, operation): ... + + +@Operations.implementation_for(DecryptColumnOp) +def decrypt_column(operations, operation): ... diff --git a/definitions/custom_scripts/schemas.py b/definitions/custom_scripts/schemas.py index c6b9fce..ac28b61 100644 --- a/definitions/custom_scripts/schemas.py +++ b/definitions/custom_scripts/schemas.py @@ -13,15 +13,19 @@ from .operations_tables import GrantRightsOp, RevokeRightsOp +# this function is exported and not used inside this file, do not delete def add_table_schema_to_model(table_schema, metadata): metadata.info.setdefault("table_schemas", set()).add(table_schema) @comparators.dispatch_for("schema") -def compare_schemas(autogen_context, upgrade_ops, schemas): +def compare_for_groups(autogen_context, upgrade_ops, schemas): + environment = "test" if os.getenv("ENVIRONMENT") != "production" else "prod" + project_prefix = os.getenv("SCHEMA_PREFIX", "dwh") all_conn_schemas = set() default_pg_schemas = ["pg_toast", "information_schema", "public", "pg_catalog"] - query = text("select schema_name from information_schema.schemata") + query = text("SELECT schema_name FROM information_schema.schemata") + # all schemas in database all_conn_schemas.update( [ sch[0] @@ -29,88 +33,50 @@ def compare_schemas(autogen_context, upgrade_ops, schemas): if sch[0] not in default_pg_schemas ] ) - + # all schemas in code metadata_schemas = autogen_context.metadata.info.setdefault("table_schemas", set()) + # Create/delete new schemas for sch in metadata_schemas - all_conn_schemas: upgrade_ops.ops.append(CreateTableSchemaOp(sch)) - for render_scope in ["read", "write", "all"]: - group_name = ( - f"test_dwh_{sch}_{render_scope}".lower() - if os.getenv("ENVIRONMENT") != "production" - else f"prod_dwh_{sch}_{render_scope}".lower() - ) - upgrade_ops.ops.append(CreateGroupOp(group_name)) - upgrade_ops.ops.append(GrantOnSchemaOp(group_name, sch)) - - tables = set( - [ - table - for table in autogen_context.metadata.tables.values() - if table.schema == sch - ] - ) - for table in tables: - for render_scope in ["read", "write", "all"]: - scopes = [] - match render_scope: - case "read": - scopes = ["SELECT"] - case "write": - scopes = ["SELECT", "UPDATE", "DELETE", "TRUNCATE", "INSERT"] - case "all": - scopes = ["ALL"] - - group_name = ( - f"test_dwh_{sch}_{render_scope}".lower() - if os.getenv("ENVIRONMENT") != "production" - else f"prod_dwh_{sch}_{render_scope}".lower() - ) - upgrade_ops.ops.append( - GrantRightsOp( - table_name=str(table), scopes=scopes, group_name=group_name - ) - ) - for sch in all_conn_schemas - metadata_schemas: upgrade_ops.ops.append(DropTableSchemaOp(sch)) + + all_groups_db = set() + all_groups_code = set() + query = text( + "SELECT grantee,table_schema FROM information_schema.role_table_grants " + "WHERE grantee LIKE :pattern" + ) + all_groups_db.update( + autogen_context.connection.execute( + query, {"pattern": f"{environment}%{project_prefix}%"} + ) + ) + for sch in metadata_schemas: + has_regular = any( + table.schema == sch and not table.info.get("sensitive", False) + for table in autogen_context.metadata.tables.values() + ) + has_sensitive = any( + table.schema == sch and table.info.get("sensitive", False) + for table in autogen_context.metadata.tables.values() + ) for render_scope in ["read", "write", "all"]: group_name = ( - f"test_dwh_{sch}_{render_scope}".lower() - if os.getenv("ENVIRONMENT") != "production" - else f"prod_dwh_{sch}_{render_scope}".lower() + f"{environment}%s_{project_prefix}_{sch}_{render_scope}".lower() ) - upgrade_ops.ops.append(DeleteGroupOp(group_name)) - upgrade_ops.ops.append(RevokeOnSchemaOp(group_name, sch)) + if has_regular: + all_groups_code.add((group_name % "", sch)) + if has_sensitive: + all_groups_code.add((group_name % "_sensitive", sch)) - query = text( - f"SELECT * FROM information_schema.tables WHERE table_schema='{sch}';" - ) - tables = set( - [ - ".".join(table[1:3]) - for table in autogen_context.connection.execute(query) - ] - ) - print(tables) - for table in tables: - for render_scope in ["read", "write", "all"]: - scopes = [] - match render_scope: - case "read": - scopes = ["SELECT"] - case "write": - scopes = ["SELECT", "UPDATE", "DELETE", "TRUNCATE", "INSERT"] - case "all": - scopes = ["ALL"] + # for all new required groups + for group, sch in all_groups_code - all_groups_db: + upgrade_ops.ops.append(CreateGroupOp(group)) + upgrade_ops.ops.append(GrantOnSchemaOp(group, sch)) - group_name = ( - f"test_dwh_{sch}_{render_scope}".lower() - if os.getenv("ENVIRONMENT") != "production" - else f"prod_dwh_{sch}_{render_scope}".lower() - ) - upgrade_ops.ops.append( - RevokeRightsOp( - table_name=str(table), scopes=scopes, group_name=group_name - ) - ) + # for all groups no longer needed + for group, sch in all_groups_db - all_groups_code: + upgrade_ops.ops.append(DeleteGroupOp(group)) + upgrade_ops.ops.append(RevokeOnSchemaOp(group, sch)) diff --git a/definitions/custom_scripts/tables.py b/definitions/custom_scripts/tables.py index 8cc4c33..293ba89 100644 --- a/definitions/custom_scripts/tables.py +++ b/definitions/custom_scripts/tables.py @@ -2,34 +2,67 @@ from alembic.autogenerate import comparators +from .operations_encrypt import DecryptTableOp, EncryptTableOp from .operations_tables import GrantRightsOp, RevokeRightsOp -@comparators.dispatch_for("table") -def compare_table( - autogen_context, modify_table_ops, s, tname, metadata_table_db, metadata_table_code -): - if str(metadata_table_db) == "None": - sensitive = metadata_table_code.info.get("sensitive", False) - for render_scope in ["read", "write", "all"]: - group_name = ( - f'test_{"sensitive_" if sensitive else ""}dwh_{s}_{render_scope}'.lower() - if os.getenv("ENVIRONMENT") != "production" - else f'prod_{"sensitive_" if sensitive else ""}dwh_{s}_{render_scope}'.lower() +def compare_for_encrypted(table_db, table_code): + db_info = dict() if table_db is None else table_db.info + code_info = dict() if table_code is None else table_code.info + encrypted_db = db_info.get("encrypted", False) + encrypted_code = code_info.get("encrypted", False) + if encrypted_code and encrypted_db: + if db_info != code_info: + print( + "Error: changing encryption parameters is not supported.\n" + "Please edit migration file manually, using EncryptTableOp and DecryptTableOp", + file=stderr, + ) + exit(1) + if encrypted_code and not encrypted_db: + modify_table_ops.append( + EncryptTableOp( + table_code.name, + code_info.encryption["keys"], + code_info.encryption["id"], + code_info.encryption["columns"], ) + ) + if encrypted_db and not encrypted_code: + modify_table_ops.append( + DecryptTableOp( + table_code.name, + code_info.encryption["keys"], + code_info.encryption["id"], + code_info.encryption["columns"], + ) + ) - scopes = [] - match render_scope: - case "read": - scopes = ["SELECT"] - case "write": - scopes = ["SELECT", "UPDATE", "DELETE", "TRUNCATE", "INSERT"] - case "all": - scopes = ["ALL"] - if sensitive: - modify_table_ops.ops.append(CreateGroupOp(group_name=group_name)) +@comparators.dispatch_for("table") +def compare_for_sensitive( + autogen_context, + modify_table_ops, + sch, + tname, + metadata_table_db, + metadata_table_code, +): + environment = "test" if os.getenv("ENVIRONMENT") != "production" else "prod" + project_prefix = os.getenv("SCHEMA_PREFIX", "dwh") + render_scope_map = { + "read": ["SELECT"], + "write": ["SELECT", "UPDATE", "DELETE", "TRUNCATE", "INSERT"], + "all": ["ALL"], + } + # опять неправильный подход от миши, который не учитывает изменяющиеся таблицы + if metadata_table_db is None: + sensitive = ( + "_sensitive" if metadata_table_code.info.get("sensitive", False) else "" + ) + for render_scope, scopes in render_scope_map.items(): + group_name = f"{environment}{sensitive}_{project_prefix}_{sch}_{render_scope}".lower() modify_table_ops.ops.append( GrantRightsOp( table_name=str(metadata_table_code), @@ -38,23 +71,12 @@ def compare_table( ) ) - elif str(metadata_table_code) == "None": - sensitive = metadata_table_db.info.get("sensitive", False) - for render_scope in ["read", "write", "all"]: - group_name = ( - f'test_{"sensitive_" if sensitive else ""}dwh_{s}_{render_scope}'.lower() - if os.getenv("ENVIRONMENT") != "production" - else f'prod_{"sensitive_" if sensitive else ""}dwh_{s}_{render_scope}'.lower() - ) - scopes = [] - match render_scope: - case "read": - scopes = ["SELECT"] - case "write": - scopes = ["SELECT", "UPDATE", "DELETE", "TRUNCATE", "INSERT"] - case "all": - scopes = ["ALL"] - + if metadata_table_code is None: + sensitive = ( + "_sensitive" if metadata_table_db.info.get("sensitive", False) else "" + ) + for render_scope, scopes in render_scope_map.items(): + group_name = f"{environment}{sensitive}_{project_prefix}_{sch}_{render_scope}".lower() modify_table_ops.ops.append( RevokeRightsOp( table_name=str(metadata_table_db), @@ -62,6 +84,3 @@ def compare_table( group_name=group_name, ) ) - - if sensitive: - modify_table_ops.ops.append(DeleteGroupOp(group_name=group_name)) From 4075fbe73ebe58c248a38304f7393f3e182f2481 Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Tue, 4 Nov 2025 01:28:59 +0300 Subject: [PATCH 3/7] Finished groups/roles refractoring --- .gitignore | 1 + definitions/custom_scripts/schemas.py | 9 ++-- definitions/custom_scripts/tables.py | 70 ++++++++++++++++++--------- 3 files changed, 52 insertions(+), 28 deletions(-) diff --git a/.gitignore b/.gitignore index 6368327..3d6abe3 100644 --- a/.gitignore +++ b/.gitignore @@ -106,6 +106,7 @@ celerybeat.pid .venv env/ venv/ +venv ENV/ env.bak/ venv.bak/ diff --git a/definitions/custom_scripts/schemas.py b/definitions/custom_scripts/schemas.py index ac28b61..d37452b 100644 --- a/definitions/custom_scripts/schemas.py +++ b/definitions/custom_scripts/schemas.py @@ -62,14 +62,13 @@ def compare_for_groups(autogen_context, upgrade_ops, schemas): table.schema == sch and table.info.get("sensitive", False) for table in autogen_context.metadata.tables.values() ) + + group_name = f"{environment}%s_{project_prefix}_{sch}_%s".lower() for render_scope in ["read", "write", "all"]: - group_name = ( - f"{environment}%s_{project_prefix}_{sch}_{render_scope}".lower() - ) if has_regular: - all_groups_code.add((group_name % "", sch)) + all_groups_code.add((group_name % ("", render_scope), sch)) if has_sensitive: - all_groups_code.add((group_name % "_sensitive", sch)) + all_groups_code.add((group_name % ("_sensitive", render_scope), sch)) # for all new required groups for group, sch in all_groups_code - all_groups_db: diff --git a/definitions/custom_scripts/tables.py b/definitions/custom_scripts/tables.py index 293ba89..735702e 100644 --- a/definitions/custom_scripts/tables.py +++ b/definitions/custom_scripts/tables.py @@ -1,12 +1,16 @@ import os from alembic.autogenerate import comparators +from sqlalchemy import text from .operations_encrypt import DecryptTableOp, EncryptTableOp from .operations_tables import GrantRightsOp, RevokeRightsOp -def compare_for_encrypted(table_db, table_code): +# @comparators.dispatch_for("table") +def compare_for_encrypted( + autogen_context, modify_table_ops, sch, tname, table_db, table_code +): db_info = dict() if table_db is None else table_db.info code_info = dict() if table_code is None else table_code.info encrypted_db = db_info.get("encrypted", False) @@ -39,6 +43,12 @@ def compare_for_encrypted(table_db, table_code): ) +# TODO: check for all existing rights, don't assume that they come from render_scope_map +# query = text("SELECT grantee, privilege_type FROM information_schema.role_table_grants " +# "WHERE table_schema=:schema AND table_name=:tablename AND grantee ILIKE :pattern") +# params = { "schema": sch, "tablename": metadata_table_db.name, "pattern": f"{environment}%{project_prefix}_{sch}_%"} +# for role, scope in autogen_context.connection.execute(query, params): +# current_rights.setdefault(role, []).append(scope) @comparators.dispatch_for("table") def compare_for_sensitive( autogen_context, @@ -55,32 +65,46 @@ def compare_for_sensitive( "write": ["SELECT", "UPDATE", "DELETE", "TRUNCATE", "INSERT"], "all": ["ALL"], } + required_rights = set() + current_rights = set() - # опять неправильный подход от миши, который не учитывает изменяющиеся таблицы - if metadata_table_db is None: + # extract rights from database + if metadata_table_db is not None: + query = text( + "SELECT DISTINCT grantee FROM information_schema.role_table_grants " + "WHERE table_schema=:schema AND table_name=:tablename AND grantee ILIKE :pattern" + ) + params = { + "schema": sch, + "tablename": metadata_table_db.name, + "pattern": f"{environment}%{project_prefix}_{sch}_%", + } + current_rights.update( + row.grantee for row in autogen_context.connection.execute(query, params) + ) + + # get required rights + if metadata_table_code is not None: sensitive = ( "_sensitive" if metadata_table_code.info.get("sensitive", False) else "" ) - for render_scope, scopes in render_scope_map.items(): - group_name = f"{environment}{sensitive}_{project_prefix}_{sch}_{render_scope}".lower() - modify_table_ops.ops.append( - GrantRightsOp( - table_name=str(metadata_table_code), - scopes=scopes, - group_name=group_name, - ) - ) + group_name = f"{environment}%s_{project_prefix}_{sch}_%s".lower() + required_rights.update( + group_name % (sensitive, render_scope) + for render_scope in render_scope_map.keys() + ) - if metadata_table_code is None: - sensitive = ( - "_sensitive" if metadata_table_db.info.get("sensitive", False) else "" + for group in required_rights - current_rights: + scope = group[group.rfind("_") + 1 :] + modify_table_ops.ops.append( + GrantRightsOp( + table_name=tname, scopes=render_scope_map[scope], group_name=group + ) ) - for render_scope, scopes in render_scope_map.items(): - group_name = f"{environment}{sensitive}_{project_prefix}_{sch}_{render_scope}".lower() - modify_table_ops.ops.append( - RevokeRightsOp( - table_name=str(metadata_table_db), - scopes=scopes, - group_name=group_name, - ) + for group in current_rights - required_rights: + scope = group[group.rfind("_") + 1 :] + modify_table_ops.ops.append( + RevokeRightsOp( + table_name=tname, scopes=render_scope_map[scope], group_name=group ) + ) From ff983a3c5e9fc70e5aa030e7805bd03e99311eb1 Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Fri, 14 Nov 2025 23:58:16 +0300 Subject: [PATCH 4/7] Started work on encryption operators --- .../custom_scripts/operations_encrypt.py | 133 ++++++++++++++++-- .../custom_scripts/operations_tables.py | 7 + definitions/custom_scripts/tables.py | 12 +- 3 files changed, 134 insertions(+), 18 deletions(-) diff --git a/definitions/custom_scripts/operations_encrypt.py b/definitions/custom_scripts/operations_encrypt.py index 198635c..9e574ed 100644 --- a/definitions/custom_scripts/operations_encrypt.py +++ b/definitions/custom_scripts/operations_encrypt.py @@ -1,19 +1,93 @@ from alembic.operations import MigrateOperation, Operations +from sqlalchemy import text +from enum import IntEnum + + +def _get_column_names(schema, table, id_column): + return [ + i.column_name + for i in operations.get_bind().execute( + text( + f"SELECT column_name FROM information_schema.columns " + f"WHERE table_schema='{schema}' AND table_name='{name}';" + ) + ) + if i != id_column + ] -@Operations.register_operation("encrypt_table") -class EncryptTableOp(MigrateOperation): - def __init__(self, table_name, key_table_name, id_column, columns): +class op_type(IntEnum): + ENCRYPT = 0x00 + DECRYPT = 0x01 + + +def _common_constructor(self, table_name, key_table_name, id_column, columns, **kw): + _name = table_name.split(".") + if len(_name) != 2: self.table_name = table_name + else: + self.table_schema, self.table_name = _name + if (sch := kw.get("schema", None)) is not None: + self.table_schema = sch + + _key_name = key_table_name.split(".") + if len(_key_name) < 2: self.key_table_name = key_table_name + else: + self.key_table_schema, self.key_table_name = _key_name + if (sch := kw.get("key_schema", None)) is not None: + self.key_table_schema = sch + + self.id_column = id_column + self.columns = columns + + +@Operations.register_operation("encrypt_table") +class EncryptTableOp(MigrateOperation): + """fields: + self.table_name: str + self.table_schema: str + self.key_table_name: str + self.key_table_schema: str + self.id_column: str + self.columns: list[str] + """ + + def __init__(self, table_name, key_table_name, id_column, columns, **kw): + _name = table_name.split(".") + if len(_name) != 2: + self.table_name = table_name + else: + self.table_schema, self.table_name = _name + if (sch := kw.get("schema", None)) is not None: + self.table_schema = sch + + _key_name = key_table_name.split(".") + if len(_key_name) < 2: + self.key_table_name = key_table_name + else: + self.key_table_schema, self.key_table_name = _key_name + if (sch := kw.get("key_schema", None)) is not None: + self.key_table_schema = sch + self.id_column = id_column self.columns = columns @classmethod def encrypt_table( - cls, operations, table_name, key_table_name=None, id_column="id", columns=None + cls, + operations, + table_name, + key_table_name=None, + id_column="id", + columns=None, + **kw, ): - op = EncryptTableOp(table_name, key_table_name, id_column) + if columns is None: + columns = _get_column_names(*table_name.split("."), id_column) + if key_table_name is None: + key_table_name = table_name + "_ekeys" + op = EncryptTableOp(table_name, key_table_name, id_column, columns, **kw) return operations.invoke(op) def reverse(self): @@ -24,15 +98,15 @@ def reverse(self): @Operations.register_operation("decrypt_table") class DecryptTableOp(MigrateOperation): - def __init__(self, table_name, key_table_name, id_column, columns): - self.table_name = table_name - self.key_table_name = key_table_name - self.id_column = id_column - self.columns = columns + "same fields as in EncryptTableOp" + + __init__ = _common_constructor @classmethod - def decrypt_table(cls, operations, table_name, key_table_name, id_column, columns): - op = DecryptTableOp(table_name, key_table_name, id_column, columns) + def decrypt_table( + cls, operations, table_name, key_table_name, id_column, columns, **kw + ): + op = DecryptTableOp(table_name, key_table_name, id_column, columns, **kw) return operations.invoke(op) def reverse(self): @@ -73,13 +147,44 @@ def reverse(self): return EncryptColumnOp(self.column_name, self.key_table_name, self.id_column) +def _generate_keygen_query(operation) -> str: + return ( + f'INSERT INTO "{operation.key_table_schema}".{operation.key_table_name} ' + f"SELECT src.{operation.id_column}, encode(gen_random_bytes(32), 'base64'), NOW() " + f'FROM "{operation.table_schema}".{operation.table_name} ' + f'AS src LEFT JOIN "{operation.key_table_schema}".{operation.key_table_name} AS dest ' + f"ON dest.id = src.{operation.id_column} WHERE dest.id IS NULL;" + ) + + +def _generate_encryption_query(operation, func: op_type) -> str: + set_cols = [] + encrypt_cols = [] + for colname in operation.columns: + set_cols.append(f"{colname} = sub.enc_{colname}") + encrypt_cols.append( + f"pgp_sym_{func.name}_bytea(dst.{colname}, keys.key) enc_{colname}" + ) + return ( + f'UPDATE "{operation.table_schema}".{operation.table_name} dst ' + f'SET {",".join(set_cols)} FROM (SELECT ' + f'dst.{operation.id_column} id,{",".join(encrypt_cols)}' + f' FROM "{operation.table_schema}".{operation.table_name} dst ' + f' LEFT JOIN "{operation.key_table_schema}".{operation.key_table_name} keys' + f" ON keys.id=dst.{operation.id_column}) sub " + f"WHERE dst.{operation.id_column}=sub.id;" + ) + + @Operations.implementation_for(EncryptTableOp) def encrypt_table(operations, operation): - print(type(operations), type(operation)) + operations.execute(_generate_keygen_query(operation)) + operations.execute(_generate_encryption_query(operation, op_type.ENCRYPT)) @Operations.implementation_for(DecryptTableOp) -def decrypt_table(operations, operation): ... +def decrypt_table(operations, operation): + operations.execute(_generate_encryption_query(operation, op_type.DECRYPT)) @Operations.implementation_for(EncryptColumnOp) diff --git a/definitions/custom_scripts/operations_tables.py b/definitions/custom_scripts/operations_tables.py index 7ff749b..80070e2 100644 --- a/definitions/custom_scripts/operations_tables.py +++ b/definitions/custom_scripts/operations_tables.py @@ -1,5 +1,12 @@ from alembic.operations import MigrateOperation, Operations +from alembic.operations.ops import CreateTableOp + +# TODO: replace CreateTableOp and DropTableOp implementations +# with custom ones +# Will be possible in alembic 1.17.2+ +# till then, we wait + @Operations.register_operation("grant_on_table") class GrantRightsOp(MigrateOperation): diff --git a/definitions/custom_scripts/tables.py b/definitions/custom_scripts/tables.py index 735702e..cb7e0b9 100644 --- a/definitions/custom_scripts/tables.py +++ b/definitions/custom_scripts/tables.py @@ -26,7 +26,7 @@ def compare_for_encrypted( if encrypted_code and not encrypted_db: modify_table_ops.append( EncryptTableOp( - table_code.name, + f"{sch}.{table_code.name}", code_info.encryption["keys"], code_info.encryption["id"], code_info.encryption["columns"], @@ -35,7 +35,7 @@ def compare_for_encrypted( if encrypted_db and not encrypted_code: modify_table_ops.append( DecryptTableOp( - table_code.name, + f"{sch}.{table_code.name}", code_info.encryption["keys"], code_info.encryption["id"], code_info.encryption["columns"], @@ -98,13 +98,17 @@ def compare_for_sensitive( scope = group[group.rfind("_") + 1 :] modify_table_ops.ops.append( GrantRightsOp( - table_name=tname, scopes=render_scope_map[scope], group_name=group + table_name=f"{sch}.{tname}", + scopes=render_scope_map[scope], + group_name=group, ) ) for group in current_rights - required_rights: scope = group[group.rfind("_") + 1 :] modify_table_ops.ops.append( RevokeRightsOp( - table_name=tname, scopes=render_scope_map[scope], group_name=group + table_name=f"{sch}.{tname}", + scopes=render_scope_map[scope], + group_name=group, ) ) From d7387208790faa7bf8502d49a8bd0f0ee6fea00a Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Sat, 15 Nov 2025 00:00:56 +0300 Subject: [PATCH 5/7] Formatting fix --- definitions/custom_scripts/operations_encrypt.py | 3 ++- definitions/custom_scripts/operations_tables.py | 1 - 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/definitions/custom_scripts/operations_encrypt.py b/definitions/custom_scripts/operations_encrypt.py index 9e574ed..93d38a4 100644 --- a/definitions/custom_scripts/operations_encrypt.py +++ b/definitions/custom_scripts/operations_encrypt.py @@ -1,6 +1,7 @@ +from enum import IntEnum + from alembic.operations import MigrateOperation, Operations from sqlalchemy import text -from enum import IntEnum def _get_column_names(schema, table, id_column): diff --git a/definitions/custom_scripts/operations_tables.py b/definitions/custom_scripts/operations_tables.py index 80070e2..27b3c30 100644 --- a/definitions/custom_scripts/operations_tables.py +++ b/definitions/custom_scripts/operations_tables.py @@ -1,5 +1,4 @@ from alembic.operations import MigrateOperation, Operations - from alembic.operations.ops import CreateTableOp # TODO: replace CreateTableOp and DropTableOp implementations From 5602226e141deb5999cb8aa2d6da15e4f3d1f55e Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Sat, 15 Nov 2025 00:40:43 +0300 Subject: [PATCH 6/7] Restored accidentally removed __all__ declaration --- definitions/custom_scripts/__init__.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/definitions/custom_scripts/__init__.py b/definitions/custom_scripts/__init__.py index 5d979b6..b1d123e 100644 --- a/definitions/custom_scripts/__init__.py +++ b/definitions/custom_scripts/__init__.py @@ -6,5 +6,18 @@ render_drop_group, render_drop_sequence, ) -from .schemas import add_table_schema_to_model, compare_for_groups +from .schemas import add_table_schema_to_model, compare_for_groups, compare_schemas from .tables import compare_for_encrypted, compare_for_sensitive + +__all__ = [ + "create_table_schema", + "drop_table_schema", + "render_create_sequence", + "render_drop_sequence", + "add_table_schema_to_model", + "compare_schemas", + "create_group", + "delete_group", + "render_create_group", + "render_drop_group", +] From cda512b35befc09735d634d402c6f15a1528d749 Mon Sep 17 00:00:00 2001 From: justanothercatgirl Date: Sat, 15 Nov 2025 00:47:50 +0300 Subject: [PATCH 7/7] modified module exports --- definitions/custom_scripts/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/definitions/custom_scripts/__init__.py b/definitions/custom_scripts/__init__.py index b1d123e..2d95528 100644 --- a/definitions/custom_scripts/__init__.py +++ b/definitions/custom_scripts/__init__.py @@ -6,7 +6,7 @@ render_drop_group, render_drop_sequence, ) -from .schemas import add_table_schema_to_model, compare_for_groups, compare_schemas +from .schemas import add_table_schema_to_model, compare_for_groups from .tables import compare_for_encrypted, compare_for_sensitive __all__ = [ @@ -15,7 +15,9 @@ "render_create_sequence", "render_drop_sequence", "add_table_schema_to_model", - "compare_schemas", + "compare_for_groups", + "compare_for_encrypted", + "compare_for_sensitive", "create_group", "delete_group", "render_create_group",