def __init__(self, app, id=None, config=None):
super().__init__(app, id=id, config=config)
self.ConnectionEvent = asyncio.Event(loop=app.Loop)
self.ConnectionEvent.clear()
self.PubSub = PubSub(app)
self.Loop = app.Loop
self._host = self.Config['host']
self._port = int(self.Config['port'])
self._user = self.Config['user']
self._password = self.Config['password']
self._connect_timeout = int(self.Config['connect_timeout'])
self._db = self.Config['db']
self._reconnect_delay = self.Config['reconnect_delay']
self._output_queue_max_size = self.Config['output_queue_max_size']
self._max_bulk_size = int(self.Config['max_bulk_size'])
self._conn_future = None
self._connection_request = False
self._pause = False
self._throttled_by_error = False
# Subscription
self._on_health_check('connection.open!')
app.PubSub.subscribe("Application.stop!", self._on_application_stop)
app.PubSub.subscribe("Application.tick!", self._on_health_check)
app.PubSub.subscribe("MySQLConnection.pause!", self._on_pause)
app.PubSub.subscribe("MySQLConnection.unpause!", self._on_unpause)
self._output_queue = asyncio.Queue(loop=app.Loop)
self._bulks = {} # We have a "bulk" per query
def __init__(self, app, connection_id, config=None):
super().__init__(app, connection_id, config=config)
self.ConnectionEvent = asyncio.Event(loop=app.Loop)
self.ConnectionEvent.clear()
self.PubSub = PubSub(app)
self.Loop = app.Loop
self._host = self.Config['host']
self._port = self.Config['port']
self._user = self.Config['user']
self._password = self.Config['password']
self._connect_timeout = self.Config['connect_timeout']
self._db = self.Config['db']
self._reconnect_delay = self.Config['reconnect_delay']
self._output_queue_max_size = self.Config['output_queue_max_size']
self._conn_future = None
self._connection_request = False
# Subscription
self._on_health_check('connection.open!')
app.PubSub.subscribe("Application.stop!", self._on_application_stop)
app.PubSub.subscribe("Application.tick!", self._on_health_check)
self._output_queue = asyncio.Queue(loop=app.Loop)
def __init__(self, app, id=None, config=None): super().__init__(app, id=id, config=config) self.Connection = None self.ConnectionEvent = asyncio.Event(loop=app.Loop) self.ConnectionEvent.clear() self.PubSub = PubSub(app) self.Loop = app.Loop self._reconnect()
class MySQLConnection(Connection):
ConfigDefaults = {
'host': 'localhost',
'port': 3306,
'user': '',
'password': '',
'db': '',
'connect_timeout': 1,
'reconnect_delay': 5.0,
'output_queue_max_size': 3,
}
def __init__(self, app, connection_id, config=None):
super().__init__(app, connection_id, config=config)
self.ConnectionEvent = asyncio.Event(loop=app.Loop)
self.ConnectionEvent.clear()
self.PubSub = PubSub(app)
self.Loop = app.Loop
self._host = self.Config['host']
self._port = self.Config['port']
self._user = self.Config['user']
self._password = self.Config['password']
self._connect_timeout = self.Config['connect_timeout']
self._db = self.Config['db']
self._reconnect_delay = self.Config['reconnect_delay']
self._output_queue_max_size = self.Config['output_queue_max_size']
self._conn_future = None
self._connection_request = False
# Subscription
self._on_health_check('connection.open!')
app.PubSub.subscribe("Application.stop!", self._on_application_stop)
app.PubSub.subscribe("Application.tick!", self._on_health_check)
self._output_queue = asyncio.Queue(loop=app.Loop)
def _on_application_stop(self, message_type, counter):
self._output_queue.put_nowait(None)
def _on_health_check(self, message_type):
if self._conn_future is not None:
# Connection future exists
if not self._conn_future.done():
# Connection future didn't result yet
# No sanitization needed
return
try:
self._conn_future.result()
except:
# Connection future threw an error
L.exception("Unexpected connection future error")
# Connection future already resulted (with or without exception)
self._conn_future = None
assert (self._conn_future is None)
self._conn_future = asyncio.ensure_future(self._connection(),
loop=self.Loop)
async def _connection(self):
try:
async with create_pool(
host=self._host,
port=self._port,
user=self._user,
password=self._password,
db=self._db,
connect_timeout=self.
_connect_timeout, #Doesn't work! See socket.timeout exception below
loop=self.Loop) as pool:
self._conn_pool = pool
self.ConnectionEvent.set()
await self._loader()
except socket.timeout:
# Socket timeout not implemented in aiomysql as it sets a keepalive to the connection
# it has been placed as an issue on GitHub: https://github.com/aio-libs/aiomysql/issues/257
L.exception("MySQL connection timeout")
pass
except BaseException:
L.exception("Unexpected MySQL connection error")
raise
def acquire(self):
assert (self._conn_pool is not None)
return self._conn_pool.acquire()
def consume(self, query):
self._output_queue.put_nowait(query)
if self._output_queue.qsize() == self._output_queue_max_size:
self.PubSub.publish("MySQLConnection.pause!", self)
async def _loader(self):
while True:
query = await self._output_queue.get()
if query is None:
break
if self._output_queue.qsize() == self._output_queue_max_size - 1:
self.PubSub.publish("MySQLConnection.unpause!",
self,
asynchronously=True)
try:
async with self.acquire() as conn:
try:
async with conn.cursor() as cur:
await cur.execute(query)
await conn.commit()
except BaseException as e:
L.exception(
"Unexpected error when processing MySQL query.")
except BaseBaseException as e:
L.exception("Couldn't acquire connection")
class MySQLConnection(Connection):
"""
MySQLConnection is a top BitSwan object, that is used to connect the application to an external MySQL database.
MySQLConnection is thus used in MySQL-related sources, sinks, processors and lookups such as MySQLLookup.
MySQLConnection is built on top of aiomysql library and utilizes its functions: https://aiomysql.readthedocs.io/en/latest/
The following code illustrates how to create and register the MySQL connection inside the application object.
The connection's ID is set to "MySQLConnection" (the second argument of the connection's constructor),
which is also then used as a section name in configuration files.
.. code:: python
app = bspump.BSPumpApplication()
svc = app.get_service("bspump.PumpService")
svc.add_connection(
bspump.mysql.MySQLConnection(app, "MySQLConnection")
)
In order to locate MySQLConnection in the constructor method of sources, sinks, processors etc., the following code can be used:
.. code:: python
svc = app.get_service("bspump.PumpService")
connection = svc.locate_connection("MySQLConnection")
"""
ConfigDefaults = {
'host': 'localhost',
'port': 3306,
'user': '',
'password': '',
'db': '',
'connect_timeout': 10,
'reconnect_delay': 5.0,
'output_queue_max_size': 10,
'max_bulk_size': 2,
}
RetryErrors = frozenset([
1047, 1053, 1077, 1078, 1079, 1080, 2003, 2006, 2012, 2013, 1152, 1205,
1213, 1223
])
def __init__(self, app, id=None, config=None):
super().__init__(app, id=id, config=config)
self.ConnectionEvent = asyncio.Event(loop=app.Loop)
self.ConnectionEvent.clear()
self.PubSub = PubSub(app)
self.Loop = app.Loop
self._host = self.Config['host']
self._port = int(self.Config['port'])
self._user = self.Config['user']
self._password = self.Config['password']
self._connect_timeout = int(self.Config['connect_timeout'])
self._db = self.Config['db']
self._reconnect_delay = self.Config['reconnect_delay']
self._output_queue_max_size = self.Config['output_queue_max_size']
self._max_bulk_size = int(self.Config['max_bulk_size'])
self._conn_future = None
self._connection_request = False
self._pause = False
self._throttled_by_error = False
# Subscription
self._on_health_check('connection.open!')
app.PubSub.subscribe("Application.stop!", self._on_application_stop)
app.PubSub.subscribe("Application.tick!", self._on_health_check)
app.PubSub.subscribe("MySQLConnection.pause!", self._on_pause)
app.PubSub.subscribe("MySQLConnection.unpause!", self._on_unpause)
self._output_queue = asyncio.Queue(loop=app.Loop)
self._bulks = {} # We have a "bulk" per query
def _on_pause(self):
self._pause = True
def _on_unpause(self):
self._pause = False
def _flush(self):
for query in self._bulks.keys():
# Break if throttling was requested during the flush,
# so that put_nowait doesn't raise
if self._pause:
break
self._flush_bulk(query)
def _flush_bulk(self, query):
# Enqueue and throttle if needed
self._output_queue.put_nowait((query, self._bulks[query]))
if self._output_queue.qsize() == self._output_queue_max_size:
self.PubSub.publish("MySQLConnection.pause!", self)
# Reset bulk
self._bulks[query] = []
def _on_application_stop(self, message_type, counter):
self._flush()
self._output_queue.put_nowait((None, None))
def _on_health_check(self, message_type):
if self._conn_future is not None:
# Connection future exists
if not self._conn_future.done():
# Connection future didn't result yet
# No sanitization needed
return
try:
self._conn_future.result()
except Exception:
# Connection future threw an error
L.exception("Unexpected connection future error")
# Connection future already resulted (with or without exception)
self._conn_future = None
assert (self._conn_future is None)
self._conn_future = asyncio.ensure_future(self._async_connection(),
loop=self.Loop)
self._sync_connection()
async def _async_connection(self):
try:
async with aiomysql.create_pool(
host=self._host,
port=self._port,
user=self._user,
password=self._password,
db=self._db,
conv=convertors,
connect_timeout=self._connect_timeout,
loop=self.Loop) as pool:
self._conn_pool = pool
self.ConnectionEvent.set()
await self._loader()
except (pymysql.err.InternalError, pymysql.err.ProgrammingError,
pymysql.err.OperationalError) as e:
if e.args[0] in self.RetryErrors:
L.warn(
"Recoverable error '{}' occurred in MySQLConnection. Retrying."
.format(e.args[0]))
self._conn_future = None
return
L.exception("Unexpected MySQL connection error")
raise e
except BaseException:
L.exception("Unexpected MySQL connection error")
raise
def _sync_connection(self):
try:
connection = pymysql.connect(host=self._host,
port=self._port,
user=self._user,
password=self._password,
database=self._db,
conv=convertors,
connect_timeout=self._connect_timeout)
self._conn_sync = connection
except (pymysql.err.InternalError, pymysql.err.ProgrammingError,
pymysql.err.OperationalError) as e:
if e.args[0] in self.RetryErrors:
L.warn(
"Recoverable error '{}' occurred in MySQLConnection. Retrying."
.format(e.args[0]))
return
L.exception("Unexpected MySQL connection error")
raise e
except BaseException:
L.exception("Unexpected MySQL connection error")
raise
def acquire_connection(self) -> aiomysql.utils._PoolAcquireContextManager:
"""
Acquire asynchronous database connection
Use with `async with` statement
.. code-block:: python
async with self.Connection.acquire_connection() as connection:
async with connection.cursor() as cursor:
await cursor.execute(query)
:return: Asynchronous Context Manager
"""
assert (self._conn_pool is not None)
return self._conn_pool.acquire()
def acquire_sync_cursor(self) -> pymysql.cursors.DictCursor:
"""
Acquire synchronous database cursor
Use with `with` statement
.. code-block:: python
with self.Connection.acquire_sync_cursor() as cursor:
await cursor.execute(query)
:return: Context Manager
"""
assert (self._conn_sync is not None)
return pymysql.cursors.DictCursor(self._conn_sync)
def consume(self, query, data):
# Create a bulk for this query if doesn't yet exist
if query not in self._bulks:
self._bulks[query] = []
# Add data to the query's bulk
self._bulks[query].append(data)
# Flush on _max_bulk_size
if len(self._bulks[query]) >= self._max_bulk_size:
self._flush_bulk(query)
async def _loader(self):
while True:
query, data = await self._output_queue.get()
if query is None:
break
if self._output_queue.qsize() == self._output_queue_max_size - 1:
self.PubSub.publish("MySQLConnection.unpause!",
self,
asynchronously=True)
async with self.acquire_connection() as connection:
async with connection.cursor() as cursor:
try:
await cursor.executemany(query, data)
await connection.commit()
if self._throttled_by_error:
self.Pipeline.throttle(self, False)
self._throttled_by_error = False
except (pymysql.err.InternalError,
pymysql.err.ProgrammingError,
pymysql.err.OperationalError) as e:
if e.args[0] in self.RetryErrors:
L.warn(
"Recoverable error '{}' occurred in MySQLConnection. Retrying."
.format(e.args[0]))
self._output_queue.put_nowait((query, data))
self.Pipeline.throttle(self, True)
self._throttled_by_error = True
raise e
class AMQPConnection(Connection):
ConfigDefaults = {
'url': 'amqp://localhost/',
'appname': 'bspump.py',
'reconnect_delay': 10.0,
}
def __init__(self, app, id=None, config=None):
super().__init__(app, id=id, config=config)
self.Connection = None
self.ConnectionEvent = asyncio.Event(loop=app.Loop)
self.ConnectionEvent.clear()
self.PubSub = PubSub(app)
self.Loop = app.Loop
self._reconnect()
def _reconnect(self):
if self.Connection is not None:
if not (self.Connection.is_closing or self.Connection.is_closed):
self.Connection.close()
self.Connection = None
parameters = pika.URLParameters(self.Config['url'])
if parameters.client_properties is None:
parameters.client_properties = dict()
parameters.client_properties['application'] = self.Config['appname']
self.Connection = pika.adapters.asyncio_connection.AsyncioConnection(
parameters=parameters,
on_open_callback=self._on_connection_open,
on_open_error_callback=self._on_connection_open_error,
on_close_callback=self._on_connection_close
)
# Connection callbacks
def _on_connection_open(self, connection):
L.info("AMQP connected")
self.ConnectionEvent.set()
self.PubSub.publish("AMQPConnection.open!")
def _on_connection_close(self, connection, code, reason):
L.warning("AMQP disconnected ({}): {}".format(code, reason))
self.ConnectionEvent.clear()
self.PubSub.publish("AMQPConnection.close!")
self.Loop.call_later(float(self.Config['reconnect_delay']), self._reconnect)
def _on_connection_open_error(self, connection, error_message=None):
L.error("AMQP error: {}".format(error_message if error_message is not None else 'Generic error'))
self.ConnectionEvent.clear()
self.PubSub.publish("AMQPConnection.close!")
self.Loop.call_later(float(self.Config['reconnect_delay']), self._reconnect)
class ODBCConnection(Connection):
# Caution: Providing incorrect connection configuration terminates the program with 'Abort trap 6'
ConfigDefaults = {
'host': 'localhost',
'port': 3306,
'user': '',
'password': '',
'driver': '',
'db': '',
'connect_timeout': 1,
'reconnect_delay': 5.0,
'output_queue_max_size': 10,
'max_bulk_size': 2,
}
def __init__(self, app, id=None, config=None):
super().__init__(app, id=id, config=config)
self.ConnectionEvent = asyncio.Event(loop=app.Loop)
self.ConnectionEvent.clear()
self.PubSub = PubSub(app)
self.Loop = app.Loop
self._host = self.Config['host']
self._port = int(self.Config['port'])
self._user = self.Config['user']
self._password = self.Config['password']
self._connect_timeout = self.Config['connect_timeout']
self._dsn = "Driver={};Database={}".format(self.Config['driver'],
self.Config['db'])
self._reconnect_delay = self.Config['reconnect_delay']
self._output_queue_max_size = self.Config['output_queue_max_size']
self._max_bulk_size = int(self.Config['max_bulk_size'])
self._conn_future = None
self._connection_request = False
self._pause = False
# Subscription
self._on_health_check('connection.open!')
app.PubSub.subscribe("Application.stop!", self._on_application_stop)
app.PubSub.subscribe("Application.tick!", self._on_health_check)
app.PubSub.subscribe("ODBCConnection.pause!", self._on_pause)
app.PubSub.subscribe("ODBCConnection.unpause!", self._on_unpause)
self._output_queue = asyncio.Queue(loop=app.Loop)
self._bulks = {} # We have a "bulk" per query
def _on_pause(self):
self._pause = True
def _on_unpause(self):
self._pause = False
def _flush(self):
for query in self._bulks.keys():
# Break if throttling was requested during the flush,
# so that put_nowait doesn't raise
if self._pause:
break
self._flush_bulk(query)
def _flush_bulk(self, query):
# Enqueue and thorttle if needed
self._output_queue.put_nowait((query, self._bulks[query]))
if self._output_queue.qsize() == self._output_queue_max_size:
self.PubSub.publish("ODBCConnection.pause!", self)
# Reset bulk
self._bulks[query] = []
def _on_application_stop(self, message_type, counter):
self._flush()
self._output_queue.put_nowait((None, None))
def _on_health_check(self, message_type):
if self._conn_future is not None:
# Connection future exists
if not self._conn_future.done():
# Connection future didn't result yet
# No sanitization needed
return
try:
self._conn_future.result()
except Exception:
# Connection future threw an error
L.exception("Unexpected connection future error")
# Connection future already resulted (with or without exception)
self._conn_future = None
assert (self._conn_future is None)
self._conn_future = asyncio.ensure_future(self._connection(),
loop=self.Loop)
async def _connection(self):
try:
async with aioodbc.create_pool(
host=self._host,
port=self._port,
user=self._user,
password=self._password,
dsn=self._dsn,
connect_timeout=self._connect_timeout,
loop=self.Loop) as pool:
self._conn_pool = pool
self.ConnectionEvent.set()
await self._loader()
except BaseException:
L.exception("Unexpected ODBC connection error")
raise
def acquire(self):
assert (self._conn_pool is not None)
return self._conn_pool.acquire()
def consume(self, query, data):
# Create a bulk for this query if doesn't yet exist
if query not in self._bulks:
self._bulks[query] = []
# Add data to the query's bulk
self._bulks[query].append(data)
# Flush on _max_bulk_size
if len(self._bulks[query]) >= self._max_bulk_size:
self._flush_bulk(query)
async def _loader(self):
while True:
query, data = await self._output_queue.get()
if query is None:
break
if self._output_queue.qsize() == self._output_queue_max_size - 1:
self.PubSub.publish("ODBCConnection.unpause!",
self,
asynchronously=True)
async with self.acquire() as conn:
async with conn.cursor() as cur:
await cur.executemany(query, data)
await conn.commit()