|
2 | 2 |
|
3 | 3 | import io |
4 | 4 | import os |
5 | | -import six |
| 5 | +import psutil |
6 | 6 | import subprocess |
7 | 7 | import time |
8 | 8 |
|
9 | 9 | from shutil import rmtree |
10 | | -from six import raise_from |
| 10 | +from six import raise_from, iteritems |
11 | 11 | from tempfile import mkstemp, mkdtemp |
12 | 12 |
|
13 | | -from .enums import NodeStatus |
| 13 | +from .enums import NodeStatus, ProcessType |
14 | 14 |
|
15 | 15 | from .cache import cached_initdb |
16 | 16 |
|
|
49 | 49 | QueryException, \ |
50 | 50 | StartNodeException, \ |
51 | 51 | TimeoutException, \ |
52 | | - InitNodeException |
| 52 | + InitNodeException, \ |
| 53 | + TestgresException |
53 | 54 |
|
54 | 55 | from .logger import TestgresLogger |
55 | 56 |
|
|
67 | 68 | from .backup import NodeBackup |
68 | 69 |
|
69 | 70 |
|
| 71 | +class ProcessProxy(object): |
| 72 | + """ |
| 73 | + Wrapper for psutil.Process |
| 74 | +
|
| 75 | + Attributes: |
| 76 | + process: wrapped psutill.Process object |
| 77 | + ptype: instance of ProcessType |
| 78 | + """ |
| 79 | + |
| 80 | + def __init__(self, process): |
| 81 | + self.process = process |
| 82 | + self.ptype = ProcessType.from_process(process) |
| 83 | + |
| 84 | + def __getattr__(self, name): |
| 85 | + return getattr(self.process, name) |
| 86 | + |
| 87 | + def __repr__(self): |
| 88 | + return '{} : {}'.format(str(self.ptype), repr(self.process)) |
| 89 | + |
| 90 | + |
70 | 91 | class PostgresNode(object): |
71 | 92 | def __init__(self, name=None, port=None, base_dir=None): |
72 | 93 | """ |
@@ -119,7 +140,84 @@ def __exit__(self, type, value, traceback): |
119 | 140 |
|
120 | 141 | @property |
121 | 142 | def pid(self): |
122 | | - return self.get_pid() |
| 143 | + """ |
| 144 | + Return postmaster's PID if node is running, else 0. |
| 145 | + """ |
| 146 | + |
| 147 | + if self.status(): |
| 148 | + pid_file = os.path.join(self.data_dir, PG_PID_FILE) |
| 149 | + with io.open(pid_file) as f: |
| 150 | + return int(f.readline()) |
| 151 | + |
| 152 | + # for clarity |
| 153 | + return 0 |
| 154 | + |
| 155 | + @property |
| 156 | + def auxiliary_pids(self): |
| 157 | + """ |
| 158 | + Returns a dict of { ProcessType : PID }. |
| 159 | + """ |
| 160 | + |
| 161 | + result = {} |
| 162 | + |
| 163 | + for process in self.auxiliary_processes: |
| 164 | + if process.ptype not in result: |
| 165 | + result[process.ptype] = [] |
| 166 | + |
| 167 | + result[process.ptype].append(process.pid) |
| 168 | + |
| 169 | + return result |
| 170 | + |
| 171 | + @property |
| 172 | + def auxiliary_processes(self): |
| 173 | + """ |
| 174 | + Returns a list of auxiliary processes. |
| 175 | + Each process is represented by ProcessProxy object. |
| 176 | + """ |
| 177 | + |
| 178 | + def is_aux(process): |
| 179 | + return process.ptype != ProcessType.Unknown |
| 180 | + |
| 181 | + return list(filter(is_aux, self.child_processes)) |
| 182 | + |
| 183 | + @property |
| 184 | + def child_processes(self): |
| 185 | + """ |
| 186 | + Returns a list of all child processes. |
| 187 | + Each process is represented by ProcessProxy object. |
| 188 | + """ |
| 189 | + |
| 190 | + # get a list of postmaster's children |
| 191 | + children = psutil.Process(self.pid).children() |
| 192 | + |
| 193 | + return [ProcessProxy(p) for p in children] |
| 194 | + |
| 195 | + @property |
| 196 | + def source_walsender(self): |
| 197 | + """ |
| 198 | + Returns master's walsender feeding this replica. |
| 199 | + """ |
| 200 | + |
| 201 | + sql = """ |
| 202 | + select pid |
| 203 | + from pg_catalog.pg_stat_replication |
| 204 | + where application_name = $1 |
| 205 | + """ |
| 206 | + |
| 207 | + if not self.master: |
| 208 | + raise TestgresException("Node doesn't have a master") |
| 209 | + |
| 210 | + # master should be on the same host |
| 211 | + assert self.master.host == self.host |
| 212 | + |
| 213 | + with self.master.connect() as con: |
| 214 | + for row in con.execute(sql, self.name): |
| 215 | + for child in self.master.auxiliary_processes: |
| 216 | + if child.pid == int(row[0]): |
| 217 | + return child |
| 218 | + |
| 219 | + msg = "Master doesn't send WAL to {}".format(self.name) |
| 220 | + raise TestgresException(msg) |
123 | 221 |
|
124 | 222 | @property |
125 | 223 | def master(self): |
@@ -428,19 +526,6 @@ def status(self): |
428 | 526 | elif e.exit_code == 4: |
429 | 527 | return NodeStatus.Uninitialized |
430 | 528 |
|
431 | | - def get_pid(self): |
432 | | - """ |
433 | | - Return postmaster's PID if node is running, else 0. |
434 | | - """ |
435 | | - |
436 | | - if self.status(): |
437 | | - pid_file = os.path.join(self.data_dir, PG_PID_FILE) |
438 | | - with io.open(pid_file) as f: |
439 | | - return int(f.readline()) |
440 | | - |
441 | | - # for clarity |
442 | | - return 0 |
443 | | - |
444 | 529 | def get_control_data(self): |
445 | 530 | """ |
446 | 531 | Return contents of pg_control file. |
@@ -895,7 +980,7 @@ def catchup(self, dbname=None, username=None): |
895 | 980 | """ |
896 | 981 |
|
897 | 982 | if not self.master: |
898 | | - raise CatchUpException("Node doesn't have a master") |
| 983 | + raise TestgresException("Node doesn't have a master") |
899 | 984 |
|
900 | 985 | if pg_version_ge('10'): |
901 | 986 | poll_lsn = "select pg_current_wal_lsn()::text" |
@@ -1042,7 +1127,7 @@ def pgbench_run(self, |
1042 | 1127 | "-U", username, |
1043 | 1128 | ] + options |
1044 | 1129 |
|
1045 | | - for key, value in six.iteritems(kwargs): |
| 1130 | + for key, value in iteritems(kwargs): |
1046 | 1131 | # rename keys for pgbench |
1047 | 1132 | key = key.replace('_', '-') |
1048 | 1133 |
|
|
0 commit comments