pg-data-yaml 0.0.0__py3-none-any.whl

pg_data_yaml/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.0.0'

pg_data_yaml/analyze_envs.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import yaml
+
+from .paths import list_table_yaml_files
+
+
+class AnalyzeEnvs:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.sources = [os.path.abspath(path) for path in args.source]
+
+    def run(self) -> None:
+        self._validate()
+        synchronized_directories = open('/tmp/synchronized_directory.txt').read().split('\n')
+        env_files = [list_table_yaml_files(path) for path in self.sources]
+        all_paths = set.union(*[set(files) for files in env_files])
+        if not all_paths:
+            print('Nothing to analyze: no table files across environments')
+            return
+
+        for rel_path in sorted(all_paths):
+            lines = {}
+            postfix = ''
+            for env in self.sources:
+                file_path = os.path.join(env, rel_path)
+                if os.path.isfile(file_path):
+                    for line in yaml.safe_load(open(file_path)):
+                        line = str(line)
+                        lines[line] = lines.get(line, 0) + 1
+                else:
+                    postfix = '*'
+            identical_lines = sum(1 for count in lines.values() if count == len(self.sources))
+            table = self.path_to_table_name(rel_path)
+            sync_dir = "да" if table in synchronized_directories else ''
+            print(f'{table}\t{identical_lines}/{len(lines) - identical_lines}{postfix}\t{sync_dir}')
+
+    # def run_base(self) -> None:
+    #     synchronized_directories=open('/tmp/synchronized_directory.txt').read().split('\n')
+    #     env_files = [list_table_yaml_files(path) for path in self.sources]
+    #     all_paths = set.union(*[set(files) for files in env_files])
+    #     if not all_paths:
+    #         print('Nothing to analyze: no table files across environments')
+    #         return
+
+    #     for rel_path in sorted(all_paths):
+    #         env = self.sources[0]
+    #         file_path = os.path.join(env, rel_path)
+    #         lines_count = len(yaml.safe_load(open(file_path)))
+    #         table = self.path_to_table_name(rel_path)
+    #         sync_dir="да" if table in synchronized_directories else ''
+    #         print(f'{table}\t{lines_count}/0\t{sync_dir}')
+
+    def _validate(self) -> None:
+        if len(self.sources) < 2:
+            print('ERROR: specify at least two --source directories', file=sys.stderr)
+            sys.exit(1)
+
+        seen = set()
+        for path in self.sources:
+            if path in seen:
+                print(f'ERROR: duplicate source directory: {path}', file=sys.stderr)
+                sys.exit(1)
+            seen.add(path)
+            if not os.path.isdir(path):
+                print(f'ERROR: source directory not found: {path}', file=sys.stderr)
+                sys.exit(1)
+
+    @staticmethod
+    def path_to_table_name(path: str) -> str:
+        return '{}.{}'.format(*path[:-5].split('/'))

pg_data_yaml/extractor.py

@@ -0,0 +1,46 @@
+import argparse
+import os
+
+from .formatter import Formatter
+from .pg import Pg
+from .registry import SyncTable, TableRegistry
+from .values import ordered_row
+
+
+class Extractor:
+    ROWS_LIMIT = 50000
+
+    def __init__(self, args: argparse.Namespace, pg: Pg):
+        self.args = args
+        self.pg = pg
+        self.registry = TableRegistry.from_args(pg, args)
+        self.formatter = Formatter()
+
+    async def export(self) -> None:
+        tables = await self.registry.load()
+        for sync_table in tables.values():
+            rows_count = await self.get_rows_count(sync_table)
+            if rows_count == self.ROWS_LIMIT:
+                print(f'Table {sync_table.table} has {rows_count} rows, skipping')
+                continue
+            rows = await self.fetch_rows(sync_table)
+            file_name = sync_table.file_path(self.args.out_dir)
+            os.makedirs(os.path.dirname(file_name), exist_ok=True)
+            self.formatter.dump(rows, file_name)
+
+    async def fetch_rows(self, sync_table: SyncTable) -> list[dict]:
+        rows = await self.pg.fetch(sync_table.select_query())
+        return [ordered_row(row) for row in rows]
+
+    async def get_rows_count(self, sync_table: SyncTable) -> list[dict]:
+        query = sync_table.select_query()
+        query = f'select count(*) from ({query} limit {self.ROWS_LIMIT}) as t'
+        rows = await self.pg.fetch(query)
+        return rows[0]['count']
+
+    async def get_tables_data(self) -> dict[tuple[str, str], list[dict]]:
+        tables = await self.registry.load()
+        return {
+            key: await self.fetch_rows(sync_table)
+            for key, sync_table in tables.items()
+        }

pg_data_yaml/formatter.py

