equal
deleted
inserted
replaced
13 import getpass |
13 import getpass |
14 import sys |
14 import sys |
15 |
15 |
16 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser |
16 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser |
17 |
17 |
18 has_psycopg2 = False |
18 import psycopg2 |
19 try: |
|
20 import psycopg2 |
|
21 has_psycopg2 = True |
|
22 except ImportError: |
|
23 try: |
|
24 from pyPgSQL import PgSQL |
|
25 except ImportError: |
|
26 sys.stderr.write('error: no suitable database module found\n') |
|
27 raise SystemExit(1) |
|
28 |
|
29 if has_psycopg2: |
|
30 DBErr = psycopg2.DatabaseError |
|
31 else: |
|
32 DBErr = PgSQL.libpq.DatabaseError |
|
33 |
19 |
34 |
20 |
35 def check_args(args, err_hdlr): |
21 def check_args(args, err_hdlr): |
36 if not args.postfix: |
22 if not args.postfix: |
37 err_hdlr('missing Postfix database user name') |
23 err_hdlr('missing Postfix database user name') |
40 if args.askp: |
26 if args.askp: |
41 args.dbpass = getpass.getpass() |
27 args.dbpass = getpass.getpass() |
42 |
28 |
43 |
29 |
44 def get_dbh(database, user, password, host, port): |
30 def get_dbh(database, user, password, host, port): |
45 if has_psycopg2: |
31 return psycopg2.connect(database=database, user=user, |
46 return psycopg2.connect(database=database, user=user, |
32 password=password, host=host, port=port) |
47 password=password, host=host, port=port) |
|
48 return PgSQL.connect(user=user, password=password, host=host, |
|
49 database=database, port=port) |
|
50 |
33 |
51 |
34 |
52 def get_argparser(): |
35 def get_argparser(): |
53 descr = 'Set permissions for Dovecot and Postfix in the vmm database.' |
36 descr = 'Set permissions for Dovecot and Postfix in the vmm database.' |
54 parser = ArgumentParser(description=descr, usage='%(prog)s OPTIONS', |
37 parser = ArgumentParser(description=descr, usage='%(prog)s OPTIONS', |
129 versions['pgsql'] = dbh.server_version |
112 versions['pgsql'] = dbh.server_version |
130 else: |
113 else: |
131 try: |
114 try: |
132 dbc.execute("SELECT current_setting('server_version_num')") |
115 dbc.execute("SELECT current_setting('server_version_num')") |
133 versions['pgsql'] = int(dbc.fetchone()[0]) |
116 versions['pgsql'] = int(dbc.fetchone()[0]) |
134 except DBErr: |
117 except psycopg2.DatabaseError: |
135 versions['pgsql'] = 80199 |
118 versions['pgsql'] = 80199 |
136 dbc.execute("SELECT relname FROM pg_stat_user_tables WHERE relname LIKE " |
119 dbc.execute("SELECT relname FROM pg_stat_user_tables WHERE relname LIKE " |
137 "'userquota%'") |
120 "'userquota%'") |
138 res = dbc.fetchall() |
121 res = dbc.fetchall() |
139 dbc.close() |
122 dbc.close() |