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
|
#!/usr/bin/env python
"""
dropdb only more so.
Cut off access, slaughter connections and burn the database to the ground.
"""
import os
import sys
import time
import psycopg
from signal import SIGTERM, SIGQUIT, SIGKILL, SIGINT
from optparse import OptionParser
def connect():
if options.user is not None:
return psycopg.connect("dbname=template1 user=%s" % options.user)
else:
return psycopg.connect("dbname=template1")
def send_signal(database, signal):
con = connect()
con.set_isolation_level(1)
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(psycopg.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"
):
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 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()
cur = con.cursor()
# Ensure the database exists. Note that the script returns success
# in this case 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
# Stop connetions to the doomed database
cur.execute(
"UPDATE pg_database SET datallowconn=false WHERE datname=%s", [database]
)
con.commit()
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()
con.set_isolation_level(0) # Required to execute commands like DROP DATABASE
cur = con.cursor()
cur.execute("DROP DATABASE %s" % database) # Not quoted
return 0
# print "Mwahahahaha!"
if __name__ == '__main__':
sys.exit(main())
|