@@ -0,0 +1,49 @@
+import re
+from typing import Any
+
+import yaml
+
+_INTERVAL_STR_RE = re.compile(r'^-?(\d+ days? )?\d+:\d{2}:\d{2}$')
+
+
+class Formatter:
+    _representers_registered = False
+
+    def __init__(self):
+        self._register_representers()
+
+    @classmethod
+    def _register_representers(cls):
+        if cls._representers_registered:
+            return
+
+        def str_presenter(dumper, data):
+            if '\n' in data:
+                return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
+            if _INTERVAL_STR_RE.fullmatch(data):
+                return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='')
+            return dumper.represent_scalar('tag:yaml.org,2002:str', data)
+
+        yaml.add_representer(str, str_presenter)
+        cls._representers_registered = True
+
+    @staticmethod
+    def dump(data: Any, file_name: str = None):
+        Formatter._register_representers()
+        if data is None:
+            return ''
+        file = None
+        if file_name:
+            file = open(file_name, 'w')
+        return yaml.dump(
+            data,
+            file,
+            allow_unicode=True,
+            sort_keys=False,
+            width=float('inf'),
+        )
+
+    @staticmethod
+    def load(file_name: str) -> Any:
+        with open(file_name) as file:
+            return yaml.safe_load(file)

pg_data_yaml/main.py

@@ -0,0 +1,198 @@
+import argparse
+import asyncio
+import os
+import shutil
+import sys
+
+from .analyze_envs import AnalyzeEnvs
+from .extractor import Extractor
+from .merge_envs import MergeEnvs
+from .pg import Pg
+from .synchronizer import Synchronizer
+from . import __version__
+
+
+async def run(args):
+    pg = Pg(args)
+    await pg.init()
+
+    if args.command == 'export':
+        await Extractor(args, pg).export()
+
+    elif args.command in ('diff', 'sync'):
+        await Synchronizer(args, pg).sync(show_diff_only=args.command == 'diff')
+
+
+def main():
+    def add_connection_args(parser):
+        parser.add_argument('-d', '--dbname',
+                            type=str, help='database name to connect to')
+        parser.add_argument('-h', '--host',
+                            type=str, help='database server host or socket directory')
+        parser.add_argument('-p', '--port',
+                            type=str, help='database server port')
+        parser.add_argument('-U', '--user',
+                            type=str, help='database user name')
+        parser.add_argument('-W', '--password',
+                            type=str, help='database user password')
+
+    def add_registry_args(parser):
+        group = parser.add_mutually_exclusive_group(required=True)
+        group.add_argument(
+            '--comment-label',
+            metavar='LABEL',
+            help='include tables whose comment contains LABEL; '
+            'optional export query or WHERE filter in parentheses after the label',
+        )
+        group.add_argument(
+            '--table-list-predicate',
+            metavar='PREDICATE',
+            help='include tables matching SQL PREDICATE (table comments are ignored); '
+            'references n.nspname and c.relname from pg_catalog',
+        )
+
+    arg_parser = argparse.ArgumentParser(
+        epilog='Report bugs: https://github.com/andruche/pg-data-yaml/issues',
+        conflict_handler='resolve',
+    )
+
+    arg_parser.add_argument(
+        '--version',
+        action='version',
+        version=__version__,
+    )
+
+    subparsers = arg_parser.add_subparsers(
+        dest='command',
+        title='commands',
+    )
+
+    parser_export = subparsers.add_parser(
+        'export',
+        help='export synchronized reference tables to yaml files',
+        conflict_handler='resolve',
+    )
+    add_connection_args(parser_export)
+    add_registry_args(parser_export)
+    parser_export.add_argument(
+        '--out-dir',
+        required=True,
+        help='directory for exporting files',
+    )
+    parser_export.add_argument(
+        '--clean',
+        action='store_true',
+        help='clean out_dir if not empty '
+        '(env variable DATA_DIRECTORY_AUTOCLEAN=true)',
+    )
+
+    parser_diff = subparsers.add_parser(
+        'diff',
+        help='show diff between database and yaml files',
+        conflict_handler='resolve',
+    )
+    add_connection_args(parser_diff)
+    add_registry_args(parser_diff)
+    parser_diff.add_argument(
+        '--source',
+        required=True,
+        help='directory or file with table data to compare with database',
+    )
+
+    parser_sync = subparsers.add_parser(
+        'sync',
+        help='sync yaml files to database tables',
+        conflict_handler='resolve',
+    )
+    add_connection_args(parser_sync)
+    add_registry_args(parser_sync)
+    parser_sync.add_argument(
+        '--source',
+        required=True,
+        help='directory or file with table data to sync to database',
+    )
+    parser_sync.add_argument(
+        '--dry-run',
+        action='store_true',
+        help='test run without real changes',
+    )
+    parser_sync.add_argument(
+        '--echo-queries',
+        action='store_true',
+        help='echo commands sent to server',
+    )
+    parser_sync.add_argument(
+        '-y', '--yes',
+        action='store_true',
+        help='do not ask confirm',
+    )
+
+    parser_merge_envs = subparsers.add_parser(
+        'merge-envs',
+        help='move identical table files from env dirs to base directory',
+        conflict_handler='resolve',
+    )
+    parser_merge_envs.add_argument(
+        '--source',
+        action='append',
+        required=True,
+        metavar='ENV_DIR',
+        help='environment directory (can be repeated)',
+    )
+    parser_merge_envs.add_argument(
+        '--out-dir',
+        required=True,
+        help='base directory for common table files',
+    )
+    parser_merge_envs.add_argument(
+        '--dry-run',
+        action='store_true',
+        help='show actions without changing files',
+    )
+
+    parser_merge_envs = subparsers.add_parser(
+        'analyze-envs',
+        help='compare data between envs',
+        conflict_handler='resolve',
+    )
+    parser_merge_envs.add_argument(
+        '--source',
+        action='append',
+        required=True,
+        metavar='ENV_DIR',
+        help='environment directory (can be repeated)',
+    )
+
+    args = arg_parser.parse_args()
+    if not args.command:
+        arg_parser.print_help()
+        sys.exit(1)
+
+    if args.command == 'export':
+        if os.path.exists(args.out_dir) and os.listdir(args.out_dir):
+            if args.clean or os.environ.get('DATA_DIRECTORY_AUTOCLEAN') == 'true':
+                shutil.rmtree(args.out_dir)
+            else:
+                parser_export.error(
+                    'out_dir directory not empty (you can use option --clean)'
+                )
+        try:
+            os.makedirs(args.out_dir, exist_ok=True)
+        except OSError:
+            arg_parser.error("can not access to directory '%s'" % args.out_dir)
+
+    if args.command in ('diff', 'sync'):
+        if not os.path.exists(args.source):
+            arg_parser.error(f'file or directory not found: {args.source}')
+
+    if args.command == 'merge-envs':
+        MergeEnvs(args).run()
+        return
+
+    if args.command == 'analyze-envs':
+        AnalyzeEnvs(args).run_base()
+        return
+
+    if os.name == 'nt':
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+    asyncio.run(run(args))

