endorphin-tabularium-docs/scripts/documents-client.py

"""Upload documents through the database API; credentials stay in a private file."""
import argparse
import json
import os
from pathlib import Path
import stat
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid


class DocumentClientError(ValueError):
    pass


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        # Never forward a bearer credential to a redirected host or protocol.
        return None


def server_url(value):
    parsed = urllib.parse.urlsplit(value)
    if (parsed.username or parsed.password or parsed.query or parsed.fragment
            or parsed.path not in ('', '/') or not parsed.hostname):
        raise ValueError('Provide a server origin without credentials, path, query or fragment')
    if parsed.scheme != 'https' and not (
            parsed.scheme == 'http' and parsed.hostname in ('localhost', '127.0.0.1', '::1')):
        raise ValueError('Use HTTPS or an HTTP loopback connection through a trusted SSH tunnel')
    return value.rstrip('/')


def read_token(filename):
    path = Path(filename)
    if stat.S_IMODE(path.stat().st_mode) & 0o077:
        raise ValueError('API key file must be private (chmod 600)')
    value = path.read_text().strip()
    if value.startswith('{'):
        value = json.loads(value)['token']
    if not isinstance(value, str) or not value.startswith('tdoc_') or len(value) != 48:
        raise ValueError('Invalid document API key file')
    return value


def request(base, token, method, path, payload=None, request_id=None):
    data = None if payload is None else json.dumps(payload, ensure_ascii=False).encode()
    if data is not None and len(data) > 8 * 1024 * 1024:
        raise ValueError('Request exceeds 8 MiB; split into smaller atomic batches')
    headers = {'Authorization': 'Bearer ' + token, 'Accept': 'application/json'}
    if data is not None:
        headers['Content-Type'] = 'application/json'
        headers['Idempotency-Key'] = request_id
    opener = urllib.request.build_opener(NoRedirect())
    for attempt in range(2):
        req = urllib.request.Request(base + path, data=data, method=method, headers=headers)
        try:
            with opener.open(req, timeout=30) as response:
                return json.load(response)
        except urllib.error.HTTPError as error:
            if error.code in (502, 503, 504) and attempt == 0:
                error.close()
                time.sleep(1)
                continue
            status = error.code
            error.close()
            explanations = {400: 'Invalid document or pagination', 401: 'Key missing, expired or revoked',
                            403: 'Read-only key', 404: 'Document or revision not found',
                            409: 'Revision or request-id conflict; read current state before retrying',
                            413: 'Request too large', 503: 'Database unavailable'}
            raise DocumentClientError(f'HTTP {status}: ' + explanations.get(status, 'Request failed')) from None
        except (urllib.error.URLError, TimeoutError):
            if attempt == 0:
                time.sleep(1)
                continue
            raise DocumentClientError('Connection failed; retry using the same request id and unchanged input') from None


def document_batches(documents, max_bytes=8 * 1024 * 1024):
    if not isinstance(documents, list) or not documents:
        raise DocumentClientError('Provide a nonempty documents array')
    batches, chunk, size, seen = [], [], 20, set()
    for document in documents:
        if not isinstance(document, dict) or not isinstance(document.get('sourceSlug'), str):
            raise DocumentClientError('Every document needs a sourceSlug')
        if document['sourceSlug'] in seen:
            raise DocumentClientError('Duplicate sourceSlug in upload')
        seen.add(document['sourceSlug'])
        encoded = len(json.dumps(document, ensure_ascii=False).encode()) + 2
        if encoded + 20 > max_bytes:
            raise DocumentClientError('A document exceeds the request limit')
        if chunk and (len(chunk) == 100 or size + encoded > max_bytes):
            batches.append(chunk)
            chunk, size = [], 20
        chunk.append(document)
        size += encoded
    batches.append(chunk)
    return batches


