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         parsed_sql = parse_sql(sql, params)
 126
 127         while isinstance(parsed_sql['from'], dict):
 128             # Pytables does not support subqueries
 129             parsed_sql['from'] = parsed_sql['from']['value']['from']
 130
 131         table = self.conn.h5file.get_node(self._deserialize_table_name(parsed_sql['from']))
 132         table_rows = []
 133         fields = parsed_sql['select']
 134         colnames = ['value']
 135         if type(table) is tables.table.Table:
 136             colnames = table.colnames
 137
 138         query = ''
 139         start = 0
 140         end = table.nrows
 141
 142         def _get_field_type(field):
 143             coltype = table.dtype.name
 144             if type(table) is tables.table.Table:
 145                 coltype = table.coltypes[field]
 146             return coltype
 147
 148         # Use 'where' statement or get all the rows
 149         def _cast_param(field, pname):
 150             # Cast value to the column type
 151             coltype = _get_field_type(field)
 152             fcast = None
 153             if coltype == 'string':
 154                 fcast = str
 155             elif coltype.startswith('int'):
 156                 fcast = int
 157             elif coltype.startswith('float'):
 158                 fcast = float
 159             if fcast:
 160                 params[pname] = fcast(params[pname])
 161
 162         def _translate_where(where):
 163             # Translate SQL to PyTables expression
 164             nonlocal start, end
 165             expr = ''
 166             operator = list(where)[0]
 167
 168             if operator in ['and', 'or']:
 169                 subexpr = [_translate_where(e) for e in where[operator]]
 170                 subexpr = filter(lambda e: e, subexpr)
 171                 subexpr = ["({})".format(e) for e in subexpr]
 172                 expr = " {} ".format(self.operators[operator]).join(subexpr)
 173             elif operator == 'exists':
 174                 pass
 175             elif where == {'eq': ['rowid', 'p0']}:
 176                 start = int(params['p0'])
 177                 end = start + 1
 178             elif where == {'gt': ['rowid', 'p0']}:
 179                 start = int(params['p0']) + 1
 180             else:
 181                 left, right = where[operator]
 182
 183                 if isinstance(left, dict):
 184                     left = "(" + _translate_where(left) + ")"
 185                 elif left in params:
 186                     _cast_param(right, left)
 187
 188                 if isinstance(right, dict):
 189                     right = "(" + _translate_where(right) + ")"
 190                 elif right in params:
 191                     _cast_param(left, right)
 192
 193                 expr = "{left} {operator} {right}".format(
 194                     left=left,
 195                     operator=self.operators.get(operator, operator),
 196                     right=right,
 197                 )
 198
 199             return expr
 200
 201         if 'where' in parsed_sql:
 202             if type(parsed_sql['where']) is dict:
 203                 query = _translate_where(parsed_sql['where'])
 204             else:
 205                 query = parsed_sql['where']
 206
 207         # Sort by column
 208         orderby = ''
 209         if 'orderby' in parsed_sql:
 210             orderby = parsed_sql['orderby']
 211             if type(orderby) is list:
 212                 orderby = orderby[0]
 213             orderby = orderby['value']
 214             if orderby == 'rowid':
 215                 orderby = ''
 216
 217         # Limit number of rows
 218         limit = None
 219         if 'limit' in parsed_sql:
 220             limit = int(parsed_sql['limit'])
 221
 222         # Offset
 223         offset = None
 224         if 'offset' in parsed_sql:
 225             offset = int(parsed_sql['offset'])
 226
 227         # Truncate if needed
 228         if page_size and max_returned_rows and truncate:
 229             if max_returned_rows == page_size:
 230                 max_returned_rows += 1
 231
 232         # Execute query
 233         if query:
 234             if not ' glob ' in query:
 235                 table_rows = table.where(query, params, start, end)
 236         elif orderby:
 237             table_rows = table.itersorted(orderby, start=start, stop=end)
 238         else:
 239             table_rows = table.iterrows(start, end)
 240
 241         # Prepare rows
 242         def normalize_field_value(value):
 243             if type(value) is bytes:
 244                 return value.decode('utf-8')
 245             elif not type(value) in (int, float, complex):
 246                 return str(value)
 247             else:
 248                 return value
 249
 250         def make_get_rowid():
 251             if type(table) is tables.table.Table:
 252                 def get_rowid(row):
 253                     return int(row.nrow)
 254             else:
 255                 rowid = start - 1
 256                 def get_rowid(row):
 257                     nonlocal rowid
 258                     rowid += 1
 259                     return rowid
 260             return get_rowid
 261
 262         def make_get_row_value():
 263             if type(table) is tables.table.Table:
 264                 def get_row_value(row, field):
 265                     return row[field]
 266             else:
 267                 def get_row_value(row, field):
 268                     return row
 269             return get_row_value
 270
 271         # Get results
 272         get_rowid = make_get_rowid()
 273         get_row_value = make_get_row_value()
 274         if offset:
 275             table_rows = table_rows[offset:]
 276         count = 0
 277         for table_row in table_rows:
 278             count += 1
 279             if limit is not None and count > limit:
 280                 break
 281             if truncate and max_returned_rows and count > max_returned_rows:
 282                 truncated = True
 283                 break
 284             row = {}
 285             for field in fields:
 286                 field_name = field
 287                 if isinstance(field, dict):
 288                     field_name = field['value']
 289                 if isinstance(field_name, dict) and 'distinct' in field_name:
 290                     field_name = field_name['distinct']
 291                 if field_name == 'rowid':
 292                     row['rowid'] = get_rowid(table_row)
 293                 elif field_name == '*':
 294                     for col in colnames:
 295                         row[col] = normalize_field_value(get_row_value(table_row, col))
 296                 elif isinstance(field_name, dict):
 297                     if field_name.get('count') == '*':
 298                         row['count(*)'] = int(table.nrows)
 299                     elif field_name.get('json_type'):
 300                         field_name = field_name.get('json_type')
 301                         row['json_type(' + field_name + ')'] = _get_field_type(field_name)
 302                     else:
 303                         raise Exception("Function not recognized")
 304                 else:
 305                     row[field_name] = normalize_field_value(get_row_value(table_row, field_name))
 306             results.append(row)
 307
 308         # Prepare query description
 309         for field in [f['value'] if isinstance(f, dict) else f for f in fields]:
 310             if field == '*':
 311                 for col in colnames:
 312                     description += ((col,),)
 313             else:
 314                 description += ((field,),)
 315
 316         return results, truncated, description