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
|
#!/usr/bin/python2.4
"""
dropdb only more so.
Cut off access, slaughter connections and burn the database to the ground.
"""
import sys
import time
import psycopg2
import psycopg2.extensions
from signal import SIGTERM, SIGQUIT, SIGKILL, SIGINT
from optparse import OptionParser
def connect(dbname='template1'):
"""Connect to the database, returning the DB-API connection."""
if options.user is not None:
return psycopg2.connect("dbname=%s user=%s" % (dbname, options.user))
else:
return psycopg2.connect("dbname=%s" % dbname)
def send_signal(database, signal):
con = connect()
con.set_isolation_level(1) # READ COMMITTED. We rollback changes we make.
cur = con.cursor()
# Install PL/PythonU if it isn't already.
cur.execute("SELECT TRUE FROM pg_language WHERE lanname = 'plpythonu'")
if cur.fetchone() is None:
cur.execute('CREATE LANGUAGE "plpythonu"')
# Create a stored procedure to kill a backend process.
qdatabase = str(psycopg2.extensions.QuotedString(database))
cur.execute("""
CREATE OR REPLACE FUNCTION _pgmassacre_killall(integer)
RETURNS Boolean AS $$
import os
signal = args[0]
for row in plpy.execute('''
SELECT procpid FROM pg_stat_activity WHERE datname=%(qdatabase)s
AND procpid != pg_backend_pid()
'''):
try:
os.kill(row['procpid'], signal)
except OSError:
pass
else:
return False
return True
$$ LANGUAGE plpythonu
""" % vars())
cur.execute("SELECT _pgmassacre_killall(%(signal)s)", vars())
con.rollback()
con.close()
def rollback_prepared_transactions(database):
"""Rollback any prepared transactions.
For some crazy reason PostgreSQL will refuse to drop a database with
outstanding prepared transactions.
"""
con = connect(database)
con.set_isolation_level(0) # Autocommit so we can ROLLBACK PREPARED.
cur = con.cursor()
# Get a list of outstanding prepared transactions.
cur.execute(
"SELECT gid FROM pg_prepared_xacts WHERE database=%(database)s",
vars())
xids = [row[0] for row in cur.fetchall()]
for xid in xids:
cur.execute("ROLLBACK PREPARED %(xid)s", vars())
con.close()
def still_open(database):
"""Return True if there are still open connections.
Waits a while to ensure that connections shutting down have a chance to.
"""
con = connect()
con.set_isolation_level(1)
cur = con.cursor()
# Wait for up to 10 seconds, returning True if all backends are gone.
start = time.time()
while time.time() < start + 10:
cur.execute("""
SELECT procpid FROM pg_stat_activity
WHERE datname=%(database)s LIMIT 1
""", vars())
if cur.fetchone() is None:
return False
time.sleep(0.6) # Stats only updated every 500ms.
con.rollback()
con.close()
return True
options = None
def main():
parser = OptionParser()
parser.add_option("-U", "--user", dest="user", default=None,
help="Connect as USER", metavar="USER")
global options
(options, args) = parser.parse_args()
if len(args) != 1:
print >> sys.stderr, (
'Must specify one, and only one, database to destroy')
sys.exit(1)
database = args[0]
if database in ('template1', 'template0'):
print >> sys.stderr, (
"Put the gun down and back away from the vehicle!")
return 666
con = connect()
con.set_isolation_level(0) # Autocommit
cur = con.cursor()
# Ensure the database exists. Note that the script returns success
# if the database does not exist to ease scripting.
cur.execute(
"SELECT count(*) FROM pg_database WHERE datname=%s", [database])
if cur.fetchone()[0] == 0:
print >> sys.stderr, (
"%s has fled the building. Database does not exist"
% database)
return 0
# Allow connections to the doomed database if something turned this off,
# such as an aborted run of this script.
cur.execute(
"UPDATE pg_database SET datallowconn=TRUE WHERE datname=%s",
[database])
# Rollback prepared transactions.
rollback_prepared_transactions(database)
try:
# Stop connections to the doomed database.
cur.execute(
"UPDATE pg_database SET datallowconn=FALSE WHERE datname=%s",
[database])
con.close()
# Terminate current statements.
send_signal(database, SIGINT)
# Shutdown current connections normally.
send_signal(database, SIGTERM)
# Shutdown current connections immediately.
if still_open(database):
send_signal(database, SIGQUIT)
# Shutdown current connections nastily.
if still_open(database):
send_signal(database, SIGKILL)
if still_open(database):
print >> sys.stderr, (
"Unable to kill all backends! Database not destroyed.")
return 9
# Destroy the database.
con = connect()
# AUTOCOMMIT required to execute commands like DROP DATABASE.
con.set_isolation_level(0)
cur = con.cursor()
cur.execute("DROP DATABASE %s" % database) # Not quoted.
return 0
finally:
# In case something messed up, allow connections again so we can
# inspect the damage.
con = connect()
con.set_isolation_level(0)
cur = con.cursor()
cur.execute(
"UPDATE pg_database SET datallowconn=TRUE WHERE datname=%s",
[database])
if __name__ == '__main__':
sys.exit(main())
|