'''
This store uses the AWS RDS data API over an RDS Postgres cluster.

Each data type is stored in its own table. In the basic case, these tables have only two columns:
- id uuid   -- For the auto-generated UUID
- data bytea   -- For the serialized data object

This approach allows us to store any un-typed data against a UUID key. If you add an index, we add another column in
the table to store the index values, named according to its ID:
- i_1 <type>   -- The type of this column is determined by type of the indexed value - see _TYPES

The client is then responsible for writing a value into the 'i_1' column every time it writes data (on 'create' or
'update'). This presents a problem when rolling out new indexes that not all clients know about - clients are not
allowed to write data without index values if there are index columns present. To solve this problem, we keep a
registry of not only all the indexes, but the serialized 'key' code that can be loaded into other clients and used to
generate those index values. This registry is stored in the table 'op_store.op_store_metadata'. This table is also used
to track the status of initializing an index.


I do not suggest mucking about with the DB directly - everything should generally be exposed through the APIs only. If
you do need to run SQL directly, the easiest way is to use the query editor in the AWS RDS console.

List all the indexes:
SELECT * FROM op_store.op_store_metadata;

Remove an index:
ALTER TABLE public.<data_type> DROP COLUMN <index> CASCADE;

Remove the whole index table (will break all clients - generally only for testing / development):
DROP TABLE op_store.op_store_metadata;

Kill long-running (hung) queries:
SELECT pid, query, * from pg_stat_activity WHERE state != 'idle' ORDER BY xact_start;
SELECT pg_terminate_backend(<pid>);


You removed this function - formerly, you were calling it before yielding results from the 'iter' function, but I'm not
sure it's actually useful and doesn't break:

@contextmanager
def suspend_transaction(self):
  existing = getattr(self._thread_local, 'transaction_ids', False)
  if (not existing) or (existing[-1] is None):
    raise Exception('No transaction to suspend')
  self._thread_local.transaction_ids.append(None)
  try:
    yield
  finally:
    self._thread_local.transaction_ids.pop()


You also removed the declaraion from the AbstractBaseStore, here's what you had there:

def suspend_transaction(self):
  """
  A context manager that allows other data functions to be called outside of an existing transaction. This can only
  be used within a 'transaction' context.
  """
  raise NotImplementedError()


TODO:
- When you build indexes on fields without modification, you are storing redundant information. You could optimize
  storage space by extracting the indexed field from the data blob, then re-combining it when you read out the data.
  However, this would make it much harder to delete indexes, as you'd have to update every data blob before removing
  the index column.
- I believe you'll have a scaling issue when you try to create indexes via SQL - the call to 'CREATE INDEX
  CONCURRENTLY' could be very slow, which would cause a timeout and fail with the AWS data API.
'''

import boto3
import botocore
import logging
import pickle
import re
import string
import time
import threading
import warnings

from contextlib import contextmanager, suppress
from datetime import date, datetime, timedelta, timezone
from functools import wraps
from uuid import UUID, uuid4

from ..abstract_base_db import AbstractBaseDB
from ..exceptions import *
from ..util.heartbeat import interval
from ..util.fn_serialize import build_fn_key, dumps_fn, loads_fn


logger = logging.getLogger(__name__)


_MAX_INDEX_REFRESH_ATTEMPTS = 3

_BATCH_ITER_SIZE = 100

_INDEX_STALE_TIMEDELTA = timedelta(hours=2)

_INDEX_IN_USE_HEARTBEAT_INTERVAL_S = 3600

# Allows us to map python types to DB types for use with index columns and tells the RDS API how to write data to those
# index columns
_TYPES = {
  bool: {
    'pg': 'boolean',
    'rds-data-api': 'booleanValue'
  },
  str: {
    'pg': 'text',
    'rds-data-api': 'stringValue'
  },
  int: {
    'pg': 'bigint',
    'rds-data-api': 'longValue'
  },
  float: {
    'pg': 'float8',
    'rds-data-api': 'doubleValue'
  },
  UUID: {
    'pg': 'uuid',
    'rds-data-api': 'stringValue',
    'cast': str
  },
  date: {
    'pg': 'date',
    'rds-data-api': 'stringValue',
    'cast': lambda i: i.isoformat()
  },
  datetime: {
    'pg': 'timestamptz',
    'rds-data-api': 'stringValue',
    'cast': lambda i: i.isoformat()
  },
  bytes: {
    'pg': 'bytea',
    'rds-data-api': 'blobValue'
  }
}