pg_data_yaml/merge_envs.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import sys
+
+from .paths import list_table_yaml_files
+
+
+class MergeEnvs:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.sources = [os.path.abspath(path) for path in args.source]
+        self.out_dir = os.path.abspath(args.out_dir)
+
+    def run(self) -> None:
+        self._validate()
+        env_files = [list_table_yaml_files(path) for path in self.sources]
+        common_paths = set.intersection(*[set(files) for files in env_files])
+        if not common_paths:
+            print('Nothing to merge: no common table files across environments')
+            return
+
+        merged = 0
+        skipped_diff = 0
+        skipped_base_conflict = 0
+
+        for rel_path in sorted(common_paths):
+            source_paths = [files[rel_path] for files in env_files]
+            if not self._files_identical(source_paths):
+                print(
+                    f'SKIP: {rel_path} differs between environments',
+                    file=sys.stderr,
+                )
+                skipped_diff += 1
+                continue
+
+            base_path = os.path.join(self.out_dir, rel_path)
+            if os.path.exists(base_path):
+                if not self._files_identical([source_paths[0], base_path]):
+                    print(
+                        f'SKIP: {rel_path} in base dir differs from environments',
+                        file=sys.stderr,
+                    )
+                    skipped_base_conflict += 1
+                    continue
+                self._remove_from_sources(rel_path, source_paths)
+                print(f'merged (already in base): {rel_path}')
+            else:
+                self._move_to_base(rel_path, source_paths)
+                print(f'merged: {rel_path}')
+            merged += 1
+
+        print(
+            f'Done: merged {merged}, '
+            f'skipped {skipped_diff} (differs between envs), '
+            f'skipped {skipped_base_conflict} (conflicts with base)'
+        )
+
+    def _validate(self) -> None:
+        if len(self.sources) < 2:
+            print('ERROR: specify at least two --source directories', file=sys.stderr)
+            sys.exit(1)
+
+        seen = set()
+        for path in self.sources:
+            if path in seen:
+                print(f'ERROR: duplicate source directory: {path}', file=sys.stderr)
+                sys.exit(1)
+            seen.add(path)
+            if not os.path.isdir(path):
+                print(f'ERROR: source directory not found: {path}', file=sys.stderr)
+                sys.exit(1)
+
+        if self.out_dir in self.sources:
+            print('ERROR: --out-dir must not be one of --source directories', file=sys.stderr)
+            sys.exit(1)
+
+        for path in self.sources:
+            if os.path.commonpath([path, self.out_dir]) == self.out_dir:
+                print(
+                    f'ERROR: source directory inside --out-dir is not supported: {path}',
+                    file=sys.stderr,
+                )
+                sys.exit(1)
+            if os.path.commonpath([path, self.out_dir]) == path:
+                print(
+                    f'ERROR: --out-dir inside source directory is not supported: {path}',
+                    file=sys.stderr,
+                )
+                sys.exit(1)
+
+        if not self.args.dry_run:
+            os.makedirs(self.out_dir, exist_ok=True)
+
+    @staticmethod
+    def _files_identical(paths: list[str]) -> bool:
+        with open(paths[0], 'rb') as first:
+            content = first.read()
+        for path in paths[1:]:
+            with open(path, 'rb') as other:
+                if other.read() != content:
+                    return False
+        return True
+
+    def _move_to_base(self, rel_path: str, source_paths: list[str]) -> None:
+        base_path = os.path.join(self.out_dir, rel_path)
+        if self.args.dry_run:
+            return
+        os.makedirs(os.path.dirname(base_path), exist_ok=True)
+        shutil.move(source_paths[0], base_path)
+        for path in source_paths[1:]:
+            os.remove(path)
+
+    def _remove_from_sources(self, rel_path: str, source_paths: list[str]) -> None:
+        if self.args.dry_run:
+            return
+        for path in source_paths:
+            os.remove(path)
+        self._remove_empty_dirs(source_paths)
+
+    @staticmethod
+    def _remove_empty_dirs(file_paths: list[str]) -> None:
+        dirs = {os.path.dirname(path) for path in file_paths}
+        for directory in sorted(dirs, key=len, reverse=True):
+            if os.path.isdir(directory) and not os.listdir(directory):
+                os.rmdir(directory)

