117 lines
3.3 KiB
Python
117 lines
3.3 KiB
Python
import aiohttp
|
|
|
|
ENABLED = False
|
|
|
|
import aiohttp
|
|
import asyncio
|
|
from aiohttp import web
|
|
|
|
import sqlite3
|
|
|
|
import dataset
|
|
from sqlalchemy import event
|
|
from sqlalchemy.engine import Engine
|
|
|
|
import json
|
|
|
|
queue = asyncio.Queue()
|
|
|
|
class State:
|
|
do_not_sync = False
|
|
|
|
async def sync_service(app):
|
|
if not ENABLED:
|
|
return
|
|
session = aiohttp.ClientSession()
|
|
async with session.ws_connect('http://localhost:3131/ws') as ws:
|
|
async def receive():
|
|
|
|
queries_synced = 0
|
|
async for msg in ws:
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
try:
|
|
data = json.loads(msg.data)
|
|
State.do_not_sync = True
|
|
app.db.execute(*data)
|
|
app.db.commit()
|
|
State.do_not_sync = False
|
|
queries_synced += 1
|
|
print("queries synced: " + str(queries_synced))
|
|
print(*data)
|
|
await app.services.socket.broadcast_event()
|
|
except Exception as e:
|
|
print(e)
|
|
pass
|
|
#print(f"Received: {msg.data}")
|
|
elif msg.type == aiohttp.WSMsgType.ERROR:
|
|
break
|
|
async def write():
|
|
while True:
|
|
msg = await queue.get()
|
|
await ws.send_str(json.dumps(msg,default=str))
|
|
queue.task_done()
|
|
|
|
await asyncio.gather(receive(), write())
|
|
|
|
await session.close()
|
|
|
|
queries_queued = 0
|
|
# Attach a listener to log all executed statements
|
|
@event.listens_for(Engine, "before_cursor_execute")
|
|
def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
|
|
if not ENABLED:
|
|
return
|
|
global queries_queued
|
|
if State.do_not_sync:
|
|
print(statement,parameters)
|
|
return
|
|
if statement.startswith("SELECT"):
|
|
return
|
|
queue.put_nowait((statement, parameters))
|
|
queries_queued += 1
|
|
print("Queries queued: " + str(queries_queued))
|
|
|
|
async def websocket_handler(request):
|
|
queries_broadcasted = 0
|
|
ws = web.WebSocketResponse()
|
|
await ws.prepare(request)
|
|
request.app['websockets'].append(ws)
|
|
async for msg in ws:
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
for client in request.app['websockets']:
|
|
if client != ws:
|
|
await client.send_str(msg.data)
|
|
cursor = request.app['db'].cursor()
|
|
data = json.loads(msg.data)
|
|
queries_broadcasted += 1
|
|
|
|
cursor.execute(*data)
|
|
cursor.close()
|
|
print("Queries broadcasted: " + str(queries_broadcasted))
|
|
elif msg.type == aiohttp.WSMsgType.ERROR:
|
|
print(f'WebSocket connection closed with exception {ws.exception()}')
|
|
|
|
request.app['websockets'].remove(ws)
|
|
return ws
|
|
|
|
app = web.Application()
|
|
app['websockets'] = []
|
|
|
|
app.router.add_get('/ws', websocket_handler)
|
|
|
|
async def on_startup(app):
|
|
app['db'] = sqlite3.connect('snek.db')
|
|
print("Server starting...")
|
|
|
|
async def on_cleanup(app):
|
|
for ws in app['websockets']:
|
|
await ws.close()
|
|
app['db'].close()
|
|
|
|
app.on_startup.append(on_startup)
|
|
app.on_cleanup.append(on_cleanup)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
web.run_app(app, host='127.0.0.1', port=3131)
|