class _ColumnDoesNotExist(Exception):
  pass


class _MissingValueForColumn(Exception):
  pass


class _HandleIndexesRetry(Exception):
  pass


def _translate_id_data_records(records):
  # A helper function that parses DB rows containing an id and data
  return [ (UUID(i['stringValue']), pickle.loads(j['blobValue'])) for i, j in records ]


_PG_ALLOWED_CHARS = string.ascii_lowercase + '_'
def _encode_pg(name):
  '''
  A one-way encoding to translate a 'name' in to a postgres-friendly string. Note that we preserve some characters to
  add a modicum of legibility to our table names - this is largely for simplifying development purposes.
  https://www.postgresql.org/docs/9.2/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
  '''
  output = 'op_'
  for char in name:
    if char in _PG_ALLOWED_CHARS:
      output += char
    else:
      output += str(ord(char))
  return output


def _build_rds_params(key_values):
  '''
  Use this function to help you escape params when building input for an _execute call

  Example:
    sql_format, params = _build_rds_params({ 'name': 'Alex', 'email': None, 'date_of_birth': date('1988-06-03') })
    # sql_format = { 'name': ':0', 'email': ':1', 'date_of_birth': ':2::date' }
    # params = [
    #   { 'name': '0', 'value': { 'stringValue': 'Alex' } },
    #   { 'name': '1', 'value': { 'isNull': True } },
    #   { 'name': '2', 'value': { 'stringValue': '1988-06-03' } }
    # ]
    sql = 'INSERT INTO user (name, email, date_of_birth) VALUES (%(name)s, %(email)s, %(date_of_birth)s)'
    self._execute(sql % sql_format, params=params)
  '''
  sql_format = {}
  params = []
  for i, (key, value) in enumerate(key_values.items()):
    sql_format_value = ':%i' % i
    if value is None:
      rds_param_value = { 'isNull': True }
    else:
      type_rules = _TYPES[type(value)]
      if 'cast' in type_rules:
        value = type_rules['cast'](value)
        sql_format_value += '::%s' % type_rules['pg']
      rds_param_value = { type_rules['rds-data-api']: value }
    sql_format[key] = sql_format_value
    params.append({ 'name': str(i), 'value': rds_param_value })
  return sql_format, params


class PostgresIndex:

  def __init__(self, index_name, key=None):
    self.name = index_name
    self.key = key or (lambda i: i[index_name])


class PostgresCondition:

  def __init__(self, index):
    self._index = index
    self.operation = None

  def _chain(self, operator, operand):
    if isinstance(operand, Condition):
      # You could probably make this work if you wanted to...
      raise MalformedConditionException('Cannot compare two conditions')
    self.operation = (self._index, operator, operand)
    return self

  def __bool__(self):
    raise MalformedConditionException('Cannot use conditions as booleans expressions (which precludes chained comparisons)')

  def __lt__(self, operand):
    return self._chain('<', operand)

  def __le__(self, operand):
    return self._chain('<=', operand)

  def __gt__(self, operand):
    return self._chain('>', operand)

  def __ge__(self, operand):
    return self._chain('>=', operand)

  def __eq__(self, operand):
    return self._chain('=', operand)

  def __ne__(self, operand):
    return self._chain('!=', operand)


