-
Notifications
You must be signed in to change notification settings - Fork 105
/
sharedpostgresql.py
274 lines (235 loc) · 8.2 KB
/
sharedpostgresql.py
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
#
# Gramps - a GTK+/GNOME based genealogy program
#
# Copyright (C) 2015-2016 Douglas S. Blank <doug.blank@gmail.com>
# Copyright (C) 2016-2017 Nick Hall
# Copyright (C) 2022 David Straub
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
"""
Backend for PostgreSQL database.
"""
# -------------------------------------------------------------------------
#
# Standard python modules
#
# -------------------------------------------------------------------------
import psycopg2
import os
import re
from uuid import uuid4
# -------------------------------------------------------------------------
#
# Gramps modules
#
# -------------------------------------------------------------------------
from gramps.gen.utils.configmanager import ConfigManager
from gramps.gen.config import config
from gramps.gen.db.dbconst import ARRAYSIZE
from gramps.gen.db.exceptions import DbConnectionError
from gramps.gen.const import GRAMPS_LOCALE as glocale
try:
_trans = glocale.get_addon_translator(__file__)
except ValueError:
_trans = glocale.translation
_ = _trans.gettext
from shareddbapi import SharedDBAPI
psycopg2.paramstyle = "format"
# -------------------------------------------------------------------------
#
# SharedPostgreSQL class
#
# -------------------------------------------------------------------------
class SharedPostgreSQL(SharedDBAPI):
def get_summary(self):
"""
Return a diction of information about this database
backend.
"""
summary = super().get_summary()
summary.update(
{
_("Database version"): psycopg2.__version__,
_("Database module location"): psycopg2.__file__,
}
)
return summary
def requires_login(self):
return True
def _initialize(self, directory, username, password):
config_file = os.path.join(directory, "settings.ini")
config_mgr = ConfigManager(config_file)
config_mgr.register("database.dbname", "")
config_mgr.register("database.host", "")
config_mgr.register("database.port", "")
config_mgr.register("tree.uuid", "")
if not os.path.exists(config_file):
config_mgr.set("database.dbname", "gramps")
config_mgr.set("database.host", config.get("database.host"))
config_mgr.set("database.port", config.get("database.port"))
config_mgr.set("tree.uuid", uuid4().hex)
config_mgr.save()
config_mgr.load()
dbkwargs = {}
for key in config_mgr.get_section_settings("database"):
value = config_mgr.get("database." + key)
if value:
dbkwargs[key] = value
if username:
dbkwargs["user"] = username
if password:
dbkwargs["password"] = password
uuid = config_mgr.get("tree.uuid")
if not uuid:
raise ValueError("No tree UUID found.")
try:
self.dbapi = Connection(uuid=uuid, **dbkwargs)
except psycopg2.OperationalError as msg:
raise DbConnectionError(str(msg), config_file)
# -------------------------------------------------------------------------
#
# Connection class
#
# -------------------------------------------------------------------------
class Connection:
def __init__(self, *args, uuid, **kwargs):
self.__connection = psycopg2.connect(*args, **kwargs)
self.__connection.autocommit = True
self.__cursor = self.__connection.cursor()
self.uuid = uuid
self._treeid = ""
self.check_collation(glocale)
@property
def treeid(self):
"""Return an integer treeid from the UUID."""
# return cached value
if self._treeid:
return self._treeid
# try to fetch ID from database
self.execute("SELECT treeid FROM trees WHERE uuid = %s", [self.uuid])
row = self.fetchone()
if row:
self._treeid = row[0]
return self.treeid
# create new ID
self.execute("INSERT INTO trees (uuid) VALUES (%s)", [self.uuid])
return self.treeid
def check_collation(self, locale):
"""
Checks that a collation exists and if not creates it.
:param locale: Locale to be checked.
:param type: A GrampsLocale object.
"""
# Duplicating system collations works, but to delete them the schema
# must be specified, so get the current schema
collation = locale.get_collation()
self.execute(
'CREATE COLLATION IF NOT EXISTS "%s"'
"(LOCALE = '%s')" % (collation, locale.collation)
)
def execute(self, *args, **kwargs):
sql = _hack_query(args[0])
if len(args) > 1:
args = args[1]
else:
args = None
try:
self.__cursor.execute(sql, args, **kwargs)
except:
self.__cursor.execute("rollback")
raise
def fetchone(self):
try:
return self.__cursor.fetchone()
except:
return None
def fetchall(self):
return self.__cursor.fetchall()
def begin(self):
self.__cursor.execute("BEGIN;")
def commit(self):
self.__cursor.execute("COMMIT;")
def rollback(self):
self.__connection.rollback()
def table_exists(self, table):
self.__cursor.execute(
"SELECT COUNT(*) " "FROM information_schema.tables " "WHERE table_name=%s;",
[table],
)
return self.fetchone()[0] != 0
def close(self):
self.__connection.close()
def cursor(self):
return Cursor(self.__connection)
# -------------------------------------------------------------------------
#
# Cursor class
#
# -------------------------------------------------------------------------
class Cursor:
def __init__(self, connection):
self.__connection = connection
def __enter__(self):
self.__cursor = self.__connection.cursor()
self.__cursor.arraysize = ARRAYSIZE
return self
def __exit__(self, *args, **kwargs):
self.__cursor.close()
def execute(self, *args, **kwargs):
"""
Executes an SQL statement.
:param args: arguments to be passed to the sqlite3 execute statement
:type args: list
:param kwargs: arguments to be passed to the sqlite3 execute statement
:type kwargs: list
"""
sql = _hack_query(args[0])
if len(args) > 1:
args = args[1]
else:
args = None
self.__cursor.execute(sql, args, **kwargs)
def fetchmany(self):
"""
Fetches the next set of rows of a query result, returning a list. An
empty list is returned when no more rows are available.
"""
try:
return self.__cursor.fetchmany()
except:
return None
def _hack_query(query):
query = query.replace("?", "%s")
query = query.replace("REGEXP", "~")
query = query.replace("desc", "desc_")
query = query.replace("BLOB", "bytea")
query = query.replace("INTEGER PRIMARY KEY", "SERIAL PRIMARY KEY")
## LIMIT offset, count
## count can be -1, for all
## LIMIT -1
## LIMIT offset, -1
query = query.replace("LIMIT -1", "LIMIT all") ##
match = re.match(".* LIMIT (.*), (.*) ", query)
if match and match.groups():
offset, count = match.groups()
if count == "-1":
count = "all"
query = re.sub(
"(.*) LIMIT (.*), (.*) ",
"\\1 LIMIT %s OFFSET %s " % (count, offset),
query,
)
return query