def main():
    os.umask(0o077)
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--url', default=os.environ.get('TABULARIUM_DOCS_URL'))
    parser.add_argument('--key-file', default=os.environ.get('TABULARIUM_DOCS_KEY_FILE'))
    sub = parser.add_subparsers(dest='command', required=True)
    listing = sub.add_parser('list')
    listing.add_argument('--limit', type=int, default=50)
    listing.add_argument('--offset', type=int, default=0)
    reading = sub.add_parser('get')
    reading.add_argument('slug')
    mode = reading.add_mutually_exclusive_group()
    mode.add_argument('--history', action='store_true')
    mode.add_argument('--revision', type=int)
    reading.add_argument('--limit', type=int, default=50)
    reading.add_argument('--offset', type=int, default=0)
    sections = sub.add_parser('outline', help='Read chapter titles and internal summaries without bodies')
    sections.add_argument('slug')
    sections.add_argument('--revision', type=int)
    section = sub.add_parser('sections', help='Read selected chapter bodies')
    section.add_argument('slug')
    section.add_argument('ids', nargs='+')
    section.add_argument('--revision', type=int)
    editing = sub.add_parser('edit-sections', help='Atomically apply chapter operations from private JSON')
    editing.add_argument('slug')
    editing.add_argument('file')
    editing.add_argument('--request-id')
    saving = sub.add_parser('save')
    saving.add_argument('file', help='Private JSON file or - for stdin; one document or {documents: [...]}')
    saving.add_argument('--request-id', help='Reuse this id only with the same input when retrying')
    args = parser.parse_args()
    if not args.url or not args.key_file:
        parser.error('Set --url and --key-file or TABULARIUM_DOCS_URL and TABULARIUM_DOCS_KEY_FILE')
    try:
        base = server_url(args.url)
        token = read_token(args.key_file)
        if args.command == 'list':
            result = request(base, token, 'GET', '/api/documents?' + urllib.parse.urlencode(
                {'limit': args.limit, 'offset': args.offset}))
        elif args.command == 'get':
            path = '/api/documents/item/' + urllib.parse.quote(args.slug, safe='/')
            if args.history:
                path += '?' + urllib.parse.urlencode({'history': 1, 'limit': args.limit, 'offset': args.offset})
            elif args.revision is not None:
                path += '?' + urllib.parse.urlencode({'revision': args.revision})
            result = request(base, token, 'GET', path)
        elif args.command in ('outline', 'sections', 'edit-sections'):
            path = '/api/documents/sections/' + urllib.parse.quote(args.slug, safe='/')
            if args.command == 'edit-sections':
                payload = json.load(sys.stdin) if args.file == '-' else json.loads(Path(args.file).read_text())
                request_id = args.request_id or str(uuid.uuid4())
                if not request_id or len(request_id) > 100 or any(c not in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._:-' for c in request_id):
                    raise ValueError('Invalid request id')
                print('Request id: ' + request_id, file=sys.stderr)
                result = request(base, token, 'POST', path, payload, request_id)
            else:
                query = []
                if args.revision is not None:
                    query.append(('revision', args.revision))
                if args.command == 'sections':
                    query.extend(('section', section_id) for section_id in args.ids)
                if query:
                    path += '?' + urllib.parse.urlencode(query)
                result = request(base, token, 'GET', path)
        else:
            payload = json.load(sys.stdin) if args.file == '-' else json.loads(Path(args.file).read_text())
            request_id = args.request_id or str(uuid.uuid4())
            if not request_id or len(request_id) > 100 or any(c not in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._:-' for c in request_id):
                raise ValueError('Invalid request id')
            print('Request id: ' + request_id, file=sys.stderr)
            if isinstance(payload, dict) and 'documents' in payload:
                if set(payload) != {'documents'}:
                    raise DocumentClientError('Batch payload must contain only documents')
                chunks = document_batches(payload['documents'])
                receipts, replayed = [], 0
                for index, chunk in enumerate(chunks, 1):
                    batch_id = f'{request_id}:{index}'
                    saved = request(base, token, 'POST', '/api/documents/batch', {'documents': chunk}, batch_id)
                    receipts.extend(saved['documents'])
                    replayed += int(saved['replayed'])
                    print(f'Batch {index}/{len(chunks)} committed; {len(receipts)} documents total; request id: {batch_id}', file=sys.stderr)
                result = {'documents': receipts, 'batches': len(chunks), 'replayedBatches': replayed, 'requestId': request_id}
            else:
                result = request(base, token, 'POST', '/api/documents', payload, request_id)
        print(json.dumps(result, ensure_ascii=False, indent=2))
    except DocumentClientError as error:
        raise SystemExit(str(error)) from None
    except (ValueError, OSError, KeyError, TypeError):
        # Never print exception objects that could contain credentials or document data.
        raise SystemExit('Document request failed. Check the URL, private key file, JSON schema and server response; credentials omitted.') from None


if __name__ == '__main__':
    main()