class DB(AbstractBaseDB):
  '''
  A store built over the RDS data API using postgres
  '''


  def index(self, index_name, key=None):
    return PostgresIndex(index_name, key=key)


  def condition(self, index_name):
    return PostgresCondition(index_name)


  _INIT_DB_SQL = '''
  CREATE EXTENSION IF NOT EXISTS "uuid-ossp";

  CREATE OR REPLACE FUNCTION set_updated_at()
    RETURNS TRIGGER AS $$
      BEGIN
        NEW.updated_at = now(); 
        RETURN NEW;
      END;
    $$ language 'plpgsql';

  CREATE SCHEMA IF NOT EXISTS op_store;

  CREATE TABLE IF NOT EXISTS op_store.op_store_metadata (
    id serial,
    table_name text NOT NULL,
    index_name text NOT NULL,
    index_fn bytea,
    index_fn_key bytea,
    init_completed timestamptz,
    last_used timestamptz DEFAULT now(),
    PRIMARY KEY (id),
    CONSTRAINT unique_index UNIQUE (table_name, index_name, index_fn_key)
  );
  '''
  def __init__(self, aws_config, database_name, rds_cluster_arn, secret_arn):
    self._rds_data = boto3.client('rds-data', **aws_config)
    self._database_name = database_name
    self._rds_cluster_arn = rds_cluster_arn
    self._secret_arn = secret_arn

    # The '_schema' stores all the data types and indexes that we know about. It is a dictionary structured as follows:
    # {
    #   # I tend to refer to this object as a 'data type schema' - see the '_get_data_type_schema' function
    #   <data_type>: {
    #     'table': str, # The name of the corresponding table in the DB - see '_encode_pg'
    #     'update_indexes_lock': lock, # A lock object that ensures thread safety when updating indexes
    #     'indexes': [
    #       # I tend to refer to this object as an 'index schema'
    #       {
    #         'name: str, # The name of the index
    #         'id': int, # The ID of this index data in the 'op_store.op_store_metadata' table
    #         'key': callable, # The single-argument index function - given the object data, it returns the index value
    #         ['discovered']: bool # If present and true, indicates that we did not explicitly initialize this index
    #       },
    #       ...
    #     }
    #   },
    #   ...
    # }
    self._schema = {}

    self._thread_local = threading.local()
    self._execute(self._INIT_DB_SQL)


  def __del__(self):
    self._heartbeat.cancel()


  def describe(self):
    output = {}
    for data_type, data_type_schema in self._schema.items():
      index_names = []
      for index_schema in data_type_schema['indexes']:
        if index_schema.get('discovered', False):
          index_names.append(index)
      output[data_type] = index_names
    return output


  def _get_data_type_schema(self, data_type):
    with suppress(KeyError):
      return self._schema[data_type]
    raise MissingDataTypeException("Data type '%s' has not been initialized" % data_type)


  def _get_initialized_index_schema(self, data_type, index_name):
    data_type_schema = self._get_data_type_schema(data_type)
    for index_schema in data_type_schema['indexes']:
      if index_schema['name'] == index_name and not index_schema.get('discovered', False):
        return index_schema
    raise MissingIndexException("Index '%s' has not been initialized for data type '%s'" % (index_name, data_type))


  _UPDATE_INDEX_LAST_USED_SQL = 'UPDATE op_store.op_store_metadata SET last_used = now() WHERE id IN (%(ids_str)s)'

  @interval(_INDEX_IN_USE_HEARTBEAT_INTERVAL_S)
  def _heartbeat(self):
    # Mark that all the indexes we know about are still in use
    for data_type in self._schema.keys():
      index_schemas = self._get_data_type_schema(data_type)['indexes']
      index_ids = [ i['id'] for i in index_schemas if not i.get('discovered') ]
      if index_ids:
        self._execute(self._UPDATE_INDEX_LAST_USED_SQL % { 'ids_str': (',').join(index_ids) })

    # Clean up unused indexes
    for data_type in self._schema.keys():
      self._drop_stale_indexes(data_type)


  def _execute(self, sql, params={}):
    # Anything we pass in through 'params' will be automatically typed and escaped - it should only be used for values,
    # not things like table names
    sql_format, rds_params = _build_rds_params(params)
    transaction_id_stack = getattr(self._thread_local, 'transaction_ids', [])
    transaction_id_arg = {}
    if transaction_id_stack and (transaction_id_stack[-1] is not None):
      transaction_id_arg = { 'transactionId': transaction_id_stack[-1] }

    try:
      response = self._rds_data.execute_statement(
        database=self._database_name,
        resourceArn=self._rds_cluster_arn,
        secretArn=self._secret_arn,
        sql=sql % sql_format,
        parameters=rds_params,
        **transaction_id_arg
      )
      return response.get('records')
    except botocore.exceptions.ClientError as boto_error:
      e = boto_error.response['Error']
      if not (e['Code'] == 'BadRequestException'):
        raise

      match = re.match('ERROR: column "i_([0-9]+)"(?: of relation "op_[_a-z0-9]+")? does not exist', e['Message'])
      if match:
        raise _ColumnDoesNotExist(match.groups()[0]) from boto_error

      match = re.match('ERROR: missing value for column "i_([0-9]+)"', e['Message'])
      if match:
        raise _MissingValueForColumn(match.groups()[0]) from boto_error

      raise


  @contextmanager
  def transaction(self):
    response = self._rds_data.begin_transaction(
      database=self._database_name,
      resourceArn=self._rds_cluster_arn,
      secretArn=self._secret_arn,
      schema='public'
    )
    if not getattr(self._thread_local, 'transaction_ids', False):
      self._thread_local.transaction_ids = []
    transaction_id = response['transactionId']
    self._thread_local.transaction_ids.append(transaction_id)
    try:
      yield
      self._rds_data.commit_transaction(
        resourceArn=self._rds_cluster_arn,
        secretArn=self._secret_arn,
        transactionId=transaction_id
      )
    except:
      self._rds_data.rollback_transaction(
        resourceArn=self._rds_cluster_arn,
        secretArn=self._secret_arn,
        transactionId=transaction_id
      )
      raise
    finally:
      self._thread_local.transaction_ids.pop()


  _LIST_STALE_INDEXES_SQL = '''
  SELECT id
  FROM op_store.op_store_metadata
  WHERE table_name = %(table)s AND last_used < %(stale_datetime)s
  '''

  _DROP_INDEX_SQL = '''
  DROP TRIGGER IF EXISTS %(table)s_i_%(index_id)i_trigger ON public.%(table)s;
  ALTER TABLE IF EXISTS public.%(table)s DROP COLUMN IF EXISTS i_%(index_id)i CASCADE;
  DROP FUNCTION IF EXISTS %(table)s_i_%(index_id)i_required;
  DELETE FROM op_store.op_store_metadata WHERE id = %(index_id)i;
  '''

  def _drop_stale_indexes(self, data_type, stale=_INDEX_STALE_TIMEDELTA):
    # We allow 'stale' as a parameter to facilitate testing
    stale_datetime = datetime.now(timezone.utc) - stale
    table = _encode_pg(data_type)
    params = { 'table': table, 'stale_datetime': stale_datetime }
    records = self._execute(self._LIST_STALE_INDEXES_SQL, params)
    with self.transaction():
      for row in records:
        index_id = row[0]['longValue']
        self._execute(self._DROP_INDEX_SQL % { 'table': table, 'index_id': index_id })
    if records and (data_type in self._schema):
      self._update_indexes(data_type)


  _INIT_TABLE_SQL = '''
  CREATE TABLE IF NOT EXISTS public.%(table)s (
    id uuid DEFAULT uuid_generate_v4(),
    data bytea,
    PRIMARY KEY (id)
  );
  '''

  _ADD_INDEX_TO_METADATA_SQL = '''
  INSERT INTO op_store.op_store_metadata (table_name, index_name, index_fn, index_fn_key)
  VALUES (%(table)s, %(index_name)s, %(index_fn)s, %(index_fn_key)s)
    ON CONFLICT(table_name, index_name, index_fn_key) DO
      UPDATE SET last_used=now()
  RETURNING id, init_completed;
  '''

  _ADD_INDEX_COLUMN_TRIGGER_SQL = '''
  CREATE OR REPLACE FUNCTION %(table)s_i_%(index_id)i_required()
    RETURNS TRIGGER AS $$
      BEGIN
        IF NEW.i_%(index_id)i IS NULL THEN
          RAISE EXCEPTION 'missing value for column "i_%(index_id)i"';
        END IF;
        RETURN NEW;
      END;
    $$ language 'plpgsql';

  LOCK TABLE public.%(table)s;
  DROP TRIGGER IF EXISTS %(table)s_i_%(index_id)i_trigger ON public.%(table)s;
  CREATE TRIGGER %(table)s_i_%(index_id)i_trigger BEFORE INSERT OR UPDATE ON public.%(table)s
    FOR EACH ROW EXECUTE PROCEDURE %(table)s_i_%(index_id)i_required();
  '''

  def init(self, data_type, indexes=[]):
    if data_type in self._schema:
      raise DuplicateDataTypeException("Data type '%s' has already been initialized" % data_type)

    # From the caller perspective, the 'indexes' name and abstraction make sense, but we rename it here to match our
    # perspective
    index_definitions = indexes

    table_name = _encode_pg(data_type)
    self._schema[data_type] = { 'table': table_name, 'indexes': [], 'update_indexes_lock': threading.Lock() }
    self._execute(self._INIT_TABLE_SQL % { 'table': table_name })

    if not index_definitions:
      self._drop_stale_indexes(data_type)
      return

    index_names = set()
    indexes = self._schema[data_type]['indexes']
    indexes_to_init = []
    
    for index_definition in index_definitions:
      if type(index_definition) == str:
        index_definition = Index(index_definition)
      if index_definition.name in index_names:
        raise DuplicateIndexException("Index '%s' was listed twice" % index_definition.name)
      index_names.add(index_definition.name)

      params = {
        'table': table_name,
        'index_name': index_definition.name,
        'index_fn': dumps_fn(index_definition.key),
        'index_fn_key': build_fn_key(index_definition.key)
      }
      records = self._execute(self._ADD_INDEX_TO_METADATA_SQL, params)
      index_id = records[0][0]['longValue']
      is_initialized = not records[0][1].get('isNull')
      index = { 'name': index_definition.name, 'id': index_id, 'key': index_definition.key }
      indexes.append(index)
      if not is_initialized:
        # This index never finished initialization, let's take it over
        indexes_to_init.append(index)

    # Before we initialize our new indexes, remove any stale ones
    self._drop_stale_indexes(data_type)

    if not indexes_to_init:
      return

    # In the case of a new index, no one else necessarily knows about it yet - we did add it to the op_store_metadata
    # table, but other callers may still be call create / update without setting the new index value(s). Now we create
    # a trigger that throws an error if those indexes are not specified. After the following block has run, everyone
    # calling create / update will be forced to include the new index values (and create the corresponding columns if
    # necessary).

    with self.transaction():
      for index in indexes_to_init:
        self._execute(self._ADD_INDEX_COLUMN_TRIGGER_SQL % { 'table': table_name, 'index_id': index['id'] })

    # We want to make sure that every single row has all of its index values set by the time this function exits so that
    # any queries we run are accurate, thus we have to backfill all existing data (ehich may take a long time).

    logger.info("Starting index initialization for data type '%s'" % data_type)

    # Start with just one value using the regular 'list' function. This will create any missing columns, at which point
    # we can use our optimized '_list_null_indexes' function, or short-circuit if there are no records to update (i.e.
    # a new table).
    to_update = self.list(data_type, limit=1)
    while to_update:
      # TODO: It's possible that other clients may also be running here, especially for big tables. It might be
      # worthwhile to think about how to optimize this process in that context. At least this version is safe.
      for data_id, data in to_update:
        self._update_item_indexes(data_type, data_id, data)
      to_update = self._list_null_indexes(data_type, indexes_to_init)
    self._execute('UPDATE op_store.op_store_metadata SET init_completed = now()')

    logger.info("Finished index initialization for data type '%s'" % data_type)


  _LIST_NULL_INDEXES_SQL = 'SELECT id, data FROM public.%(table)s WHERE %(index_is_null_clauses)s LIMIT %(limit)i'

  def _list_null_indexes(self, data_type, indexes):
    params = {
      'table': self._schema[data_type]['table'],
      'index_is_null_clauses': (' OR ').join([ 'i_%i IS NULL' % i['id'] for i in indexes ]),
      'limit': _BATCH_ITER_SIZE
    }
    records = self._execute(self._LIST_NULL_INDEXES_SQL % params)
    return _translate_id_data_records(records)


  _UPDATE_INDEXES_SQL = '''
  UPDATE public.%(table)s SET %(update_clauses)s WHERE id = %%(id)s AND (%(index_is_null_clauses)s) RETURNING *
  '''

  def _update_item_indexes(self, data_type, data_id, data):
    initial_params = { 'id': data_id }
    build_sql = lambda params: self._UPDATE_INDEXES_SQL % {
      'table': self._schema[data_type]['table'],
      'update_clauses': (',').join([ '%s=%%(%s)s' % (i, i) for i in params.keys() if i != 'id' ]),
      'index_is_null_clauses': (' OR ').join([ '%s IS NULL' % i for i in params.keys() if i != 'id' ])
    }
    self._handle_indexes(data_type, data, initial_params, build_sql)


  # The call to '_create_columns' could in theory be very slow because of the need to create an index. It is only
  # called in the following cases:
  # A) We're calling in the context of initialization, so it can be slow without issue
  # B) Our initialialization didn't have any values, so it couldn't create the column since it didn't know the type.
  #    Therefore, we will only have a few values in the table, so the call will be fast.
  # C) While caller A is initializing the index, caller B happens to be the first one to actually try to create a
  #    value. In that case, we definitely don't want caller B to do the work (the table could be huge).
  # We can ensure case C never happens by only creating the index if it's explicitly initialized and not 'discovered'.

  # Note that I do want to create the index 'CONCURRENTLY' so that the table can still be used while it is building.
  # https://www.postgresql.org/docs/9.1/sql-createindex.html#SQL-CREATEINDEX-CONCURRENTLY
  # TODO: this will likely run into issues with timeouts with the RDS data API...

  _ADD_INDEX_COLUMN_SQL = '''
  ALTER TABLE public.%(table)s ADD COLUMN IF NOT EXISTS i_%(index_id)i %(pg_type)s;
  '''

  _ADD_INDEX_SQL = '''
  CREATE INDEX CONCURRENTLY IF NOT EXISTS "public.%(table)s-i_%(index_id)i"
  ON public.%(table)s(i_%(index_id)i);
  '''

  def _create_columns(self, data_type, params, is_in_initialization=False):
    data_type_schema = self._schema[data_type]
    indexes = data_type_schema['indexes']
    table = data_type_schema['table']

    batch_sql = ''
    for index in indexes:
      py_type = type(params['i_%i' % index['id']])
      if py_type not in _TYPES:
        raise Exception('Unsupported index type: %s' % py_type)
      pg_type = _TYPES[py_type]['pg']
      batch_sql += self._ADD_INDEX_COLUMN_SQL % { 'table': table, 'index_id': index['id'], 'pg_type': pg_type }
      if not index.get('discovered', False):
        batch_sql += self._ADD_INDEX_SQL % { 'table': table, 'index_id': index['id'] }
    self._execute(batch_sql)


  def _is_discovered(self, data_type, index_id):
    indexes = self._schema[data_type]['indexes']
    for index in indexes:
      if index['id'] == index_id:
        return index.get('discovered', False)
    raise Exception("Index ID '%i' not found" % index_id)


  def _execute_update_indexes(self, data_type, sql, params):
    try:
      # Try to run the statement with the indexes we know about
      return self._execute(sql, params)

    except _ColumnDoesNotExist as e:
      # This could happen for two reasons:
      # - If the index was 'discovered', then it has finally been removed by someone else
      # - Otherwise, we just couldn't create the column yet because we didn't have a value to determine its type
      index_id = int(str(e))
      if self._is_discovered(data_type, index_id):
        self._update_indexes(data_type)
      else:
        # In this case, our params haven't changed, so let's try running again directly
        self._create_columns(data_type, params)
        return self._execute_update_indexes(data_type, sql, params)

    except _MissingValueForColumn as e:
      # This happens when someone else adds an index we haven't discovered yet
      self._update_indexes(data_type)

    except Exception:
      raise

    raise _HandleIndexesRetry()


  def _handle_indexes(self, data_type, data, params, build_sql):
    '''
    A helper function to manage to caveats associated with indexes when calling 'create' or 'update'

    build_sql: a single-argument callback that will be passed the complete 'params' dictionary which includes a key /
      value for every index in addition any initial key values. Should return the SQL ready to be executed with the
      specified 'params'
    initial_params: a dictionary containing params to inject into the query - will be supplemented by the index params

    If the number of columns in the built SQL do not match the number of columns we know about, this function will
    automatically update our knowledge of indexes for this data type and try again. Note: this behavior relies on the
    built SQL returning all the columns (i.e. 'RETURNING *').
    '''

    for attempt in range(_MAX_INDEX_REFRESH_ATTEMPTS):

      # Compute the index values and add them to the params
      for index in self._get_data_type_schema(data_type)['indexes']:
        value = index['key'](data)
        if value is None:
          raise Exception("Index values cannot be null (on data type '%s' for index '%s')" % (data_type, index['name']))
        params['i_%i' % index['id']] = value

      sql = build_sql(params)
      with suppress(_HandleIndexesRetry):
        return self._execute_update_indexes(data_type, sql, params)

    # This could technically happen if lots other calls are initializing the data type with new indexes in between our
    # attempts to update the indexes, but that is reasonably never going to happen
    raise Exception('Unexpected columns - tried updating indexes %i times' % (attempt + 1))


  _CREATE_SQL = 'INSERT INTO public.%(table)s (%(columns)s) VALUES (%(values)s) RETURNING *'

  def create(self, data_type, data):
    initial_params = { 'data': pickle.dumps(data) }
    build_sql = lambda params: self._CREATE_SQL % {
      'table': self._schema[data_type]['table'],
      'columns': (',').join(params.keys()),
      'values': (',').join([ '%%(%s)s' % i for i in params.keys() ])
    }
    records = self._handle_indexes(data_type, data, initial_params, build_sql)
    return UUID(records[0][0]['stringValue'])


  _UPDATE_SQL = 'UPDATE public.%(table)s SET %(update_clauses)s WHERE id = %%(id)s RETURNING *'

  def update(self, data_type, data_id, data):
    initial_params = { 'id': data_id, 'data': pickle.dumps(data) }
    build_sql = lambda params: self._UPDATE_SQL % {
      'table': self._schema[data_type]['table'],
      'update_clauses': (',').join([ '%s=%%(%s)s' % (i, i) for i in params.keys() if i != 'id' ])
    }
    self._handle_indexes(data_type, data, initial_params, build_sql)
 

  _GET_SQL = 'SELECT data FROM public.%(table)s WHERE id = %%(id)s'

  def get(self, data_type, data_id):
    data_type_schema = self._get_data_type_schema(data_type)
    records = self._execute(self._GET_SQL % { 'table': data_type_schema['table'] }, { 'id': data_id })
    if not records:
      raise NotFoundException("ID '%s' not found" % data_id)
    return pickle.loads(records[0][0]['blobValue'])


  _GET_BATCH_SQL = 'SELECT id, data FROM public.%(table)s WHERE id IN (%(id_params)s)'

  def get_batch(self, data_type, data_ids):
    params = { str(i): data_id for i, data_id in enumerate(data_ids) }
    if not params:
      # Save ourselves an API call by short-circuiting here
      return []
    table = self._get_data_type_schema(data_type)['table']
    data_ids_str = (', ').join([ '%%(%s)s' % i for i in params.keys() ])
    records = self._execute(self._GET_BATCH_SQL % { 'table': table, 'id_params': data_ids_str }, params)
    return _translate_id_data_records(records)


  _DELETE_SQL = 'DELETE FROM public.%(table)s WHERE id = %%(id)s'

  def delete(self, data_type, data_id):
    self._execute(self._DELETE_SQL % { 'table': self._schema[data_type]['table'] }, { 'id': data_id })


  def create_condition(self, index_name):
    # TODO: this should be static (but still attached to the class)
    return PostgresCondition(index_name)


  def _build_query_params(self, data_type, conditions, order_by, descending, limit):
    data_type_schema = self._get_data_type_schema(data_type)
    table = data_type_schema['table']
    sql = 'SELECT id, data FROM public.%s' % table

    condition_clauses = []
    params = {}
    for i, condition in enumerate(conditions):
      # Protect against invalid conditions - ordinarily I would put the onus to call the function correctly on the
      # caller, but it's not necessarily obvious that an expression like "index('key') is 3" is not allowed
      if not isinstance(condition, Condition):
        raise MalformedConditionException('Expression is not a condition: %r' % condition)
      if condition.operation is None:
        raise MalformedConditionException('Condition does not contain a comparison')
      index_name, operator, operand = condition.operation
      index = self._get_initialized_index_schema(data_type, index_name)
      params[str(i)] = operand
      condition_clause = 'i_%i %s %%(%i)s' % (index['id'], operator, i)
      condition_clauses.append(condition_clause)

    if condition_clauses:
      sql += ' WHERE ' + (' AND ').join(condition_clauses)

    if order_by is None:
      sql += ' ORDER BY id'
      sql += ' DESC' if descending else ''
    else:
      index = self._get_initialized_index_schema(data_type, order_by)
      sql += ' ORDER BY i_%i%s,id' % (index['id'], ' DESC' if descending else '')

    if limit is not None:
      sql += ' LIMIT %i' % limit

    return sql, params


  def iter(self, data_type, *conditions, order_by=None, descending=False, limit=None):
    # TODO: I don't want to expose a 'batch_size' parameter, but that might be useful...
    sql, params = self._build_query_params(data_type, conditions, order_by, descending, limit)
    with self.transaction():
      cursor_name = 'op_' + uuid4().hex
      declare_cursor_sql = 'DECLARE "%s" CURSOR FOR %s' % (cursor_name, sql)
      self._execute(declare_cursor_sql, params)
      fetch_cursor_sql = 'FETCH %i FROM "%s"' % (_BATCH_ITER_SIZE, cursor_name)
      while True:
        records = self._execute(fetch_cursor_sql)
        if not records:
          break
        yield from _translate_id_data_records(records)


  def list(self, data_type, *conditions, order_by=None, descending=False, limit=None):
    sql, params = self._build_query_params(data_type, conditions, order_by, descending, limit)
    try:
      records = self._execute(sql, params)
    except _ColumnDoesNotExist:
      return []
    return _translate_id_data_records(records)


  _LIST_INDEXES_SQL = 'SELECT id, index_name, index_fn FROM op_store.op_store_metadata WHERE table_name = %(table)s'

  def _update_indexes(self, data_type):
    '''
    This function is called in the midst of a 'create' or 'update' call when we find we don't know about all of the
    index columns for the specified data type.

    Note that we limit ths scope of this function to just the specified 'data_type' because we are blocking an
    operation and want to finish ASAP, so we just get the information we need.
    '''

    data_type_schema = self._schema[data_type]
    data_type_schema['update_indexes_lock'].acquire()
    try:
      # First remove all of our 'discovered' indexes
      indexes = data_type_schema['indexes']
      indexes[:] = [ i for i in indexes if not i.get('discovered', False) ]

      # Now we can discover (or perhaps re-discover) any other indexes
      params = { 'table': self._schema[data_type]['table'] }
      records = self._execute(self._LIST_INDEXES_SQL, params)
      for row in records:
        index_id = row[0]['longValue']
        if [ i for i in indexes if i['id'] == index_id ]:
          continue
        index_name = row[1]['stringValue']
        index_key = loads_fn(row[2]['blobValue'])
        indexes.append({ 'id': index_id, 'name': index_name, 'key': index_key, 'discovered': True })
    finally:
      data_type_schema['update_indexes_lock'].release()

  def _set_up_tests(self, data_type):
    # Hide useless warning messages for the rds-data-postgres type. This line only works if placed at this level. I
    # don't really understand what's going on here and hate having implementation-specific details here, but it only
    # seems to be a problm when running the tests:
    # https://github.com/boto/boto3/issues/454
    warnings.filterwarnings('ignore', category=ResourceWarning, message='unclosed.*<ssl.SSLSocket.*>')

    table = _encode_pg(data_type)
    sql = '''
    DROP TABLE IF EXISTS public.%(table_name)s CASCADE;
    DROP TABLE IF EXISTS op_store.op_store_metadata CASCADE;
    '''
    self._execute(sql % { 'table_name': table })
