197 lines
9.9 KiB
Python
197 lines
9.9 KiB
Python
"""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()
|