"""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()