pg_data_yaml/paths.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import glob
+import os
+
+
+def list_table_yaml_files(root_dir: str) -> dict[str, str]:
+    """Return mapping rel_path (schema/table.yaml) -> absolute path."""
+    files = {}
+    for path in glob.glob(os.path.join(root_dir, '*', '*.yaml')):
+        rel_path = os.path.relpath(path, root_dir)
+        files[rel_path] = path
+    return files

pg_data_yaml/pg.py

@@ -0,0 +1,75 @@
+import asyncio
+import datetime
+import decimal
+import json
+import signal
+import uuid
+
+import asyncpg
+
+
+def quote_ident(name: str) -> str:
+    return '"' + name.replace('"', '""') + '"'
+
+
+def quote_literal(value):
+    if value is None:
+        return 'null'
+    if isinstance(value, bool):
+        return 'true' if value else 'false'
+    if isinstance(value, (int, float)):
+        return str(value)
+    if isinstance(value, decimal.Decimal):
+        return str(value)
+    if isinstance(value, uuid.UUID):
+        return quote_literal(str(value))
+    if isinstance(value, (datetime.datetime, datetime.date, datetime.time)):
+        return quote_literal(value.isoformat())
+    if isinstance(value, str):
+        return "'" + value.replace("'", "''") + "'"
+    if isinstance(value, (dict, list)):
+        return quote_literal(json.dumps(value))
+    if isinstance(value, (bytes, memoryview)):
+        return quote_literal(bytes(value).hex())
+    raise TypeError(f'Unknown type for quote value: {value!r}')
+
+
+class Pg:
+    con: asyncpg.Connection
+
+    def __init__(self, args):
+        self.args = args
+
+    async def init(self):
+        self.con = await asyncpg.connect(
+            database=self.args.dbname,
+            user=self.args.user,
+            password=self.args.password,
+            host=self.args.host,
+            port=self.args.port,
+            statement_cache_size=0,
+        )
+
+    async def fetch(self, query: str, *params) -> list[dict]:
+        query_task = asyncio.create_task(self.con.fetch(query, *params))
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGINT, query_task.cancel)
+        try:
+            rows = await query_task
+            return [{key: row[key] for key in row.keys()} for row in rows]
+        except asyncio.CancelledError:
+            await asyncio.sleep(0.5)
+            return []
+        finally:
+            loop.remove_signal_handler(signal.SIGINT)
+
+    async def execute(self, query: str, *params) -> None:
+        query_task = asyncio.create_task(self.con.execute(query, *params))
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGINT, query_task.cancel)
+        try:
+            await query_task
+        except asyncio.CancelledError:
+            await asyncio.sleep(0.5)
+        finally:
+            loop.remove_signal_handler(signal.SIGINT)