|
| 1 | +import getpass |
| 2 | +import json |
| 3 | +import os |
| 4 | +import os.path |
| 5 | +import urllib.parse |
| 6 | + |
| 7 | +import pyarrow as pa |
| 8 | +import pyarrow.ipc |
| 9 | +import requests |
| 10 | + |
| 11 | + |
| 12 | +class Client(): |
| 13 | + def __init__(self, |
| 14 | + base_url=os.environ.get('SUPER_DB', 'http://localhost:9867'), |
| 15 | + config_dir=os.path.expanduser('~/.super')): |
| 16 | + self.base_url = base_url.rstrip('/') |
| 17 | + self.session = requests.Session() |
| 18 | + self.session.headers.update({'Accept': 'application/vnd.apache.arrow.stream'}) |
| 19 | + token = self.__get_auth_token(config_dir) |
| 20 | + if token is not None: |
| 21 | + self.session.headers.update({'Authorization': 'Bearer ' + token}) |
| 22 | + |
| 23 | + def __get_auth_token(self, config_dir): |
| 24 | + creds_path = os.path.join(config_dir, 'credentials.json') |
| 25 | + try: |
| 26 | + with open(creds_path) as f: |
| 27 | + data = f.read() |
| 28 | + except FileNotFoundError: |
| 29 | + return None |
| 30 | + creds = json.loads(data) |
| 31 | + if self.base_url in creds['services']: |
| 32 | + return creds['services'][self.base_url]['access'] |
| 33 | + return None |
| 34 | + |
| 35 | + def create_pool(self, name, layout={'order': 'desc', 'keys': [['ts']]}, |
| 36 | + thresh=0): |
| 37 | + r = self.session.post(self.base_url + '/pool', json={ |
| 38 | + 'name': name, |
| 39 | + 'layout': layout, |
| 40 | + 'thresh': thresh, |
| 41 | + }) |
| 42 | + self.__raise_for_status(r) |
| 43 | + |
| 44 | + def load(self, pool_name_or_id, data, branch_name='main', |
| 45 | + commit_author=getpass.getuser(), commit_body='', |
| 46 | + mime_type=None): |
| 47 | + pool = urllib.parse.quote(pool_name_or_id, safe='') |
| 48 | + branch = urllib.parse.quote(branch_name, safe='') |
| 49 | + url = self.base_url + '/pool/' + pool + '/branch/' + branch |
| 50 | + commit_message = {'author': commit_author, 'body': commit_body} |
| 51 | + headers = {'SuperDB-Commit': json.dumps(commit_message)} |
| 52 | + if mime_type is not None: |
| 53 | + headers['Content-Type'] = mime_type |
| 54 | + r = self.session.post(url, headers=headers, data=data) |
| 55 | + self.__raise_for_status(r) |
| 56 | + |
| 57 | + def delete_pool(self, pool_name_or_id): |
| 58 | + pool = urllib.parse.quote(pool_name_or_id, safe='') |
| 59 | + r = self.session.delete(self.base_url + '/pool/' + pool) |
| 60 | + self.__raise_for_status(r) |
| 61 | + |
| 62 | + def query(self, query, safe=True): |
| 63 | + if safe: |
| 64 | + # Pre-flight: verify all top-level values are records of a single |
| 65 | + # type. Arrow requires top-level records and silently truncates on |
| 66 | + # type changes, so we detect both problems before issuing the real |
| 67 | + # query. |
| 68 | + safety_r = self.query_raw( |
| 69 | + query + ' | union(typeof(this)) by kind(this)', |
| 70 | + headers={'Accept': 'application/x-ndjson'}, |
| 71 | + ) |
| 72 | + rows = [ |
| 73 | + json.loads(line) |
| 74 | + for line in safety_r.iter_lines(decode_unicode=True) |
| 75 | + if line |
| 76 | + ] |
| 77 | + if rows: |
| 78 | + if any(row['kind'] != 'record' for row in rows): |
| 79 | + kinds = sorted({row['kind'] for row in rows}) |
| 80 | + raise NonRecordError( |
| 81 | + f"Query result contains non-record values " |
| 82 | + f"(kind: {', '.join(repr(k) for k in kinds)}). " |
| 83 | + f"Arrow requires top-level records.", |
| 84 | + kinds, |
| 85 | + ) |
| 86 | + type_count = len(rows[0]['union']) |
| 87 | + if type_count > 1: |
| 88 | + raise MixedTypesError( |
| 89 | + f'Query result contains {type_count} distinct types; results ' |
| 90 | + f'would be silently truncated. Use \'| blend\' to merge types ' |
| 91 | + f'into one, or pass safe=False to skip this check and accept ' |
| 92 | + f'partial results.', |
| 93 | + type_count, |
| 94 | + ) |
| 95 | + r = self.query_raw(query) |
| 96 | + try: |
| 97 | + reader = pa.ipc.open_stream(r.raw) |
| 98 | + except pa.lib.ArrowInvalid as e: |
| 99 | + # An empty response body (no schema) means either the pool has no |
| 100 | + # data or the data contains a type the Arrow encoder can't handle |
| 101 | + # (e.g. an empty record). Both cases are indistinguishable at the |
| 102 | + # HTTP level when streaming, so both are silently treated as an |
| 103 | + # empty result. Any other ArrowInvalid (wrong format, mid-stream |
| 104 | + # corruption, etc.) is re-raised. |
| 105 | + if 'null or length 0' in str(e): |
| 106 | + return |
| 107 | + raise |
| 108 | + for batch in reader: |
| 109 | + yield from batch.to_pylist(maps_as_pydicts='strict') |
| 110 | + |
| 111 | + def query_raw(self, query, headers=None): |
| 112 | + r = self.session.post(self.base_url + '/query', headers=headers, |
| 113 | + json={'query': query}, stream=True) |
| 114 | + self.__raise_for_status(r) |
| 115 | + r.raw.decode_content = True |
| 116 | + return r |
| 117 | + |
| 118 | + @staticmethod |
| 119 | + def __raise_for_status(response): |
| 120 | + if response.status_code >= 400: |
| 121 | + try: |
| 122 | + error = response.json()['error'] |
| 123 | + except Exception: |
| 124 | + response.raise_for_status() |
| 125 | + else: |
| 126 | + raise RequestError(error, response) |
| 127 | + |
| 128 | + |
| 129 | +class RequestError(Exception): |
| 130 | + """Raised by Client methods when an HTTP request fails.""" |
| 131 | + def __init__(self, message, response): |
| 132 | + super(RequestError, self).__init__(message) |
| 133 | + self.response = response |
| 134 | + |
| 135 | + |
| 136 | +class MixedTypesError(Exception): |
| 137 | + """Raised by query() when the result contains more than one distinct type.""" |
| 138 | + def __init__(self, message, type_count): |
| 139 | + super().__init__(message) |
| 140 | + self.type_count = type_count |
| 141 | + |
| 142 | + |
| 143 | +class NonRecordError(Exception): |
| 144 | + """Raised by query() when the result contains non-record top-level values.""" |
| 145 | + def __init__(self, message, kinds): |
| 146 | + super().__init__(message) |
| 147 | + self.kinds = kinds |
| 148 | + |
| 149 | + |
| 150 | +if __name__ == '__main__': |
| 151 | + import argparse |
| 152 | + import pprint |
| 153 | + |
| 154 | + parser = argparse.ArgumentParser( |
| 155 | + description='Query default SuperDB service and print results.', |
| 156 | + formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| 157 | + parser.add_argument('query') |
| 158 | + args = parser.parse_args() |
| 159 | + |
| 160 | + c = Client() |
| 161 | + for record in c.query(args.query): |
| 162 | + pprint.pprint(record) |
0 commit comments