localtools-postgres-migration-checker 1.0.0__py3-none-any.whl

localtools_postgres_migration_checker/__main__.py

@@ -0,0 +1,257 @@
+import argparse
+import json
+import os
+import re
+import select
+import shutil
+import socket
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from pathlib import Path
+
+import multipart
+from . import engine
+from .localtools import HTTPError, Tool, main as run
+from .limits import limit_memory
+
+USAGE = '''localtools-postgres-migration-checker: offline PostgreSQL migration safety estimates.
+
+Usage:
+  localtools-postgres-migration-checker [serve] [--port N] [--no-open]
+  localtools-postgres-migration-checker check FILE... [--pg 11..18] [--transaction auto|always|never]
+      [--format human|json|markdown|sarif] [--fail-on high|medium|low|never] [-o FILE]
+  localtools-postgres-migration-checker update | uninstall | --version | --help
+
+FILE '-' reads UTF-8 SQL from stdin. Files run in supplied order. No SQL is executed.
+Limits: 100 files, 4 MiB aggregate SQL, 5000 statements, 16 MiB report, 30 seconds, 2 concurrent checks.
+POST /api/check: multipart sql, name, file (repeatable), pg, transaction, format.
+'''
+UI = Path(__file__).with_name('ui')
+if not UI.is_dir():
+    UI = Path(__file__).parents[1] / 'ui'
+DEADLINE = 30
+MAX_REPORT = 16 << 20
+SLOTS = threading.BoundedSemaphore(2)
+PREFIX = 'lt-pgmc-'
+
+
+def sweep():
+    for p in Path(tempfile.gettempdir()).glob(PREFIX + '*'):
+        try:
+            if (re.fullmatch(r'lt-pgmc-[a-z0-9_]{8}', p.name) and not p.is_symlink()
+                    and not (getattr(p.lstat(), 'st_file_attributes', 0) & 1024)
+                    and p.is_dir() and time.time() - p.stat().st_mtime > 3600):
+                shutil.rmtree(p)
+        except OSError:
+            pass
+
+
+def decode(data, name):
+    if b'\x00' in data:
+        raise HTTPError(400, f'{name}: NUL bytes are not valid SQL text')
+    try:
+        return engine.newlines(data.decode('utf-8-sig'))
+    except UnicodeDecodeError:
+        raise HTTPError(400, f'{name}: expected UTF-8 SQL') from None
+
+
+def read_form(req):
+    ctype, opts = multipart.parse_options_header(req.headers.get('Content-Type', ''))
+    if ctype != 'multipart/form-data' or 'boundary' not in opts:
+        raise HTTPError(400, 'expected multipart/form-data')
+    try:
+        size = int(req.headers.get('Content-Length', '0'))
+    except ValueError:
+        raise HTTPError(400, 'invalid Content-Length') from None
+    if size <= 0:
+        raise HTTPError(400, 'send SQL or a .sql file')
+    if size > engine.MAX_BYTES + (256 << 10):
+        raise HTTPError(413, 'limit: 4 MiB aggregate SQL plus 256 KiB form overhead')
+    parts, files, fields = [], [], {}
+    try:
+        for p in multipart.MultipartParser(req.rfile, opts['boundary'], size,
+                part_limit=engine.MAX_FILES + 8, memory_limit=1 << 20, partsize_limit=engine.MAX_BYTES):
+            parts.append(p)
+            if p.filename:
+                if p.name != 'file':
+                    raise HTTPError(400, 'uploaded SQL must use the file field')
+                name = p.filename.replace('\\', '/').split('/')[-1][:200]
+                files.append((name, decode(p.raw, name)))
+            else:
+                if p.name not in ('sql', 'name', 'pg', 'transaction', 'format') or p.name in fields:
+                    raise HTTPError(400, 'unknown or repeated form field')
+                fields[p.name] = decode(p.raw, p.name)
+        if fields.get('sql', '').strip():
+            files = [(fields.get('name', 'pasted.sql')[:200], fields['sql'])]
+        return files, fields
+    except multipart.ParserLimitReached:
+        raise HTTPError(413, 'upload exceeds file, field, or 4 MiB SQL limit') from None
+    except multipart.MultipartError:
+        raise HTTPError(400, 'malformed multipart upload or field limit exceeded') from None
+    finally:
+        for p in parts:
+            p.close()
+
+
+def disconnected(req):
+    if req is None:
+        return False
+    try:
+        ready, _, _ = select.select([req.connection], [], [], 0)
+        return bool(ready) and req.connection.recv(1, socket.MSG_PEEK) == b''
+    except OSError:
+        return True
+
+
+def worker(args):
+    # EOF is reliable on all platforms and kills the worker if its parent dies.
+    def parent_watch():
+        os.read(sys.stdin.fileno(), 1)
+        os._exit(1)
+    threading.Thread(target=parent_watch, daemon=True).start()
+    path = Path(args[0])
+    try:
+        memory_job = limit_memory()
+        job = json.loads(path.read_text(encoding='utf-8'))
+        rep = engine.check(job['files'], pg=job['pg'], txn_mode=job['transaction'])
+        fmt = job['format']
+        output = (engine.markdown(rep) if fmt == 'markdown' else engine.human(rep) if fmt == 'human'
+                  else engine.dumps(engine.sarif(rep, '1.0.0')) if fmt == 'sarif' else engine.dumps(rep))
+        if len(output.encode('utf-8')) > MAX_REPORT:
+            raise ValueError('report exceeds 16 MiB; split the migration')
+        res = {'output': output, 'failed': engine.failing(rep, job.get('fail_on', 'high'))}
+    except MemoryError:
+        res = {'error': 'SQL exceeds the 512 MiB parser worker memory budget; split it', 'code': 422}
+    except (ValueError, RecursionError) as e:
+        res = {'error': str(e) if not isinstance(e, RecursionError) else 'SQL is nested too deeply', 'code': 422}
+    except Exception:
+        res = {'error': 'unsupported SQL structure; split the migration or report this input', 'code': 422}
+    path.with_suffix('.out').write_text(json.dumps(res, ensure_ascii=False), encoding='utf-8')
+
+
+def execute(files, fields, req=None):
+    try:
+        pg = int(fields.get('pg') or engine.PG_DEFAULT)
+    except ValueError:
+        raise HTTPError(400, 'pg must be an integer from 11 to 18') from None
+    txn = fields.get('transaction', 'auto')
+    fmt = fields.get('format', 'json')
+    if not engine.PG_MIN <= pg <= engine.PG_MAX or txn not in engine.TXN_MODES:
+        raise HTTPError(400, 'pg must be 11-18; transaction must be auto, always or never')
+    if fmt not in ('json', 'markdown', 'sarif', 'human'):
+        raise HTTPError(400, 'format must be json, markdown, sarif or human')
+    if not files or not any(t.strip() for _, t in files):
+        raise HTTPError(400, 'paste SQL or upload a non-empty .sql file')
+    if len(files) > engine.MAX_FILES or sum(len(t.encode('utf-8')) for _, t in files) > engine.MAX_BYTES:
+        raise HTTPError(413, 'limit: 100 files and 4 MiB aggregate UTF-8 SQL')
+    if not SLOTS.acquire(blocking=False):
+        raise HTTPError(503, 'busy: two checks are running; retry when one finishes')
+    try:
+        with tempfile.TemporaryDirectory(prefix=PREFIX) as d:
+            path = Path(d, 'job.json')
+            path.write_text(json.dumps({'files': files, 'pg': pg, 'transaction': txn, 'format': fmt,
+                                        'fail_on': fields.get('fail_on', 'high')}), encoding='utf-8')
+            cmd = [sys.executable, *(['-I'] if sys.flags.isolated else []), '-m', __package__, '_worker', str(path)]
+            p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+                                 creationflags=getattr(subprocess, 'CREATE_NO_WINDOW', 0))
+            end = time.monotonic() + DEADLINE
+            try:
+                while p.poll() is None:
+                    if time.monotonic() > end:
+                        raise HTTPError(422, 'check exceeded 30 seconds; split the migration')
+                    if disconnected(req):
+                        raise HTTPError(499, 'client disconnected')
+                    try:
+                        p.wait(.1)
+                    except subprocess.TimeoutExpired:
+                        pass
+            finally:
+                if p.poll() is None:
+                    p.kill()
+                p.wait()
+                p.stdin.close()
+            out = path.with_suffix('.out')
+            if not out.is_file() or out.stat().st_size > MAX_REPORT * 2 + 1024:
+                raise HTTPError(422, 'parser could not process this SQL; nesting may be too deep')
+            res = json.loads(out.read_text(encoding='utf-8'))
+    finally:
+        SLOTS.release()
+    if 'error' in res:
+        raise HTTPError(res['code'], res['error'])
+    return res
+
+
+def api_check(req):
+    files, fields = read_form(req)
+    res = execute(files, fields, req)
+    fmt = fields.get('format', 'json')
+    types = {'json': 'application/json', 'sarif': 'application/sarif+json', 'markdown': 'text/markdown', 'human': 'text/plain'}
+    req.send_body(200, res['output'].encode('utf-8'), types[fmt] + '; charset=utf-8')
+
+
+def cli_check(args):
+    p = argparse.ArgumentParser(prog='localtools-postgres-migration-checker check')
+    p.add_argument('files', nargs='+')
+    p.add_argument('--pg', type=int, default=engine.PG_DEFAULT)
+    p.add_argument('--transaction', choices=engine.TXN_MODES, default='auto')
+    p.add_argument('--format', choices=('human', 'json', 'markdown', 'sarif'), default='human')
+    p.add_argument('--fail-on', choices=('high', 'medium', 'low', 'never'), default='high')
+    p.add_argument('-o', '--output')
+    a = p.parse_args(args)
+    if len(a.files) > engine.MAX_FILES:
+        raise ValueError('limit: 100 files')
+    files, total = [], 0
+    for name in a.files:
+        try:
+            stream = sys.stdin.buffer if name == '-' else open(name, 'rb')
+            try:
+                data = stream.read(engine.MAX_BYTES - total + 1)
+            finally:
+                if name != '-':
+                    stream.close()
+            total += len(data)
+            if total > engine.MAX_BYTES:
+                raise ValueError('limit: 4 MiB aggregate UTF-8 SQL')
+            files.append((Path(name).name if name != '-' else 'stdin.sql', decode(data, name)))
+        except OSError:
+            raise ValueError(f'cannot read {Path(name).name}') from None
+    res = execute(files, vars(a))
+    output = res['output']
+    if a.output:
+        atomic_write(Path(a.output), output)
+    else:
+        sys.stdout.reconfigure(encoding='utf-8', newline='')
+        print(output, end='')
+    return int(res['failed'])
+
+
+def atomic_write(path, output):
+    temp = None
+    try:
+        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', newline='', dir=path.parent,
+                                         prefix='.pgmc-', delete=False) as f:
+            temp = Path(f.name)
+            f.write(output)
+        os.replace(temp, path)
+    except OSError:
+        raise ValueError('cannot write report; check the output folder and permissions') from None
+    finally:
+        if temp is not None:
+            temp.unlink(missing_ok=True)
+
+
+tool = Tool(slug='postgres-migration-checker', port=47116, usage=USAGE, ui=UI,
+    routes={'POST /api/check': api_check}, cli={'check': cli_check, '_worker': worker})
+
+
+def main():
+    if len(sys.argv) < 2 or sys.argv[1] != '_worker':
+        sweep()
+    sys.exit(run(tool))
+
+
+if __name__ == '__main__':
+    main()