def __init__(self,
dsn=None,
partition=None,
read_only=False,
name=None,
pool_size=13,
transaction_strategy='resolve_readcommitted',
conn_acquire_timeout=20,
db_schema='public',
store_json=True,
objects_table_name='objects',
blobs_table_name='blobs',
connection_manager=None,
autovacuum=True,
**options):
super(PostgresqlStorage,
self).__init__(read_only,
transaction_strategy=transaction_strategy)
self._dsn = dsn
self._pool_size = pool_size
self._partition_class = partition
self._read_only = read_only
self.__name__ = name
self._conn_acquire_timeout = conn_acquire_timeout
self._options = options
self._store_json = store_json
self._connection_options = {}
self._connection_initialized_on = time.time()
self._db_schema = db_schema
self._objects_table_name = f'{db_schema}.{objects_table_name}'
self._blobs_table_name = f'{db_schema}.{blobs_table_name}'
self._sql = SQLStatements()
self._connection_manager = connection_manager
self._autovacuum = autovacuum
def __init__(self, manager, loop):
self._manager = manager
self._loop = loop
self._queue = asyncio.Queue(loop=loop)
self._closed = False
self._active = False
self._sql = SQLStatements()
def __init__(self,
dsn=None,
partition=None,
read_only=False,
name=None,
pool_size=13,
transaction_strategy='resolve_readcommitted',
conn_acquire_timeout=20,
cache_strategy='dummy',
objects_table_name='objects',
blobs_table_name='blobs',
**options):
super(PostgresqlStorage,
self).__init__(read_only,
transaction_strategy=transaction_strategy,
cache_strategy=cache_strategy)
self._dsn = dsn
self._pool_size = pool_size
self._partition_class = partition
self._read_only = read_only
self.__name__ = name
self._read_conn = None
self._lock = asyncio.Lock()
self._conn_acquire_timeout = conn_acquire_timeout
self._options = options
self._connection_options = {}
self._connection_initialized_on = time.time()
self._objects_table_name = objects_table_name
self._blobs_table_name = blobs_table_name
self._sql = SQLStatements()
def __init__(self, manager, loop):
self._manager = manager
self._loop = loop
self._queue = asyncio.Queue(loop=loop)
self._closed = False
self._active = False
self._sql = SQLStatements()
def __init__(self, dsn=None, partition=None, read_only=False, name=None,
pool_size=13, transaction_strategy='resolve_readcommitted',
conn_acquire_timeout=20, cache_strategy='dummy', db_schema='public',
objects_table_name='objects', blobs_table_name='blobs',
connection_manager=None, autovacuum=True, **options):
super(PostgresqlStorage, self).__init__(
read_only, transaction_strategy=transaction_strategy,
cache_strategy=cache_strategy)
self._dsn = dsn
self._pool_size = pool_size
self._partition_class = partition
self._read_only = read_only
self.__name__ = name
self._conn_acquire_timeout = conn_acquire_timeout
self._options = options
self._connection_options = {}
self._connection_initialized_on = time.time()
self._db_schema = db_schema
self._objects_table_name = f'{db_schema}.{objects_table_name}'
self._blobs_table_name = f'{db_schema}.{blobs_table_name}'
self._sql = SQLStatements()
self._connection_manager = connection_manager
self._autovacuum = autovacuum
class PostgresqlStorage(BaseStorage):
"""Storage to a relational database, based on invalidation polling"""
_dsn = None
_partition_class = None
_large_record_size = 1 << 24
_vacuum_class = PGVacuum
_objects_table_name = 'objects'
_blobs_table_name = 'blobs'
_object_schema = {
'zoid': f'VARCHAR({MAX_UID_LENGTH}) NOT NULL PRIMARY KEY',
'tid': 'BIGINT NOT NULL',
'state_size': 'BIGINT NOT NULL',
'part': 'BIGINT NOT NULL',
'resource': 'BOOLEAN NOT NULL',
'of':
f'VARCHAR({MAX_UID_LENGTH}) REFERENCES {{objects_table_name}} ON DELETE CASCADE',
'otid': 'BIGINT',
'parent_id':
f'VARCHAR({MAX_UID_LENGTH}) REFERENCES {{objects_table_name}} ON DELETE CASCADE', # noqa
'id': 'TEXT',
'type': 'TEXT NOT NULL',
'json': 'JSONB',
'state': 'BYTEA'
}
_blob_schema = {
'bid': f'VARCHAR({MAX_UID_LENGTH}) NOT NULL',
'zoid':
f'VARCHAR({MAX_UID_LENGTH}) NOT NULL REFERENCES {{objects_table_name}} ON DELETE CASCADE',
'chunk_index': 'INT NOT NULL',
'data': 'BYTEA'
}
_initialize_statements = [
'CREATE INDEX IF NOT EXISTS {object_table_name}_tid ON {objects_table_name} (tid);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_of ON {objects_table_name} (of);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_part ON {objects_table_name} (part);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_parent ON {objects_table_name} (parent_id);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_id ON {objects_table_name} (id);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_type ON {objects_table_name} (type);',
'CREATE INDEX IF NOT EXISTS {blob_table_name}_bid ON {blobs_table_name} (bid);',
'CREATE INDEX IF NOT EXISTS {blob_table_name}_zoid ON {blobs_table_name} (zoid);',
'CREATE INDEX IF NOT EXISTS {blob_table_name}_chunk ON {blobs_table_name} (chunk_index);',
'CREATE SEQUENCE IF NOT EXISTS {schema}.tid_sequence;'
'ALTER TABLE {objects_table_name} ADD CONSTRAINT {object_table_name}_parent_id_zoid_check CHECK (parent_id != zoid) NOT VALID;' # noqa
]
_unique_constraint = """CREATE UNIQUE INDEX CONCURRENTLY {constraint_name}_parent_id_id_key
ON {objects_table_name} (parent_id, id)
WHERE parent_id != '{TRASHED_ID}'"""
def __init__(self,
dsn=None,
partition=None,
read_only=False,
name=None,
pool_size=13,
transaction_strategy='resolve_readcommitted',
conn_acquire_timeout=20,
db_schema='public',
store_json=True,
objects_table_name='objects',
blobs_table_name='blobs',
connection_manager=None,
autovacuum=True,
**options):
super(PostgresqlStorage,
self).__init__(read_only,
transaction_strategy=transaction_strategy)
self._dsn = dsn
self._pool_size = pool_size
self._partition_class = partition
self._read_only = read_only
self.__name__ = name
self._conn_acquire_timeout = conn_acquire_timeout
self._options = options
self._store_json = store_json
self._connection_options = {}
self._connection_initialized_on = time.time()
self._db_schema = db_schema
self._objects_table_name = f'{db_schema}.{objects_table_name}'
self._blobs_table_name = f'{db_schema}.{blobs_table_name}'
self._sql = SQLStatements()
self._connection_manager = connection_manager
self._autovacuum = autovacuum
async def finalize(self):
await self._connection_manager.close()
@property
def read_conn(self):
return self._connection_manager.read_conn
@property
def pool(self):
return self._connection_manager.pool
@property
def connection_manager(self):
return self._connection_manager
@property
def lock(self):
return self._connection_manager.lock
@property
def objects_table_name(self):
return self._objects_table_name
async def vacuum(self):
await self.connection_manager.vacuum.run(self._objects_table_name)
while self.connection_manager.vacuum.size > 0:
await asyncio.sleep(0.1)
async def create(self):
# Check DB
log.info('Creating initial database objects')
statements = []
if self._db_schema and self._db_schema != 'public':
statements.extend(
[f'CREATE SCHEMA IF NOT EXISTS {self._db_schema}'])
statements.extend([
get_table_definition(self._objects_table_name,
self._object_schema),
get_table_definition(self._blobs_table_name,
self._blob_schema,
primary_keys=('bid', 'zoid', 'chunk_index'))
])
statements.extend(self._initialize_statements)
for statement in statements:
otable_name = clear_table_name(self._objects_table_name)
if otable_name == 'objects':
otable_name = 'object'
btable_name = clear_table_name(self._blobs_table_name)
if btable_name == 'blobs':
btable_name = 'blob'
statement = statement.format(
objects_table_name=self._objects_table_name,
blobs_table_name=self._blobs_table_name,
# singular, index names
object_table_name=otable_name,
blob_table_name=btable_name,
schema=self._db_schema)
try:
await self.read_conn.execute(statement)
except asyncpg.exceptions.UniqueViolationError:
# this is okay on creation, means 2 getting created at same time
pass
await self.initialize_tid_statements()
async def restart_connection(self, timeout=0.1):
log.error('Connection potentially lost to pg, restarting')
await self._connection_manager.restart()
await self.initialize_tid_statements()
self._connection_initialized_on = time.time()
raise ConflictError('Restarting connection to postgresql')
async def has_unique_constraint(self):
table_name = clear_table_name(self._objects_table_name)
result = await self.read_conn.fetch('''
SELECT * FROM pg_indexes
WHERE tablename = '{}' AND indexname = '{}_parent_id_id_key';
'''.format(table_name, table_name))
return len(result) > 0
async def initialize(self, loop=None, **kw):
self._connection_options = kw
if self._connection_manager is None:
self._connection_manager = PGConnectionManager(
dsn=self._dsn,
pool_size=self._pool_size,
connection_options=self._connection_options,
conn_acquire_timeout=self._conn_acquire_timeout,
vacuum_class=self._vacuum_class)
await self._connection_manager.initialize(loop, **kw)
async with self.lock:
if await self.has_unique_constraint():
self._supports_unique_constraints = True
trash_sql = self._sql.get('CREATE_TRASH', self._objects_table_name)
try:
await self.initialize_tid_statements()
await self.read_conn.execute(trash_sql)
except asyncpg.exceptions.ReadOnlySQLTransactionError:
# Not necessary for read-only pg
pass
except (asyncpg.exceptions.UndefinedTableError,
asyncpg.exceptions.InvalidSchemaNameError):
await self.create()
# only available on new databases
await self.read_conn.execute(
self._unique_constraint.format(
objects_table_name=self._objects_table_name,
constraint_name=clear_table_name(
self._objects_table_name),
TRASHED_ID=TRASHED_ID))
self._supports_unique_constraints = True
await self.initialize_tid_statements()
await self.read_conn.execute(trash_sql)
# migrate to larger VARCHAR size...
result = await self.read_conn.fetch("""
select * from information_schema.columns
where table_name='{}'""".format(self._objects_table_name))
if len(result) > 0 and result[0][
'character_maximum_length'] != MAX_UID_LENGTH:
log.warn('Migrating VARCHAR key length')
await self.read_conn.execute(f'''
ALTER TABLE {self._objects_table_name} ALTER COLUMN zoid TYPE varchar({MAX_UID_LENGTH})'''
)
await self.read_conn.execute(f'''
ALTER TABLE {self._objects_table_name} ALTER COLUMN of TYPE varchar({MAX_UID_LENGTH})'''
)
await self.read_conn.execute(f'''
ALTER TABLE {self._objects_table_name} ALTER COLUMN parent_id TYPE varchar({MAX_UID_LENGTH})'''
)
await self.read_conn.execute(f'''
ALTER TABLE {self._blobs_table_name} ALTER COLUMN bid TYPE varchar({MAX_UID_LENGTH})'''
)
await self.read_conn.execute(f'''
ALTER TABLE {self._blobs_table_name} ALTER COLUMN zoid TYPE varchar({MAX_UID_LENGTH})'''
)
self._connection_initialized_on = time.time()
async def initialize_tid_statements(self):
self._stmt_next_tid = await self.read_conn.prepare(
NEXT_TID.format(schema=self._db_schema))
self._stmt_max_tid = await self.read_conn.prepare(
MAX_TID.format(schema=self._db_schema))
async def remove(self):
"""Reset the tables"""
async with self.pool.acquire() as conn:
await conn.execute("DROP TABLE IF EXISTS {};".format(
self._blobs_table_name))
await conn.execute("DROP TABLE IF EXISTS {};".format(
self._objects_table_name))
@restart_conn_on_exception
async def open(self):
conn = await self.pool.acquire(timeout=self._conn_acquire_timeout)
return conn
async def close(self, con):
try:
await shield(asyncio.wait_for(self.pool.release(con, timeout=1),
1))
except (asyncio.CancelledError, RuntimeError, asyncio.TimeoutError,
asyncpg.exceptions.ConnectionDoesNotExistError):
pass
async def terminate(self, conn):
conn.terminate()
async def load(self, txn, oid):
sql = self._sql.get('GET_OID', self._objects_table_name)
async with txn._lock:
objects = await self.get_one_row(txn, sql, oid)
if objects is None:
raise KeyError(oid)
return objects
@profilable
async def store(self, oid, old_serial, writer, obj, txn):
assert oid is not None
pickled = writer.serialize() # This calls __getstate__ of obj
if len(pickled) >= self._large_record_size:
log.info(f"Large object {obj.__class__}: {len(pickled)}")
if self._store_json:
json_dict = await writer.get_json()
json = ujson.dumps(json_dict)
else:
json = {}
part = writer.part
if part is None:
part = 0
update = False
statement_sql = self._sql.get('NAIVE_UPSERT', self._objects_table_name)
if not obj.__new_marker__ and obj.__serial__ is not None:
# we should be confident this is an object update
statement_sql = self._sql.get('UPDATE', self._objects_table_name)
update = True
conn = await txn.get_connection()
async with txn._lock:
try:
result = await conn.fetch(
statement_sql,
oid, # The OID of the object
txn._tid, # Our TID
len(pickled), # Len of the object
part, # Partition indicator
writer.resource, # Is a resource ?
writer.of, # It belogs to a main
old_serial, # Old serial
writer.parent_id, # Parent OID
writer.id, # Traversal ID
writer.type, # Guillotina type
json, # JSON catalog
pickled # Pickle state)
)
except asyncpg.exceptions.UniqueViolationError as ex:
if 'Key (parent_id, id)' in ex.detail:
raise ConflictIdOnContainer(ex)
raise
except asyncpg.exceptions.ForeignKeyViolationError:
txn.deleted[obj.__uuid__] = obj
raise TIDConflictError(
f'Bad value inserting into database that could be caused '
f'by a bad cache value. This should resolve on request retry.',
oid, txn, old_serial, writer)
except asyncpg.exceptions._base.InterfaceError as ex:
if 'another operation is in progress' in ex.args[0]:
raise ConflictError(
f'asyncpg error, another operation in progress.', oid,
txn, old_serial, writer)
raise
except asyncpg.exceptions.DeadlockDetectedError:
raise ConflictError(f'Deadlock detected.', oid, txn,
old_serial, writer)
if len(result) != 1 or result[0]['count'] != 1:
if update:
# raise tid conflict error
raise TIDConflictError(
f'Mismatch of tid of object being updated. This is likely '
f'caused by a cache invalidation race condition and should '
f'be an edge case. This should resolve on request retry.',
oid, txn, old_serial, writer)
else:
log.error('Incorrect response count from database update. '
'This should not happen. tid: {}'.format(
txn._tid))
await txn._cache.store_object(obj, pickled)
async def _txn_oid_commit_hook(self, status, oid):
await self._connection_manager._vacuum.add_to_queue(
oid, self._objects_table_name)
async def delete(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('TRASH_PARENT_ID', self._objects_table_name)
async with txn._lock:
# for delete, we reassign the parent id and delete in the vacuum task
await conn.execute(sql, oid)
if self._autovacuum:
txn.add_after_commit_hook(self._txn_oid_commit_hook, oid)
async def _check_bad_connection(self, ex):
# we do not use transaction lock here but a storage lock because
# a storage object has a shard conn for reads
for err in ('connection is closed', 'pool is closed',
'connection was closed'):
if err in str(ex):
if (time.time() - self._connection_initialized_on
) > BAD_CONNECTION_RESTART_DELAY:
# we need to make sure we aren't calling this over and over again
async with self.lock:
return await self.restart_connection()
@restart_conn_on_exception
async def get_next_tid(self, txn):
async with self.lock:
return await self._stmt_next_tid.fetchval()
@restart_conn_on_exception
async def get_current_tid(self, txn):
async with self.lock:
return await self._stmt_max_tid.fetchval()
async def get_one_row(self, txn, sql, *args, prepare=False):
conn = await txn.get_connection()
# Helper function to provide easy adaptation to cockroach
if prepare:
# latest version of asyncpg has prepare bypassing statement cache
smt = await conn.prepare(sql)
return await smt.fetchrow(*args)
else:
return await conn.fetchrow(sql, *args)
def _db_transaction_factory(self, txn):
# make sure asycpg knows this is a new transaction
if txn._db_conn._con is not None:
txn._db_conn._con._top_xact = None
return txn._db_conn.transaction(
readonly=txn._manager._storage._read_only)
@restart_conn_on_exception
async def _async_db_transaction_factory(self, txn):
return self._db_transaction_factory(txn)
async def start_transaction(self, txn, retries=0):
error = None
conn = await txn.get_connection()
async with txn._lock:
txn._db_txn = await self._async_db_transaction_factory(txn)
try:
await txn._db_txn.start()
return
except (asyncpg.exceptions.InterfaceError,
asyncpg.exceptions.InternalServerError) as ex:
error = ex
if error is not None:
if retries > 2:
raise error # pylint: disable=E0702
restart = rollback = False
if isinstance(error, asyncpg.exceptions.InternalServerError):
restart = True
if error.sqlstate == 'XX000':
rollback = True
elif ('manually started transaction' in error.args[0]
or 'connection is closed' in error.args[0]):
restart = True
if 'manually started transaction' in error.args[0]:
rollback = True
if rollback:
try:
# thinks we're manually in txn, manually rollback and try again...
await conn.execute('ROLLBACK;')
except asyncpg.exceptions._base.InterfaceError:
# we're okay with this error here...
pass
if restart:
await self.close(conn)
txn._db_conn = await self.open()
return await self.start_transaction(txn, retries + 1)
async def get_conflicts(self, txn):
async with self.lock:
if len(txn.modified) == 0:
return []
# use storage lock instead of transaction lock
if len(txn.modified) < 1000:
# if it's too large, we're not going to check on object ids
modified_oids = [k for k in txn.modified.keys()]
sql = self._sql.get('TXN_CONFLICTS_ON_OIDS',
self._objects_table_name)
return await self.read_conn.fetch(sql, txn._tid, modified_oids)
else:
sql = self._sql.get('TXN_CONFLICTS', self._objects_table_name)
return await self.read_conn.fetch(sql, txn._tid)
async def commit(self, transaction):
if transaction._db_txn is not None:
async with transaction._lock:
await transaction._db_txn.commit()
elif (self._transaction_strategy not in ('none', 'tidonly')
and not transaction._skip_commit):
log.warning('Do not have db transaction to commit')
return transaction._tid
async def abort(self, transaction):
if transaction._db_txn is not None:
async with transaction._lock:
try:
await transaction._db_txn.rollback()
except asyncpg.exceptions._base.InterfaceError:
# we're okay with this error here...
pass
# reads don't need transaction necessarily so don't log
# else:
# log.warning('Do not have db transaction to rollback')
# Introspection
async def get_page_of_keys(self, txn, oid, page=1, page_size=1000):
conn = await txn.get_connection()
keys = []
sql = self._sql.get('BATCHED_GET_CHILDREN_KEYS',
self._objects_table_name)
for record in await conn.fetch(sql, oid, page_size,
(page - 1) * page_size):
keys.append(record['id'])
return keys
async def keys(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('GET_CHILDREN_KEYS', self._objects_table_name)
async with txn._lock:
result = await conn.fetch(sql, oid)
return result
async def get_child(self, txn, parent_oid, id):
sql = self._sql.get('GET_CHILD', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, parent_oid, id)
return result
async def get_children(self, txn, parent_oid, ids):
conn = await txn.get_connection()
sql = self._sql.get('GET_CHILDREN_BATCH', self._objects_table_name)
async with txn._lock:
return await conn.fetch(sql, parent_oid, ids)
async def has_key(self, txn, parent_oid, id):
sql = self._sql.get('EXIST_CHILD', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, parent_oid, id)
if result is None:
return False
else:
return True
async def len(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('NUM_CHILDREN', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql, oid)
return result
async def items(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('GET_CHILDREN', self._objects_table_name)
async for record in conn.cursor(sql, oid):
# locks are dangerous in cursors since comsuming code might do
# sub-queries and they you end up with a deadlock
yield record
async def get_annotation(self, txn, oid, id):
sql = self._sql.get('GET_ANNOTATION', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, oid, id, prepare=True)
if result is not None and result['parent_id'] == TRASHED_ID:
result = None
return result
async def get_annotation_keys(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('GET_ANNOTATIONS_KEYS', self._objects_table_name)
async with txn._lock:
result = await conn.fetch(sql, oid)
items = []
for item in result:
if item['parent_id'] != TRASHED_ID:
items.append(item)
return items
async def write_blob_chunk(self, txn, bid, oid, chunk_index, data):
sql = self._sql.get('HAS_OBJECT', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, oid)
if result is None:
# check if we have a referenced ob, could be new and not in db yet.
# if so, create a stub for it here...
conn = await txn.get_connection()
async with txn._lock:
await conn.execute(
f'''INSERT INTO {self._objects_table_name}
(zoid, tid, state_size, part, resource, type)
VALUES ($1::varchar({MAX_UID_LENGTH}), -1, 0, 0, TRUE, 'stub')''', oid)
conn = await txn.get_connection()
sql = self._sql.get('INSERT_BLOB_CHUNK', self._blobs_table_name)
async with txn._lock:
return await conn.execute(sql, bid, oid, chunk_index, data)
async def read_blob_chunk(self, txn, bid, chunk=0):
sql = self._sql.get('READ_BLOB_CHUNK', self._blobs_table_name)
async with txn._lock:
return await self.get_one_row(txn, sql, bid, chunk)
async def read_blob_chunks(self, txn, bid):
conn = await txn.get_connection()
async for record in conn.cursor(bid):
# locks are dangerous in cursors since comsuming code might do
# sub-queries and they you end up with a deadlock
yield record
async def del_blob(self, txn, bid):
conn = await txn.get_connection()
sql = self._sql.get('DELETE_BLOB', self._blobs_table_name)
async with txn._lock:
await conn.execute(sql, bid)
async def get_total_number_of_objects(self, txn):
conn = await txn.get_connection()
sql = self._sql.get('NUM_ROWS', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql)
return result
async def get_total_number_of_resources(self, txn):
conn = await txn.get_connection()
sql = self._sql.get('NUM_RESOURCES', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql)
return result
async def get_total_resources_of_type(self, txn, type_):
conn = await txn.get_connection()
sql = self._sql.get('NUM_RESOURCES_BY_TYPE', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql, type_)
return result
# Massive treatment without security
async def _get_page_resources_of_type(self, txn, type_, page, page_size):
conn = await txn.get_connection()
async with txn._lock:
keys = []
sql = self._sql.get('RESOURCES_BY_TYPE', self._objects_table_name)
for record in await conn.fetch(sql, type_, page_size,
(page - 1) * page_size):
keys.append(record)
return keys
class PGVacuum:
def __init__(self, manager, loop):
self._manager = manager
self._loop = loop
self._queue = asyncio.Queue(loop=loop)
self._closed = False
self._active = False
self._sql = SQLStatements()
async def initialize(self):
while not self._closed:
try:
await self._initialize()
except (concurrent.futures.CancelledError, RuntimeError):
# we're okay with the task getting cancelled
return
async def _initialize(self):
while not self._closed:
oid = None
try:
oid, table_name = await self._queue.get()
self._active = True
await shield(self.vacuum(oid, table_name))
except (concurrent.futures.CancelledError, RuntimeError):
raise
except Exception:
log.warning(f'Error vacuuming oid {oid}', exc_info=True)
finally:
self._active = False
try:
self._queue.task_done()
except ValueError:
pass
async def run(self, table_name):
'''
get existing trashed objects, push them on the queue...
there might be contention, but that is okay
'''
async with self._manager.pool.acquire(
timeout=self._manager._conn_acquire_timeout) as conn:
try:
sql = self._sql.get('GET_TRASHED_OBJECTS', table_name)
for record in await conn.fetch(sql):
self._queue.put_nowait((record['zoid'], table_name))
except concurrent.futures.TimeoutError:
log.info('Timed out connecting to storage')
except Exception:
log.warning('Error deleting trashed object', exc_info=True)
async def add_to_queue(self, oid, table_name):
if self._closed:
raise Exception('Closing down')
await self._queue.put((oid, table_name))
async def vacuum(self, oid, table_name):
'''
DELETED objects has parent id changed to the trashed ob for the oid...
'''
async with self._manager.pool.acquire(
timeout=self._manager._conn_acquire_timeout) as conn:
sql = self._sql.get('DELETE_OBJECT', table_name)
try:
await conn.execute(sql, oid)
except Exception:
log.warning('Error deleting trashed object', exc_info=True)
async def finalize(self):
self._closed = True
try:
await asyncio.wait_for(self._queue.join(), 2)
except asyncio.TimeoutError:
pass
@property
def size(self):
return self._queue.qsize()
class PostgresqlStorage(BaseStorage):
"""Storage to a relational database, based on invalidation polling"""
_dsn = None
_partition_class = None
_large_record_size = 1 << 24
_vacuum_class = PGVacuum
_connection_manager_class = PGConnectionManager
_objects_table_name = "objects"
_blobs_table_name = "blobs"
_object_schema = {
"zoid": f"VARCHAR({MAX_UID_LENGTH}) NOT NULL PRIMARY KEY",
"tid": "BIGINT NOT NULL",
"state_size": "BIGINT NOT NULL",
"part": "BIGINT NOT NULL",
"resource": "BOOLEAN NOT NULL",
"of": f"VARCHAR({MAX_UID_LENGTH}) REFERENCES {{objects_table_name}} ON DELETE CASCADE",
"otid": "BIGINT",
"parent_id": f"VARCHAR({MAX_UID_LENGTH}) REFERENCES {{objects_table_name}} ON DELETE CASCADE", # noqa
"id": "TEXT",
"type": "TEXT NOT NULL",
"json": "JSONB",
"state": "BYTEA",
}
_blob_schema = {
"bid": f"VARCHAR({MAX_UID_LENGTH}) NOT NULL",
"zoid": f"VARCHAR({MAX_UID_LENGTH}) NOT NULL REFERENCES {{objects_table_name}} ON DELETE CASCADE",
"chunk_index": "INT NOT NULL",
"data": "BYTEA",
}
_initialize_statements = [
"CREATE INDEX IF NOT EXISTS {object_table_name}_tid ON {objects_table_name} (tid);",
"CREATE INDEX IF NOT EXISTS {object_table_name}_of ON {objects_table_name} (of);",
"CREATE INDEX IF NOT EXISTS {object_table_name}_part ON {objects_table_name} (part);",
"CREATE INDEX IF NOT EXISTS {object_table_name}_parent ON {objects_table_name} (parent_id);",
"CREATE INDEX IF NOT EXISTS {object_table_name}_id ON {objects_table_name} (id);",
"CREATE INDEX IF NOT EXISTS {object_table_name}_type ON {objects_table_name} (type);",
"CREATE INDEX IF NOT EXISTS {blob_table_name}_bid ON {blobs_table_name} (bid);",
"CREATE INDEX IF NOT EXISTS {blob_table_name}_zoid ON {blobs_table_name} (zoid);",
"CREATE INDEX IF NOT EXISTS {blob_table_name}_chunk ON {blobs_table_name} (chunk_index);",
"ALTER TABLE {objects_table_name} ADD CONSTRAINT {object_table_name}_parent_id_zoid_check CHECK (parent_id != zoid) NOT VALID;", # noqa
]
_unique_constraints = [
"""CREATE UNIQUE INDEX CONCURRENTLY {constraint_name}_parent_id_id_key
ON {objects_table_name} (parent_id, id)
WHERE parent_id != '{TRASHED_ID}' """,
"""CREATE UNIQUE INDEX CONCURRENTLY {constraint_name}_annotations_unique ON {objects_table_name} (of, id);""",
]
def __init__(
self,
dsn=None,
partition=None,
read_only=False,
name=None,
pool_size=13,
transaction_strategy="resolve_readcommitted",
conn_acquire_timeout=20,
db_schema="public",
store_json=True,
objects_table_name="objects",
blobs_table_name="blobs",
connection_manager=None,
autovacuum=True,
**options,
):
super(PostgresqlStorage, self).__init__(read_only, transaction_strategy=transaction_strategy)
self._dsn = dsn
self._pool_size = pool_size
self._partition_class = partition
self._read_only = read_only
self.__name__ = name
self._conn_acquire_timeout = conn_acquire_timeout
self._options = options
self._store_json = store_json
self._connection_options = {}
self._connection_initialized_on = time.time()
self._db_schema = db_schema
self._objects_table_name = f"{db_schema}.{objects_table_name}"
self._blobs_table_name = f"{db_schema}.{blobs_table_name}"
self._sql = SQLStatements()
self._connection_manager = connection_manager
self._autovacuum = autovacuum
async def finalize(self):
await self._connection_manager.close()
@property
def sql(self):
return self._sql
@property
def read_conn(self):
return self._connection_manager.read_conn
@property
def pool(self):
return self._connection_manager.pool
@property
def connection_manager(self):
return self._connection_manager
@property
def lock(self):
return self._connection_manager.lock
@property
def stmt_next_tid(self):
return self._connection_manager._stmt_next_tid
@property
def stmt_max_tid(self):
return self._connection_manager._stmt_max_tid
@property
def objects_table_name(self):
return self._objects_table_name
async def create(self, conn=None):
if conn is None:
conn = self.read_conn
# Check DB
log.info("Creating initial database objects")
statements = []
if self._db_schema and self._db_schema != "public":
statements.extend([f"CREATE SCHEMA IF NOT EXISTS {self._db_schema}"])
statements.extend(
[
get_table_definition(self._objects_table_name, self._object_schema),
get_table_definition(
self._blobs_table_name, self._blob_schema, primary_keys=("bid", "zoid", "chunk_index")
),
]
)
statements.extend(self._initialize_statements)
for statement in statements:
otable_name = clear_table_name(self._objects_table_name)
if otable_name == "objects":
otable_name = "object"
btable_name = clear_table_name(self._blobs_table_name)
if btable_name == "blobs":
btable_name = "blob"
statement = statement.format(
objects_table_name=self._objects_table_name,
blobs_table_name=self._blobs_table_name,
# singular, index names
object_table_name=otable_name,
blob_table_name=btable_name,
schema=self._db_schema,
)
try:
await conn.execute(statement)
except asyncpg.exceptions.UniqueViolationError:
# this is okay on creation, means 2 getting created at same time
pass
async def restart_connection(self, timeout=0.1):
log.error("Connection potentially lost to pg, restarting")
await self._connection_manager.restart()
self._connection_initialized_on = time.time()
raise ConflictError("Restarting connection to postgresql")
async def has_unique_constraint(self, conn):
table_name = clear_table_name(self._objects_table_name)
result = await conn.fetch(
"""
SELECT * FROM pg_indexes
WHERE tablename = '{}' AND indexname = '{}_parent_id_id_key';
""".format(
table_name, table_name
)
)
return len(result) > 0
async def initialize(self, loop=None, **kw):
self._connection_options = kw
if self._connection_manager is None:
self._connection_manager = self._connection_manager_class(
dsn=self._dsn,
pool_size=self._pool_size,
connection_options=self._connection_options,
conn_acquire_timeout=self._conn_acquire_timeout,
vacuum_class=self._vacuum_class,
autovacuum=self._autovacuum,
db_schema=self._db_schema,
)
await self._connection_manager.initialize(loop, **kw)
async with self.pool.acquire(timeout=self._conn_acquire_timeout) as conn:
if await self.has_unique_constraint(conn):
self._supports_unique_constraints = True
trash_sql = self._sql.get("CREATE_TRASH", self._objects_table_name)
try:
await conn.execute(trash_sql)
except asyncpg.exceptions.ReadOnlySQLTransactionError:
# Not necessary for read-only pg
pass
except (asyncpg.exceptions.UndefinedTableError, asyncpg.exceptions.InvalidSchemaNameError):
async with conn.transaction():
await self.create(conn)
# only available on new databases
for constraint in self._unique_constraints:
await conn.execute(
constraint.format(
objects_table_name=self._objects_table_name,
constraint_name=clear_table_name(self._objects_table_name),
TRASHED_ID=TRASHED_ID,
).replace("CONCURRENTLY", "")
)
self._supports_unique_constraints = True
await conn.execute(trash_sql)
await notify(StorageCreatedEvent(self, db_conn=conn))
self._connection_initialized_on = time.time()
async def remove(self):
"""Reset the tables"""
async with self.pool.acquire(timeout=self._conn_acquire_timeout) as conn:
await conn.execute("DROP TABLE IF EXISTS {} CASCADE;".format(self._blobs_table_name))
await conn.execute("DROP TABLE IF EXISTS {} CASCADE;".format(self._objects_table_name))
@restart_conn_on_exception
async def open(self):
return await self.pool.acquire(timeout=self._conn_acquire_timeout)
async def close(self, con):
try:
await shield(self.pool.release(con, timeout=1))
except (asyncio.CancelledError, asyncio.TimeoutError, asyncpg.exceptions.ConnectionDoesNotExistError):
log.warning("Exception on connection close", exc_info=True)
async def terminate(self, conn):
log.warning(f"Terminate connection {conn}", exc_info=True)
conn.terminate()
async def load(self, txn, oid):
sql = self._sql.get("GET_OID", self._objects_table_name)
async with txn._lock:
objects = await self.get_one_row(txn, sql, oid)
if objects is None:
raise KeyError(oid)
return objects
@profilable
async def store(self, oid, old_serial, writer, obj, txn):
assert oid is not None
pickled = writer.serialize() # This calls __getstate__ of obj
if len(pickled) >= self._large_record_size:
log.info(f"Large object {obj.__class__}: {len(pickled)}")
if self._store_json:
json_dict = await writer.get_json()
json = orjson.dumps(json_dict).decode("utf-8")
else:
json = None
part = writer.part
if part is None:
part = 0
update = False
statement_sql = self._sql.get("NAIVE_UPSERT", self._objects_table_name)
if not obj.__new_marker__ and obj.__serial__ is not None:
# we should be confident this is an object update
statement_sql = self._sql.get("UPDATE", self._objects_table_name)
update = True
conn = await txn.get_connection()
async with txn._lock:
try:
result = await conn.fetch(
statement_sql,
oid, # The OID of the object
txn._tid, # Our TID
len(pickled), # Len of the object
part, # Partition indicator
writer.resource, # Is a resource ?
writer.of, # It belogs to a main
old_serial, # Old serial
writer.parent_id, # Parent OID
writer.id, # Traversal ID
writer.type, # Guillotina type
json, # JSON catalog
pickled, # Pickle state)
)
except asyncpg.exceptions.UniqueViolationError as ex:
if "Key (parent_id, id)" in ex.detail or "Key (of, id)" in ex.detail:
raise ConflictIdOnContainer(ex)
raise
except asyncpg.exceptions.ForeignKeyViolationError:
txn.deleted[obj.__uuid__] = obj
raise TIDConflictError(
f"Bad value inserting into database that could be caused "
f"by a bad cache value. This should resolve on request retry.",
oid,
txn,
old_serial,
writer,
)
except asyncpg.exceptions._base.InterfaceError as ex:
if "another operation is in progress" in ex.args[0]:
raise ConflictError(
f"asyncpg error, another operation in progress.", oid, txn, old_serial, writer
)
raise
except asyncpg.exceptions.DeadlockDetectedError:
raise ConflictError(f"Deadlock detected.", oid, txn, old_serial, writer)
if len(result) != 1 or result[0]["count"] != 1:
if update:
# raise tid conflict error
raise TIDConflictError(
f"Mismatch of tid of object being updated. This is likely "
f"caused by a cache invalidation race condition and should "
f"be an edge case. This should resolve on request retry.",
oid,
txn,
old_serial,
writer,
)
else:
log.error(
"Incorrect response count from database update. "
"This should not happen. tid: {}".format(txn._tid)
)
await txn._cache.store_object(obj, pickled)
async def _txn_oid_commit_hook(self, status, oid):
if self._connection_manager._vacuum is not None:
await self._connection_manager._vacuum.add_to_queue(oid, self._objects_table_name)
async def delete(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get("TRASH_PARENT_ID", self._objects_table_name)
async with txn._lock:
# for delete, we reassign the parent id and delete in the vacuum task
await conn.execute(sql, oid)
if self._autovacuum:
txn.add_after_commit_hook(self._txn_oid_commit_hook, oid)
async def _check_bad_connection(self, ex):
# we do not use transaction lock here but a storage lock because
# a storage object has a shard conn for reads
for err in ("connection is closed", "pool is closed", "connection was closed"):
if err in str(ex):
if (time.time() - self._connection_initialized_on) > BAD_CONNECTION_RESTART_DELAY:
# we need to make sure we aren't calling this over and over again
async with self.lock:
return await self.restart_connection()
@restart_conn_on_exception
async def get_next_tid(self, txn):
async with self.lock:
return await self.stmt_next_tid.fetchval()
@restart_conn_on_exception
async def get_current_tid(self, txn):
async with self.lock:
return await self.stmt_max_tid.fetchval()
async def get_one_row(self, txn, sql, *args, prepare=False):
conn = await txn.get_connection()
# Helper function to provide easy adaptation to cockroach
if prepare:
# latest version of asyncpg has prepare bypassing statement cache
smt = await conn.prepare(sql)
return await smt.fetchrow(*args)
else:
return await conn.fetchrow(sql, *args)
def _db_transaction_factory(self, txn):
# make sure asycpg knows this is a new transaction
if txn._db_conn._con is not None:
txn._db_conn._con._top_xact = None
return txn._db_conn.transaction(readonly=txn._manager._storage._read_only)
@restart_conn_on_exception
async def _async_db_transaction_factory(self, txn):
return self._db_transaction_factory(txn)
async def start_transaction(self, txn, retries=0):
error = None
conn = await txn.get_connection()
async with txn._lock:
txn._db_txn = await self._async_db_transaction_factory(txn)
try:
await txn._db_txn.start()
return
except (asyncpg.exceptions.InterfaceError, asyncpg.exceptions.InternalServerError) as ex:
error = ex
if error is not None:
if retries > 2:
raise error # pylint: disable=E0702
restart = rollback = False
if isinstance(error, asyncpg.exceptions.InternalServerError):
restart = True
if error.sqlstate == "XX000":
rollback = True
elif "manually started transaction" in error.args[0] or "connection is closed" in error.args[0]:
restart = True
if "manually started transaction" in error.args[0]:
rollback = True
if rollback:
try:
# thinks we're manually in txn, manually rollback and try again...
await conn.execute("ROLLBACK;")
except asyncpg.exceptions._base.InterfaceError:
# we're okay with this error here...
pass
if restart:
await self.close(conn)
txn._db_conn = await self.open()
return await self.start_transaction(txn, retries + 1)
async def get_conflicts(self, txn):
async with self.lock:
if len(txn.modified) == 0:
return []
# use storage lock instead of transaction lock
if len(txn.modified) < 1000:
# if it's too large, we're not going to check on object ids
modified_oids = [k for k in txn.modified.keys()]
sql = self._sql.get("TXN_CONFLICTS_ON_OIDS", self._objects_table_name)
return await self.read_conn.fetch(sql, txn._tid, modified_oids)
else:
sql = self._sql.get("TXN_CONFLICTS", self._objects_table_name)
return await self.read_conn.fetch(sql, txn._tid)
async def commit(self, transaction):
if transaction._db_txn is not None:
async with transaction._lock:
await transaction._db_txn.commit()
elif self._transaction_strategy not in ("none", "tidonly") and not transaction._skip_commit:
log.warning("Do not have db transaction to commit")
return transaction._tid
async def abort(self, transaction):
if transaction._db_txn is not None:
async with transaction._lock:
try:
await transaction._db_txn.rollback()
except asyncpg.exceptions._base.InterfaceError:
# we're okay with this error here...
pass
# reads don't need transaction necessarily so don't log
# else:
# log.warning('Do not have db transaction to rollback')
# Introspection
async def get_page_of_keys(self, txn, oid, page=1, page_size=1000):
conn = await txn.get_connection()
keys = []
sql = self._sql.get("BATCHED_GET_CHILDREN_KEYS", self._objects_table_name)
for record in await conn.fetch(sql, oid, page_size, (page - 1) * page_size):
keys.append(record["id"])
return keys
async def keys(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get("GET_CHILDREN_KEYS", self._objects_table_name)
async with txn._lock:
result = await conn.fetch(sql, oid)
return result
async def get_child(self, txn, parent_oid, id):
sql = self._sql.get("GET_CHILD", self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, parent_oid, id)
return result
async def get_children(self, txn, parent_oid, ids):
conn = await txn.get_connection()
sql = self._sql.get("GET_CHILDREN_BATCH", self._objects_table_name)
async with txn._lock:
return await conn.fetch(sql, parent_oid, ids)
async def has_key(self, txn, parent_oid, id):
sql = self._sql.get("EXIST_CHILD", self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, parent_oid, id)
if result is None:
return False
else:
return True
async def len(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get("NUM_CHILDREN", self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql, oid)
return result
async def items(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get("GET_CHILDREN", self._objects_table_name)
async for record in conn.cursor(sql, oid):
# locks are dangerous in cursors since comsuming code might do
# sub-queries and they you end up with a deadlock
yield record
async def get_annotation(self, txn, oid, id):
sql = self._sql.get("GET_ANNOTATION", self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, oid, id, prepare=True)
if result is not None and result["parent_id"] == TRASHED_ID:
result = None
return result
async def get_annotation_keys(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get("GET_ANNOTATIONS_KEYS", self._objects_table_name)
async with txn._lock:
result = await conn.fetch(sql, oid)
items = []
for item in result:
if item["parent_id"] != TRASHED_ID:
items.append(item)
return items
async def write_blob_chunk(self, txn, bid, oid, chunk_index, data):
sql = self._sql.get("HAS_OBJECT", self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, oid)
if result is None:
# check if we have a referenced ob, could be new and not in db yet.
# if so, create a stub for it here...
conn = await txn.get_connection()
async with txn._lock:
await conn.execute(
f"""INSERT INTO {self._objects_table_name}
(zoid, tid, state_size, part, resource, type)
VALUES ($1::varchar({MAX_UID_LENGTH}), -1, 0, 0, TRUE, 'stub')""",
oid,
)
conn = await txn.get_connection()
sql = self._sql.get("INSERT_BLOB_CHUNK", self._blobs_table_name)
async with txn._lock:
return await conn.execute(sql, bid, oid, chunk_index, data)
async def read_blob_chunk(self, txn, bid, chunk=0):
sql = self._sql.get("READ_BLOB_CHUNK", self._blobs_table_name)
async with txn._lock:
return await self.get_one_row(txn, sql, bid, chunk)
async def read_blob_chunks(self, txn, bid):
conn = await txn.get_connection()
async for record in conn.cursor(bid):
# locks are dangerous in cursors since comsuming code might do
# sub-queries and they you end up with a deadlock
yield record
async def del_blob(self, txn, bid):
conn = await txn.get_connection()
sql = self._sql.get("DELETE_BLOB", self._blobs_table_name)
async with txn._lock:
await conn.execute(sql, bid)
async def get_total_number_of_objects(self, txn):
conn = await txn.get_connection()
sql = self._sql.get("NUM_ROWS", self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql)
return result
async def get_total_number_of_resources(self, txn):
conn = await txn.get_connection()
sql = self._sql.get("NUM_RESOURCES", self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql)
return result
async def get_total_resources_of_type(self, txn, type_):
conn = await txn.get_connection()
sql = self._sql.get("NUM_RESOURCES_BY_TYPE", self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql, type_)
return result
# Massive treatment without security
async def _get_page_resources_of_type(self, txn, type_, page, page_size):
conn = await txn.get_connection()
async with txn._lock:
keys = []
sql = self._sql.get("RESOURCES_BY_TYPE", self._objects_table_name)
for record in await conn.fetch(sql, type_, page_size, (page - 1) * page_size):
keys.append(record)
return keys
class PostgresqlStorage(BaseStorage):
"""Storage to a relational database, based on invalidation polling"""
_dsn = None
_partition_class = None
_large_record_size = 1 << 24
_vacuum_class = PGVacuum
_objects_table_name = 'objects'
_blobs_table_name = 'blobs'
_object_schema = {
'zoid': f'VARCHAR({MAX_OID_LENGTH}) NOT NULL PRIMARY KEY',
'tid': 'BIGINT NOT NULL',
'state_size': 'BIGINT NOT NULL',
'part': 'BIGINT NOT NULL',
'resource': 'BOOLEAN NOT NULL',
'of': f'VARCHAR({MAX_OID_LENGTH}) REFERENCES {{objects_table_name}} ON DELETE CASCADE',
'otid': 'BIGINT',
'parent_id': f'VARCHAR({MAX_OID_LENGTH}) REFERENCES {{objects_table_name}} ON DELETE CASCADE', # noqa
'id': 'TEXT',
'type': 'TEXT NOT NULL',
'json': 'JSONB',
'state': 'BYTEA'
}
_blob_schema = {
'bid': f'VARCHAR({MAX_OID_LENGTH}) NOT NULL',
'zoid': f'VARCHAR({MAX_OID_LENGTH}) NOT NULL REFERENCES {{objects_table_name}} ON DELETE CASCADE',
'chunk_index': 'INT NOT NULL',
'data': 'BYTEA'
}
_initialize_statements = [
'CREATE INDEX IF NOT EXISTS {object_table_name}_tid ON {objects_table_name} (tid);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_of ON {objects_table_name} (of);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_part ON {objects_table_name} (part);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_parent ON {objects_table_name} (parent_id);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_id ON {objects_table_name} (id);',
'CREATE INDEX IF NOT EXISTS {object_table_name}_type ON {objects_table_name} (type);',
'CREATE INDEX IF NOT EXISTS {blob_table_name}_bid ON {blobs_table_name} (bid);',
'CREATE INDEX IF NOT EXISTS {blob_table_name}_zoid ON {blobs_table_name} (zoid);',
'CREATE INDEX IF NOT EXISTS {blob_table_name}_chunk ON {blobs_table_name} (chunk_index);',
'CREATE SEQUENCE IF NOT EXISTS {schema}.tid_sequence;'
]
_unique_constraint = """CREATE UNIQUE INDEX CONCURRENTLY {constraint_name}_parent_id_id_key
ON {objects_table_name} (parent_id, id)
WHERE parent_id != '{TRASHED_ID}'"""
def __init__(self, dsn=None, partition=None, read_only=False, name=None,
pool_size=13, transaction_strategy='resolve_readcommitted',
conn_acquire_timeout=20, cache_strategy='dummy', db_schema='public',
objects_table_name='objects', blobs_table_name='blobs',
connection_manager=None, autovacuum=True, **options):
super(PostgresqlStorage, self).__init__(
read_only, transaction_strategy=transaction_strategy,
cache_strategy=cache_strategy)
self._dsn = dsn
self._pool_size = pool_size
self._partition_class = partition
self._read_only = read_only
self.__name__ = name
self._conn_acquire_timeout = conn_acquire_timeout
self._options = options
self._connection_options = {}
self._connection_initialized_on = time.time()
self._db_schema = db_schema
self._objects_table_name = f'{db_schema}.{objects_table_name}'
self._blobs_table_name = f'{db_schema}.{blobs_table_name}'
self._sql = SQLStatements()
self._connection_manager = connection_manager
self._autovacuum = autovacuum
async def finalize(self):
await self._connection_manager.close()
@property
def read_conn(self):
return self._connection_manager.read_conn
@property
def pool(self):
return self._connection_manager.pool
@property
def connection_manager(self):
return self._connection_manager
@property
def lock(self):
return self._connection_manager.lock
async def vacuum(self):
await self.connection_manager.vacuum.run(self._objects_table_name)
while self.connection_manager.vacuum.size > 0:
await asyncio.sleep(0.1)
async def create(self):
# Check DB
log.info('Creating initial database objects')
statements = []
if self._db_schema and self._db_schema != 'public':
statements.extend([f'CREATE SCHEMA IF NOT EXISTS {self._db_schema}'])
statements.extend([
get_table_definition(self._objects_table_name, self._object_schema),
get_table_definition(self._blobs_table_name, self._blob_schema,
primary_keys=('bid', 'zoid', 'chunk_index'))
])
statements.extend(self._initialize_statements)
for statement in statements:
otable_name = clear_table_name(self._objects_table_name)
if otable_name == 'objects':
otable_name = 'object'
btable_name = clear_table_name(self._blobs_table_name)
if btable_name == 'blobs':
btable_name = 'blob'
statement = statement.format(
objects_table_name=self._objects_table_name,
blobs_table_name=self._blobs_table_name,
# singular, index names
object_table_name=otable_name,
blob_table_name=btable_name,
schema=self._db_schema
)
try:
await self.read_conn.execute(statement)
except asyncpg.exceptions.UniqueViolationError:
# this is okay on creation, means 2 getting created at same time
pass
await self.initialize_tid_statements()
async def restart_connection(self, timeout=0.1):
log.error('Connection potentially lost to pg, restarting')
await self._connection_manager.restart()
await self.initialize_tid_statements()
self._connection_initialized_on = time.time()
raise ConflictError('Restarting connection to postgresql')
async def has_unique_constraint(self):
table_name = clear_table_name(self._objects_table_name)
result = await self.read_conn.fetch('''
SELECT * FROM pg_indexes
WHERE tablename = '{}' AND indexname = '{}_parent_id_id_key';
'''.format(table_name, table_name))
return len(result) > 0
async def initialize(self, loop=None, **kw):
self._connection_options = kw
if self._connection_manager is None:
self._connection_manager = PGConnectionManager(
dsn=self._dsn, pool_size=self._pool_size,
connection_options=self._connection_options,
conn_acquire_timeout=self._conn_acquire_timeout,
vacuum_class=self._vacuum_class)
await self._connection_manager.initialize(loop, **kw)
async with self.lock:
if await self.has_unique_constraint():
self._supports_unique_constraints = True
trash_sql = self._sql.get('CREATE_TRASH', self._objects_table_name)
try:
await self.initialize_tid_statements()
await self.read_conn.execute(trash_sql)
except asyncpg.exceptions.ReadOnlySQLTransactionError:
# Not necessary for read-only pg
pass
except (asyncpg.exceptions.UndefinedTableError,
asyncpg.exceptions.InvalidSchemaNameError):
await self.create()
# only available on new databases
await self.read_conn.execute(self._unique_constraint.format(
objects_table_name=self._objects_table_name,
constraint_name=clear_table_name(self._objects_table_name),
TRASHED_ID=TRASHED_ID
))
self._supports_unique_constraints = True
await self.initialize_tid_statements()
await self.read_conn.execute(trash_sql)
# migrate to larger VARCHAR size...
result = await self.read_conn.fetch("""
select * from information_schema.columns
where table_name='{}'""".format(self._objects_table_name))
if len(result) > 0 and result[0]['character_maximum_length'] != MAX_OID_LENGTH:
log.warn('Migrating VARCHAR key length')
await self.read_conn.execute(f'''
ALTER TABLE {self._objects_table_name} ALTER COLUMN zoid TYPE varchar({MAX_OID_LENGTH})''')
await self.read_conn.execute(f'''
ALTER TABLE {self._objects_table_name} ALTER COLUMN of TYPE varchar({MAX_OID_LENGTH})''')
await self.read_conn.execute(f'''
ALTER TABLE {self._objects_table_name} ALTER COLUMN parent_id TYPE varchar({MAX_OID_LENGTH})''')
await self.read_conn.execute(f'''
ALTER TABLE {self._blobs_table_name} ALTER COLUMN bid TYPE varchar({MAX_OID_LENGTH})''')
await self.read_conn.execute(f'''
ALTER TABLE {self._blobs_table_name} ALTER COLUMN zoid TYPE varchar({MAX_OID_LENGTH})''')
self._connection_initialized_on = time.time()
async def initialize_tid_statements(self):
self._stmt_next_tid = await self.read_conn.prepare(
NEXT_TID.format(schema=self._db_schema))
self._stmt_max_tid = await self.read_conn.prepare(
MAX_TID.format(schema=self._db_schema))
async def remove(self):
"""Reset the tables"""
async with self.pool.acquire() as conn:
await conn.execute("DROP TABLE IF EXISTS {};".format(self._blobs_table_name))
await conn.execute("DROP TABLE IF EXISTS {};".format(self._objects_table_name))
@restart_conn_on_exception
async def open(self):
conn = await self.pool.acquire(timeout=self._conn_acquire_timeout)
return conn
async def close(self, con):
try:
await shield(
asyncio.wait_for(self.pool.release(con, timeout=1), 1))
except (asyncio.CancelledError, RuntimeError, asyncio.TimeoutError,
asyncpg.exceptions.ConnectionDoesNotExistError):
pass
async def terminate(self, conn):
conn.terminate()
async def load(self, txn, oid):
sql = self._sql.get('GET_OID', self._objects_table_name)
async with txn._lock:
objects = await self.get_one_row(txn, sql, oid)
if objects is None:
raise KeyError(oid)
return objects
@profilable
async def store(self, oid, old_serial, writer, obj, txn):
assert oid is not None
pickled = writer.serialize() # This calls __getstate__ of obj
if len(pickled) >= self._large_record_size:
log.info(f"Large object {obj.__class__}: {len(pickled)}")
json_dict = await writer.get_json()
json = ujson.dumps(json_dict)
part = writer.part
if part is None:
part = 0
update = False
statement_sql = self._sql.get('NAIVE_UPSERT', self._objects_table_name)
if not obj.__new_marker__ and obj._p_serial is not None:
# we should be confident this is an object update
statement_sql = self._sql.get('UPDATE', self._objects_table_name)
update = True
conn = await txn.get_connection()
async with txn._lock:
try:
result = await conn.fetch(
statement_sql,
oid, # The OID of the object
txn._tid, # Our TID
len(pickled), # Len of the object
part, # Partition indicator
writer.resource, # Is a resource ?
writer.of, # It belogs to a main
old_serial, # Old serial
writer.parent_id, # Parent OID
writer.id, # Traversal ID
writer.type, # Guillotina type
json, # JSON catalog
pickled # Pickle state)
)
except asyncpg.exceptions.UniqueViolationError as ex:
if 'Key (parent_id, id)' in ex.detail:
raise ConflictIdOnContainer(ex)
raise
except asyncpg.exceptions.ForeignKeyViolationError:
txn.deleted[obj._p_oid] = obj
raise TIDConflictError(
f'Bad value inserting into database that could be caused '
f'by a bad cache value. This should resolve on request retry.',
oid, txn, old_serial, writer)
except asyncpg.exceptions._base.InterfaceError as ex:
if 'another operation is in progress' in ex.args[0]:
raise ConflictError(
f'asyncpg error, another operation in progress.',
oid, txn, old_serial, writer)
raise
except asyncpg.exceptions.DeadlockDetectedError:
raise ConflictError(f'Deadlock detected.',
oid, txn, old_serial, writer)
if len(result) != 1 or result[0]['count'] != 1:
if update:
# raise tid conflict error
raise TIDConflictError(
f'Mismatch of tid of object being updated. This is likely '
f'caused by a cache invalidation race condition and should '
f'be an edge case. This should resolve on request retry.',
oid, txn, old_serial, writer)
else:
log.error('Incorrect response count from database update. '
'This should not happen. tid: {}'.format(txn._tid))
await txn._cache.store_object(obj, pickled)
async def _txn_oid_commit_hook(self, status, oid):
await self._connection_manager._vacuum.add_to_queue(oid, self._objects_table_name)
async def delete(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('TRASH_PARENT_ID', self._objects_table_name)
async with txn._lock:
# for delete, we reassign the parent id and delete in the vacuum task
try:
await conn.execute(sql, oid)
except asyncpg.exceptions.UniqueViolationError:
# we already have ob with id, it has not been vacuumed yet,
# fallback.
sql = self._sql.get('DELETE_OBJECT', self._objects_table_name)
await conn.execute(sql, oid)
if self._autovacuum:
txn.add_after_commit_hook(self._txn_oid_commit_hook, oid)
async def _check_bad_connection(self, ex):
# we do not use transaction lock here but a storage lock because
# a storage object has a shard conn for reads
for err in ('connection is closed', 'pool is closed', 'connection was closed'):
if err in str(ex):
if (time.time() - self._connection_initialized_on) > BAD_CONNECTION_RESTART_DELAY:
# we need to make sure we aren't calling this over and over again
async with self.lock:
return await self.restart_connection()
@restart_conn_on_exception
async def get_next_tid(self, txn):
async with self.lock:
return await self._stmt_next_tid.fetchval()
@restart_conn_on_exception
async def get_current_tid(self, txn):
async with self.lock:
return await self._stmt_max_tid.fetchval()
async def get_one_row(self, txn, sql, *args, prepare=False):
conn = await txn.get_connection()
# Helper function to provide easy adaptation to cockroach
if prepare:
# latest version of asyncpg has prepare bypassing statement cache
smt = await conn.prepare(sql)
return await smt.fetchrow(*args)
else:
return await conn.fetchrow(sql, *args)
def _db_transaction_factory(self, txn):
# make sure asycpg knows this is a new transaction
if txn._db_conn._con is not None:
txn._db_conn._con._top_xact = None
return txn._db_conn.transaction(readonly=txn._manager._storage._read_only)
@restart_conn_on_exception
async def _async_db_transaction_factory(self, txn):
return self._db_transaction_factory(txn)
async def start_transaction(self, txn, retries=0):
error = None
conn = await txn.get_connection()
async with txn._lock:
txn._db_txn = await self._async_db_transaction_factory(txn)
try:
await txn._db_txn.start()
return
except (asyncpg.exceptions.InterfaceError,
asyncpg.exceptions.InternalServerError) as ex:
error = ex
if error is not None:
if retries > 2:
raise error # pylint: disable=E0702
restart = rollback = False
if isinstance(error, asyncpg.exceptions.InternalServerError):
restart = True
if error.sqlstate == 'XX000':
rollback = True
elif ('manually started transaction' in error.args[0] or
'connection is closed' in error.args[0]):
restart = True
if 'manually started transaction' in error.args[0]:
rollback = True
if rollback:
try:
# thinks we're manually in txn, manually rollback and try again...
await conn.execute('ROLLBACK;')
except asyncpg.exceptions._base.InterfaceError:
# we're okay with this error here...
pass
if restart:
await self.close(conn)
txn._db_conn = await self.open()
return await self.start_transaction(txn, retries + 1)
async def get_conflicts(self, txn):
async with self.lock:
if len(txn.modified) == 0:
return []
# use storage lock instead of transaction lock
if len(txn.modified) < 1000:
# if it's too large, we're not going to check on object ids
modified_oids = [k for k in txn.modified.keys()]
sql = self._sql.get('TXN_CONFLICTS_ON_OIDS', self._objects_table_name)
return await self.read_conn.fetch(
sql, txn._tid, modified_oids)
else:
sql = self._sql.get('TXN_CONFLICTS', self._objects_table_name)
return await self.read_conn.fetch(sql, txn._tid)
async def commit(self, transaction):
if transaction._db_txn is not None:
async with transaction._lock:
await transaction._db_txn.commit()
elif (self._transaction_strategy not in ('none', 'tidonly') and
not transaction._skip_commit):
log.warning('Do not have db transaction to commit')
return transaction._tid
async def abort(self, transaction):
if transaction._db_txn is not None:
async with transaction._lock:
try:
await transaction._db_txn.rollback()
except asyncpg.exceptions._base.InterfaceError:
# we're okay with this error here...
pass
# reads don't need transaction necessarily so don't log
# else:
# log.warning('Do not have db transaction to rollback')
# Introspection
async def get_page_of_keys(self, txn, oid, page=1, page_size=1000):
conn = await txn.get_connection()
keys = []
sql = self._sql.get('BATCHED_GET_CHILDREN_KEYS', self._objects_table_name)
for record in await conn.fetch(
sql, oid, page_size, (page - 1) * page_size):
keys.append(record['id'])
return keys
async def keys(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('GET_CHILDREN_KEYS', self._objects_table_name)
async with txn._lock:
result = await conn.fetch(sql, oid)
return result
async def get_child(self, txn, parent_oid, id):
sql = self._sql.get('GET_CHILD', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, parent_oid, id)
return result
async def get_children(self, txn, parent_oid, ids):
conn = await txn.get_connection()
sql = self._sql.get('GET_CHILDREN_BATCH', self._objects_table_name)
async with txn._lock:
return await conn.fetch(sql, parent_oid, ids)
async def has_key(self, txn, parent_oid, id):
sql = self._sql.get('EXIST_CHILD', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, parent_oid, id)
if result is None:
return False
else:
return True
async def len(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('NUM_CHILDREN', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql, oid)
return result
async def items(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('GET_CHILDREN', self._objects_table_name)
async for record in conn.cursor(sql, oid):
# locks are dangerous in cursors since comsuming code might do
# sub-queries and they you end up with a deadlock
yield record
async def get_annotation(self, txn, oid, id):
sql = self._sql.get('GET_ANNOTATION', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, oid, id, prepare=True)
if result is not None and result['parent_id'] == TRASHED_ID:
result = None
return result
async def get_annotation_keys(self, txn, oid):
conn = await txn.get_connection()
sql = self._sql.get('GET_ANNOTATIONS_KEYS', self._objects_table_name)
async with txn._lock:
result = await conn.fetch(sql, oid)
items = []
for item in result:
if item['parent_id'] != TRASHED_ID:
items.append(item)
return items
async def write_blob_chunk(self, txn, bid, oid, chunk_index, data):
sql = self._sql.get('HAS_OBJECT', self._objects_table_name)
async with txn._lock:
result = await self.get_one_row(txn, sql, oid)
if result is None:
# check if we have a referenced ob, could be new and not in db yet.
# if so, create a stub for it here...
conn = await txn.get_connection()
async with txn._lock:
await conn.execute(f'''INSERT INTO {self._objects_table_name}
(zoid, tid, state_size, part, resource, type)
VALUES ($1::varchar({MAX_OID_LENGTH}), -1, 0, 0, TRUE, 'stub')''', oid)
conn = await txn.get_connection()
sql = self._sql.get('INSERT_BLOB_CHUNK', self._blobs_table_name)
async with txn._lock:
return await conn.execute(
sql, bid, oid, chunk_index, data)
async def read_blob_chunk(self, txn, bid, chunk=0):
sql = self._sql.get('READ_BLOB_CHUNK', self._blobs_table_name)
async with txn._lock:
return await self.get_one_row(txn, sql, bid, chunk)
async def read_blob_chunks(self, txn, bid):
conn = await txn.get_connection()
async for record in conn.cursor(bid):
# locks are dangerous in cursors since comsuming code might do
# sub-queries and they you end up with a deadlock
yield record
async def del_blob(self, txn, bid):
conn = await txn.get_connection()
sql = self._sql.get('DELETE_BLOB', self._blobs_table_name)
async with txn._lock:
await conn.execute(sql, bid)
async def get_total_number_of_objects(self, txn):
conn = await txn.get_connection()
sql = self._sql.get('NUM_ROWS', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql)
return result
async def get_total_number_of_resources(self, txn):
conn = await txn.get_connection()
sql = self._sql.get('NUM_RESOURCES', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql)
return result
async def get_total_resources_of_type(self, txn, type_):
conn = await txn.get_connection()
sql = self._sql.get('NUM_RESOURCES_BY_TYPE', self._objects_table_name)
async with txn._lock:
result = await conn.fetchval(sql, type_)
return result
# Massive treatment without security
async def _get_page_resources_of_type(self, txn, type_, page, page_size):
conn = await txn.get_connection()
async with txn._lock:
keys = []
sql = self._sql.get('RESOURCES_BY_TYPE', self._objects_table_name)
for record in await conn.fetch(
sql, type_, page_size, (page - 1) * page_size):
keys.append(record)
return keys
class PGVacuum:
def __init__(self, manager, loop):
self._manager = manager
self._loop = loop
self._queue = asyncio.Queue(loop=loop)
self._closed = False
self._active = False
self._sql = SQLStatements()
async def initialize(self):
while not self._closed:
try:
await self._initialize()
except (concurrent.futures.CancelledError, RuntimeError):
# we're okay with the task getting cancelled
return
async def _initialize(self):
while not self._closed:
oid = None
try:
oid, table_name = await self._queue.get()
self._active = True
await shield(self.vacuum(oid, table_name))
except (concurrent.futures.CancelledError, RuntimeError):
raise
except Exception:
log.warning(f'Error vacuuming oid {oid}', exc_info=True)
finally:
self._active = False
try:
self._queue.task_done()
except ValueError:
pass
async def run(self, table_name):
'''
get existing trashed objects, push them on the queue...
there might be contention, but that is okay
'''
async with self._manager.pool.acquire(
timeout=self._manager._conn_acquire_timeout) as conn:
try:
sql = self._sql.get('GET_TRASHED_OBJECTS', table_name)
for record in await conn.fetch(sql):
self._queue.put_nowait((record['zoid'], table_name))
except concurrent.futures.TimeoutError:
log.info('Timed out connecting to storage')
except Exception:
log.warning('Error deleting trashed object', exc_info=True)
async def add_to_queue(self, oid, table_name):
if self._closed:
raise Exception('Closing down')
await self._queue.put((oid, table_name))
async def vacuum(self, oid, table_name):
'''
DELETED objects has parent id changed to the trashed ob for the oid...
'''
async with self._manager.pool.acquire(
timeout=self._manager._conn_acquire_timeout) as conn:
sql = self._sql.get('DELETE_OBJECT', table_name)
try:
await conn.execute(sql, oid)
except Exception:
log.warning('Error deleting trashed object', exc_info=True)
async def finalize(self):
self._closed = True
try:
await asyncio.wait_for(self._queue.join(), 2)
except asyncio.TimeoutError:
pass
@property
def size(self):
return self._queue.qsize()
