datasette_pytables/__init__.py

   1 import tables
   2 import datasette_connectors as dc
   3 from .utils import parse_sql
   4
   5
   6 class PyTablesConnection(dc.Connection):
   7     def __init__(self, path, connector):
   8         super().__init__(path, connector)
   9         self.h5file = tables.open_file(path)
  10
  11
  12 class PyTablesConnector(dc.Connector):
  13     connector_type = 'pytables'
  14     connection_class = PyTablesConnection
  15
  16     operators = {
  17         'eq': '==',
  18         'neq': '!=',
  19         'gt': '>',
  20         'gte': '>=',
  21         'lt': '<',
  22         'lte': '<=',
  23         'and': '&',
  24         'or': '|',
  25         'binary_and': '&',
  26         'binary_or': '|',
  27     }
  28
  29     def _serialize_table_name(self, table_name):
  30         return table_name.replace('/', '%')
  31
  32     def _deserialize_table_name(self, table_name):
  33         return table_name.replace('%', '/')
  34
  35     def table_names(self):
  36         return [
  37             self._serialize_table_name(node._v_pathname)
  38             for node in self.conn.h5file
  39             if not(isinstance(node, tables.group.Group))
  40         ]
  41
  42     def table_count(self, table_name):
  43         table = self.conn.h5file.get_node(self._deserialize_table_name(table_name))
  44         return int(table.nrows)
  45
  46     def table_info(self, table_name):
  47         table = self.conn.h5file.get_node(self._deserialize_table_name(table_name))
  48         columns = [
  49             {
  50                 'name': 'value',
  51                 'type': table.dtype.name,
  52             }
  53         ]
  54         if isinstance(table, tables.table.Table):
  55             columns = [
  56                 {
  57                     'name': colname,
  58                     'type': table.coltypes[colname],
  59                 }
  60                 for colname in table.colnames
  61             ]
  62
  63         return [
  64             {
  65                 'cid': cid,
  66                 'name': column['name'],
  67                 'type': column['type'],
  68                 'notnull': True,
  69                 'default_value': None,
  70                 'is_pk': False,
  71             }
  72             for cid, column in enumerate(columns)
  73         ]
  74
  75     def hidden_table_names(self):
  76         return []
  77
  78     def detect_spatialite(self):
  79         return False
  80
  81     def view_names(self):
  82         return []
  83
  84     def detect_fts(self, table_name):
  85         return False
  86
  87     def foreign_keys(self, table_name):
  88         return []
  89
  90     def table_exists(self, table_name):
  91         try:
  92             self.conn.h5file.get_node(self._deserialize_table_name(table_name))
  93             return True
  94         except:
  95             return False
  96
  97     def table_definition(self, table_type, table_name):
  98         table_name = self._deserialize_table_name(table_name)
  99         table = self.conn.h5file.get_node(table_name)
 100         colnames = ['value']
 101         if isinstance(table, tables.table.Table):
 102             colnames = table.colnames
 103
 104         return 'CREATE TABLE {} ({})'.format(
 105             table_name,
 106             ', '.join(colnames),
 107         )
 108
 109     def indices_definition(self, table_name):
 110         return []
 111
 112     def execute(
 113         self,
 114         sql,
 115         params=None,
 116         truncate=False,
 117         custom_time_limit=None,
 118         page_size=None,
 119         log_sql_errors=True,
 120     ):
 121         results = []
 122         truncated = False
 123         description = ()
 124
 125         # Some Datasette queries uses glob operand, not supported by Pytables
 126         if ' glob ' in sql:
 127             return results, truncated, description
 128
 129         parsed_sql = parse_sql(sql, params)
 130
 131         while isinstance(parsed_sql['from'], dict):
 132             # Pytables does not support subqueries
 133             parsed_sql['from'] = parsed_sql['from']['value']['from']
 134
 135         table = self.conn.h5file.get_node(self._deserialize_table_name(parsed_sql['from']))
 136         table_rows = []
 137         fields = parsed_sql['select']
 138         colnames = ['value']
 139         if type(table) is tables.table.Table:
 140             colnames = table.colnames
 141
 142         query = ''
 143         start = 0
 144         end = table.nrows
 145
 146         def _get_field_type(field):
 147             coltype = table.dtype.name
 148             if type(table) is tables.table.Table:
 149                 coltype = table.coltypes[field]
 150             return coltype
 151
 152         # Use 'where' statement or get all the rows
 153         def _cast_param(field, pname):
 154             # Cast value to the column type
 155             coltype = _get_field_type(field)
 156             fcast = None
 157             if coltype == 'string':
 158                 fcast = str
 159             elif coltype.startswith('int'):
 160                 fcast = int
 161             elif coltype.startswith('float'):
 162                 fcast = float
 163             if fcast:
 164                 params[pname] = fcast(params[pname])
 165
 166         def _translate_where(where):
 167             # Translate SQL to PyTables expression
 168             nonlocal start, end
 169             expr = ''
 170             operator = list(where)[0]
 171
 172             if operator in ['and', 'or']:
 173                 subexpr = [_translate_where(e) for e in where[operator]]
 174                 subexpr = filter(lambda e: e, subexpr)
 175                 subexpr = ["({})".format(e) for e in subexpr]
 176                 expr = " {} ".format(self.operators[operator]).join(subexpr)
 177             elif operator == 'exists':
 178                 pass
 179             elif where == {'eq': ['rowid', 'p0']}:
 180                 start = int(params['p0'])
 181                 end = start + 1
 182             elif where == {'gt': ['rowid', 'p0']}:
 183                 start = int(params['p0']) + 1
 184             else:
 185                 left, right = where[operator]
 186
 187                 if isinstance(left, dict):
 188                     left = "(" + _translate_where(left) + ")"
 189                 elif left in params:
 190                     _cast_param(right, left)
 191
 192                 if isinstance(right, dict):
 193                     right = "(" + _translate_where(right) + ")"
 194                 elif right in params:
 195                     _cast_param(left, right)
 196
 197                 expr = "{left} {operator} {right}".format(
 198                     left=left,
 199                     operator=self.operators.get(operator, operator),
 200                     right=right,
 201                 )
 202
 203             return expr
 204
 205         if 'where' in parsed_sql:
 206             if type(parsed_sql['where']) is dict:
 207                 query = _translate_where(parsed_sql['where'])
 208             else:
 209                 query = parsed_sql['where']
 210
 211         # Sort by column
 212         orderby = ''
 213         if 'orderby' in parsed_sql:
 214             orderby = parsed_sql['orderby']
 215             if type(orderby) is list:
 216                 orderby = orderby[0]
 217             orderby = orderby['value']
 218             if orderby == 'rowid':
 219                 orderby = ''
 220
 221         # Limit number of rows
 222         limit = None
 223         if 'limit' in parsed_sql:
 224             limit = int(parsed_sql['limit'])
 225
 226         # Offset
 227         offset = None
 228         if 'offset' in parsed_sql:
 229             offset = int(parsed_sql['offset'])
 230
 231         # Truncate if needed
 232         if page_size and max_returned_rows and truncate:
 233             if max_returned_rows == page_size:
 234                 max_returned_rows += 1
 235
 236         # Execute query
 237         if query:
 238             table_rows = table.where(query, params, start, end)
 239         elif orderby:
 240             table_rows = table.itersorted(orderby, start=start, stop=end)
 241         else:
 242             table_rows = table.iterrows(start, end)
 243
 244         # Prepare rows
 245         def normalize_field_value(value):
 246             if type(value) is bytes:
 247                 return value.decode('utf-8')
 248             elif not type(value) in (int, float, complex):
 249                 return str(value)
 250             else:
 251                 return value
 252
 253         def make_get_rowid():
 254             if type(table) is tables.table.Table:
 255                 def get_rowid(row):
 256                     return int(row.nrow)
 257             else:
 258                 rowid = start - 1
 259                 def get_rowid(row):
 260                     nonlocal rowid
 261                     rowid += 1
 262                     return rowid
 263             return get_rowid
 264
 265         def make_get_row_value():
 266             if type(table) is tables.table.Table:
 267                 def get_row_value(row, field):
 268                     return row[field]
 269             else:
 270                 def get_row_value(row, field):
 271                     return row
 272             return get_row_value
 273
 274         # Get results
 275         get_rowid = make_get_rowid()
 276         get_row_value = make_get_row_value()
 277         if offset:
 278             table_rows = table_rows[offset:]
 279         count = 0
 280         for table_row in table_rows:
 281             count += 1
 282             if limit is not None and count > limit:
 283                 break
 284             if truncate and max_returned_rows and count > max_returned_rows:
 285                 truncated = True
 286                 break
 287             row = {}
 288             for field in fields:
 289                 field_name = field
 290                 if isinstance(field, dict):
 291                     field_name = field['value']
 292                 if isinstance(field_name, dict) and 'distinct' in field_name:
 293                     field_name = field_name['distinct']
 294                 if field_name == 'rowid':
 295                     row['rowid'] = get_rowid(table_row)
 296                 elif field_name == '*':
 297                     for col in colnames:
 298                         row[col] = normalize_field_value(get_row_value(table_row, col))
 299                 elif isinstance(field_name, dict):
 300                     if field_name.get('count') == '*':
 301                         row['count(*)'] = int(table.nrows)
 302                     elif field_name.get('json_type'):
 303                         field_name = field_name.get('json_type')
 304                         row['json_type(' + field_name + ')'] = _get_field_type(field_name)
 305                     else:
 306                         raise Exception("Function not recognized")
 307                 else:
 308                     row[field_name] = normalize_field_value(get_row_value(table_row, field_name))
 309             results.append(row)
 310
 311         # Prepare query description
 312         for field in [f['value'] if isinstance(f, dict) else f for f in fields]:
 313             if field == '*':
 314                 for col in colnames:
 315                     description += ((col,),)
 316             else:
 317                 description += ((field,),)
 318
 319         return results, truncated, description