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
|
#
# Copyright (c) 2011, Canonical Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
__all__ = [
'PGBouncerFixture',
]
import itertools
import os.path
import signal
import socket
import subprocess
import time
from fixtures import Fixture, TempDir
import psycopg2
from testtools.content import content_from_file
def countdown(duration=60, sleep=0.1):
"""Provide a countdown iterator that sleeps between iterations.
Yields the current iteration count, starting from 1. The duration can be
in fractional seconds.
"""
start = time.time()
stop = start + duration
for iteration in itertools.count(1):
now = time.time()
if now < stop:
yield iteration
time.sleep(sleep)
else:
break
def _allocate_ports(n=1):
"""Allocate `n` unused ports.
There is a small race condition here (between the time we allocate the
port, and the time it actually gets used), but for the purposes for which
this function gets used it isn't a problem in practice.
"""
sockets = [socket.socket() for _ in range(n)]
try:
for s in sockets:
s.bind(('localhost', 0))
return [s.getsockname()[1] for s in sockets]
finally:
for s in sockets:
s.close()
class PGBouncerFixture(Fixture):
"""Programmatically configure and run pgbouncer.
Minimal usage:
>>> bouncer = PGBouncerFixture()
>>> bouncer.databases['mydb'] = 'host=hostname dbname=foo'
>>> bouncer.users['user1'] = 'credentials'
>>> with bouncer:
... connection = psycopg2.connect(
... database="mydb", host=bouncer.host, port=bouncer.port,
... user="user1", password="credentials")
"""
def __init__(self):
super().__init__()
# defaults
# pgbouncer -> path to pgbouncer executable
self.pgbouncer = 'pgbouncer'
# dbname -> connect string
self.databases = {}
# username -> details
self.users = {}
# list of usernames that can run all console queries
self.admin_users = []
# list of usernames that can run readonly console queries
self.stats_users = []
self.pool_mode = 'session'
self.unix_socket_dir = None
self.process = None
# Username and password that the fixture itself uses to run console
# queries. Since this fixture is only for testing, there's no harm
# in hardcoding these.
self.fixture_admin_username = '_pgbouncerfixture'
self.fixture_admin_password = 'trusted'
def setUp(self):
super().setUp()
self.addCleanup(self.stop)
self.host = '127.0.0.1'
self.port = _allocate_ports()[0]
self.configdir = self.useFixture(TempDir())
self.auth_type = 'trust'
self.setUpConf()
self.start()
def setUpConf(self):
"""Create a pgbouncer.ini file."""
self.inipath = os.path.join(self.configdir.path, 'pgbouncer.ini')
self.authpath = os.path.join(self.configdir.path, 'users.txt')
self.logpath = os.path.join(self.configdir.path, 'pgbouncer.log')
self.pidpath = os.path.join(self.configdir.path, 'pgbouncer.pid')
self.outputpath = os.path.join(self.configdir.path, 'output')
with open(self.inipath, 'w') as inifile:
inifile.write('[databases]\n')
for item in self.databases.items():
inifile.write('%s = %s\n' % item)
inifile.write('[pgbouncer]\n')
inifile.write('pool_mode = {}\n'.format(self.pool_mode))
inifile.write('listen_port = {}\n'.format(self.port))
inifile.write('listen_addr = {}\n'.format(self.host))
if self.unix_socket_dir is not None:
inifile.write(
'unix_socket_dir = {}\n'.format(self.unix_socket_dir))
inifile.write('auth_type = {}\n'.format(self.auth_type))
inifile.write('auth_file = {}\n'.format(self.authpath))
inifile.write('logfile = {}\n'.format(self.logpath))
inifile.write('pidfile = {}\n'.format(self.pidpath))
adminusers = ','.join(
[self.fixture_admin_username] + self.admin_users)
inifile.write('admin_users = {}\n'.format(adminusers))
statsusers = ','.join(self.stats_users)
inifile.write('stats_users = {}\n'.format(statsusers))
with open(self.authpath, 'w') as authfile:
authfile.write(
'"%s" "%s"\n' %
(self.fixture_admin_username, self.fixture_admin_password))
for user_creds in self.users.items():
authfile.write('"%s" "%s"\n' % user_creds)
@property
def is_running(self):
return (
# pgbouncer has been started.
self.process is not None and
# pgbouncer has not yet exited.
self.process.poll() is None)
def get_pgbouncer_version(self):
assert self.is_running
with psycopg2.connect(
host=self.host,
port=self.port,
database="pgbouncer",
user=self.fixture_admin_username,
password=self.fixture_admin_password,
) as con:
return con.server_version
def stop(self):
if not self.is_running:
return
if self.get_pgbouncer_version() >= 12300:
self.process.send_signal(signal.SIGQUIT)
else:
self.process.terminate()
for iteration in countdown():
if self.process.poll() is not None:
break
else:
raise Exception(
'Time-out waiting for pgbouncer to exit.')
def start(self):
if self.is_running:
return
# Add /usr/sbin if necessary to the PATH for magic just-works
# behavior with Ubuntu.
env = os.environ.copy()
if not self.pgbouncer.startswith('/'):
path = env['PATH'].split(os.pathsep)
if '/usr/sbin' not in path:
path.append('/usr/sbin')
env['PATH'] = os.pathsep.join(path)
with open(self.outputpath, "wb") as outputfile:
with open(os.devnull, "rb") as devnull:
self.process = subprocess.Popen(
[self.pgbouncer, self.inipath], env=env, stdin=devnull,
stdout=outputfile, stderr=outputfile)
self.addDetail(
os.path.basename(self.outputpath),
content_from_file(self.outputpath))
# Wait for the PID file to appear.
for iteration in countdown():
if os.path.isfile(self.pidpath):
with open(self.pidpath, "rb") as pidfile:
if pidfile.read().strip().isdigit():
break
else:
raise Exception(
'Time-out waiting for pgbouncer to create PID file.')
|
