1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
|
from dbutils.pooled_db import PooledDB
import threading
import time as _time
import persistqueue
from .sqlbase import SQLBase
from typing import Any, Optional
class MySQLQueue(SQLBase):
"""Mysql(or future standard dbms) based FIFO queue."""
_TABLE_NAME = 'queue'
_KEY_COLUMN = '_id' # the name of the key column, used in DB CRUD
# SQL to create a table
_SQL_CREATE = (
'CREATE TABLE IF NOT EXISTS {table_name} ('
'{key_column} INTEGER PRIMARY KEY AUTO_INCREMENT, '
'data BLOB, timestamp FLOAT)')
# SQL to insert a record
_SQL_INSERT = 'INSERT INTO {table_name} (data, timestamp) VALUES (%s, %s)'
# SQL to select a record
_SQL_SELECT_ID = (
'SELECT {key_column}, data, timestamp FROM {table_name} WHERE'
' {key_column} = {rowid}'
)
_SQL_SELECT = (
'SELECT {key_column}, data, timestamp FROM {table_name} '
'ORDER BY {key_column} ASC LIMIT 1'
)
_SQL_SELECT_WHERE = (
'SELECT {key_column}, data, timestamp FROM {table_name} WHERE'
' {column} {op} %s ORDER BY {key_column} ASC LIMIT 1 '
)
_SQL_UPDATE = 'UPDATE {table_name} SET data = %s WHERE {key_column} = %s'
_SQL_DELETE = 'DELETE FROM {table_name} WHERE {key_column} {op} %s'
def __init__(
self,
host: str,
user: str,
passwd: str,
db_name: str,
name: Optional[str] = None,
port: int = 3306,
charset: str = 'utf8mb4',
auto_commit: bool = True,
serializer: Any = persistqueue.serializers.pickle,
) -> None:
super(MySQLQueue, self).__init__()
self.name = name if name else "sql"
self.host = host
self.user = user
self.passwd = passwd
self.db_name = db_name
self.port = port
self.charset = charset
self._serializer = serializer
self.auto_commit = auto_commit
self.tran_lock = threading.Lock()
self.put_event = threading.Event()
self.action_lock = threading.Lock()
self._connection_pool = None
self._getter = None
self._putter = None
self._new_db_connection()
self._init()
def _new_db_connection(self) -> None:
try:
import pymysql
except ImportError:
print("Please install mysql library via 'pip install PyMySQL'")
raise
db_pool = PooledDB(pymysql, 2, 10, 5, 10, True,
host=self.host, port=self.port, user=self.user,
passwd=self.passwd, database=self.db_name,
charset=self.charset)
self._connection_pool = db_pool
conn = db_pool.connection()
cursor = conn.cursor()
cursor.execute("SELECT VERSION()")
_ = cursor.fetchone()
cursor.execute(self._sql_create)
conn.commit()
cursor.execute("use %s" % self.db_name)
self._putter = MySQLConn(queue=self)
self._getter = self._putter
def put(self, item: Any, block: bool = True) -> int:
# block kwarg is noop and only here to align with python's queue
obj = self._serializer.dumps(item)
_id = self._insert_into(obj, _time.time())
self.total += 1
self.put_event.set()
return _id
def put_nowait(self, item: Any) -> int:
return self.put(item, block=False)
def _init(self) -> None:
self.action_lock = threading.Lock()
if not self.auto_commit:
head = self._select()
if head:
self.cursor = head[0] - 1
else:
self.cursor = 0
self.total = self._count()
def get_pooled_conn(self) -> Any:
return self._connection_pool.connection()
class MySQLConn:
"""MySqlConn defines a common structure for
both mysql and sqlite3 connections.
used to mitigate the interface differences between drivers/db
"""
def __init__(self,
queue: Optional[MySQLQueue] = None,
conn: Optional[Any] = None) -> None:
self._queue = queue
if queue is not None:
self._conn = queue.get_pooled_conn()
else:
self._conn = conn
self._cursor = None
self.closed = False
def __enter__(self) -> Any:
self._cursor = self._conn.cursor()
return self._conn
def __exit__(self,
exc_type: Optional[type],
exc_val: Optional[BaseException],
exc_tb: Optional[Any]) -> None:
# do not commit() but to close() , keep same behavior
# with dbutils
self._cursor.close()
def execute(self, *args: Any, **kwargs: Any) -> Any:
if self._queue is not None:
conn = self._queue.get_pooled_conn()
else:
conn = self._conn
cursor = conn.cursor()
cursor.execute(*args, **kwargs)
return cursor
def close(self) -> None:
if not self.closed:
self._conn.close()
self.closed = True
def commit(self) -> None:
if not self.closed:
self._conn.commit()
|
