5 from pgtoolkit import pgmanager, pgbrowser, config |
5 from pgtoolkit import pgmanager, pgbrowser, config |
6 from pgtoolkit.coloredformatter import ColoredFormatter |
6 from pgtoolkit.coloredformatter import ColoredFormatter |
7 from pgtoolkit.highlight import highlight |
7 from pgtoolkit.highlight import highlight |
8 |
8 |
9 |
9 |
10 class ConnectionInfoNotFound(Exception): |
10 class ConnectionInfoNotFound(Exception): |
11 pass |
11 pass |
12 |
12 |
13 |
13 |
14 class BadArgsError(Exception): |
14 class BadArgsError(Exception): |
15 pass |
15 pass |
16 |
16 |
17 |
17 |
18 class ToolBase: |
18 class ToolBase: |
19 def __init__(self, name, desc): |
19 def __init__(self, name, desc): |
20 self.parser = argparse.ArgumentParser(description=desc) |
20 self.parser = argparse.ArgumentParser(description=desc) |
21 self.parser.add_argument('-d', dest='debug', action='store_true', |
21 self.parser.add_argument('-d', dest='debug', action='store_true', |
22 help='Debug mode - print database queries.') |
22 help='Debug mode - print database queries.') |
23 |
23 |
24 self.config = config.ConfigParser() |
24 self.config = config.ConfigParser() |
25 self.config.add_argument('databases', type=dict) |
25 self.config.add_argument('databases', type=dict) |
26 self.config.add_argument('meta_db') |
26 self.config.add_argument('meta_db') |
27 self.config.add_argument('meta_query') |
27 self.config.add_argument('meta_query') |
28 |
28 |
29 self.pgm = pgmanager.get_instance() |
29 self.pgm = pgmanager.get_instance() |
30 self.target_isolation_level = None |
30 self.target_isolation_level = None |
31 |
31 |
32 def init(self): |
32 def init(self): |
33 self.config.load('pgtoolkit.conf') |
33 self.config.load('pgtoolkit.conf') |
34 self.args = self.parser.parse_args() |
34 self.args = self.parser.parse_args() |
35 self.init_logging() |
35 self.init_logging() |
36 |
36 |
37 def init_logging(self): |
37 def init_logging(self): |
38 # logging |
38 # logging |
39 handler = logging.StreamHandler() |
39 handler = logging.StreamHandler() |
40 format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S') |
40 format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S') |
41 handler.setFormatter(format) |
41 handler.setFormatter(format) |
42 handler.setLevel(logging.DEBUG) |
42 handler.setLevel(logging.DEBUG) |
43 self.log = logging.getLogger('main') |
43 self.log = logging.getLogger('main') |
44 self.log.addHandler(handler) |
44 self.log.addHandler(handler) |
45 self.log.setLevel(logging.DEBUG) |
45 self.log.setLevel(logging.DEBUG) |
46 |
46 |
|
47 logger_notices = logging.getLogger('pgmanager_notices') |
|
48 logger_notices.addHandler(handler) |
|
49 logger_notices.setLevel(logging.DEBUG) |
|
50 |
47 if self.args.debug: |
51 if self.args.debug: |
48 handler = logging.StreamHandler() |
52 logger_sql = logging.getLogger('pgmanager_sql') |
49 handler.setFormatter(format) |
53 logger_sql.addHandler(handler) |
50 handler.setLevel(logging.DEBUG) |
54 logger_sql.setLevel(logging.DEBUG) |
51 logger = logging.getLogger('pgmanager_sql') |
|
52 logger.addHandler(handler) |
|
53 logger.setLevel(logging.DEBUG) |
|
54 |
55 |
55 def prepare_conn_from_metadb(self, name, lookup_name): |
56 def prepare_conn_from_metadb(self, name, lookup_name): |
56 '''Create connection in pgmanager using meta DB. |
57 '''Create connection in pgmanager using meta DB. |
57 |
58 |
58 name -- Name for connection in pgmanager. |
59 name -- Name for connection in pgmanager. |
59 lookup_name -- Name of connection in meta DB. |
60 lookup_name -- Name of connection in meta DB. |
60 |
61 |
61 ''' |
62 ''' |
62 with self.pgm.cursor('meta') as curs: |
63 with self.pgm.cursor('meta') as curs: |
63 curs.execute(self.config.meta_query, [lookup_name]) |
64 curs.execute(self.config.meta_query, [lookup_name]) |
64 row = curs.fetchone_dict() |
65 row = curs.fetchone_dict() |
65 curs.connection.commit() |
66 curs.connection.commit() |
80 return True |
81 return True |
81 |
82 |
82 def prepare_conns_from_cmdline_args(self, *pgm_names): |
83 def prepare_conns_from_cmdline_args(self, *pgm_names): |
83 if self.config.meta_db: |
84 if self.config.meta_db: |
84 self.pgm.create_conn(name='meta', dsn=self.config.meta_db) |
85 self.pgm.create_conn(name='meta', dsn=self.config.meta_db) |
85 |
86 |
86 for name in pgm_names: |
87 for name in pgm_names: |
87 lookup_name = self.args.__dict__[name] |
88 lookup_name = self.args.__dict__[name] |
88 found = self.prepare_conn_from_config(name, lookup_name) |
89 found = self.prepare_conn_from_config(name, lookup_name) |
89 if not found and self.config.meta_db: |
90 if not found and self.config.meta_db: |
90 found = self.prepare_conn_from_metadb(name, lookup_name) |
91 found = self.prepare_conn_from_metadb(name, lookup_name) |
91 if not found: |
92 if not found: |
92 raise ConnectionInfoNotFound('Connection name "%s" not found in config nor in meta DB.' % lookup_name) |
93 raise ConnectionInfoNotFound('Connection name "%s" not found in config nor in meta DB.' % lookup_name) |
93 |
94 |
94 if self.config.meta_db: |
95 if self.config.meta_db: |
95 self.pgm.close_conn('meta') |
96 self.pgm.close_conn('meta') |
96 |
97 |
97 |
98 |
98 class SimpleTool(ToolBase): |
99 class SimpleTool(ToolBase): |
99 def __init__(self, name, desc): |
100 def __init__(self, name, desc): |
100 ToolBase.__init__(self, name, desc) |
101 ToolBase.__init__(self, name, desc) |
101 self.parser.add_argument('target', metavar='target', type=str, help='Target database') |
102 self.parser.add_argument('target', metavar='target', type=str, help='Target database') |
102 |
103 |
103 def init(self): |
104 def init(self): |
104 ToolBase.init(self) |
105 ToolBase.init(self) |
105 self.prepare_conns_from_cmdline_args('target') |
106 self.prepare_conns_from_cmdline_args('target') |
106 |
107 |
107 |
108 |
108 class SrcDstTool(ToolBase): |
109 class SrcDstTool(ToolBase): |
109 def __init__(self, name, desc): |
110 def __init__(self, name, desc): |
110 ToolBase.__init__(self, name, desc) |
111 ToolBase.__init__(self, name, desc) |
111 self.parser.add_argument('src', metavar='source', type=str, help='Source database') |
112 self.parser.add_argument('src', metavar='source', type=str, help='Source database') |
112 self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database') |
113 self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database') |
113 |
114 |
114 def init(self): |
115 def init(self): |
115 ToolBase.init(self) |
116 ToolBase.init(self) |
116 self.prepare_conns_from_cmdline_args('src', 'dst') |
117 self.prepare_conns_from_cmdline_args('src', 'dst') |
117 |
118 |
118 |
119 |
126 self.parser.add_argument('--dst-table', metavar='destination_table', |
127 self.parser.add_argument('--dst-table', metavar='destination_table', |
127 dest='dsttable', type=str, default='', help='Destination table name (default=source_table).') |
128 dest='dsttable', type=str, default='', help='Destination table name (default=source_table).') |
128 self.parser.add_argument('--dst-schema', metavar='destination_schema', |
129 self.parser.add_argument('--dst-schema', metavar='destination_schema', |
129 dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).') |
130 dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).') |
130 self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.") |
131 self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.") |
131 |
132 |
132 def init(self): |
133 def init(self): |
133 SrcDstTool.init(self) |
134 SrcDstTool.init(self) |
134 |
135 |
135 self.schema1 = self.args.srcschema |
136 self.schema1 = self.args.srcschema |
136 self.table1 = self.args.srctable |
137 self.table1 = self.args.srctable |
137 self.schema2 = self.args.dstschema |
138 self.schema2 = self.args.dstschema |
138 self.table2 = self.args.dsttable |
139 self.table2 = self.args.dsttable |
139 |
140 |
140 # check regex - it applies to source name, dest name must not be specified |
141 # check regex - it applies to source name, dest name must not be specified |
141 # applies to only one - schema or table name |
142 # applies to only one - schema or table name |
142 if self.args.regex: |
143 if self.args.regex: |
143 if self.table2 or (self.schema2 and not self.table1): |
144 if self.table2 or (self.schema2 and not self.table1): |
144 raise BadArgsError('Cannot specify both --regex and --dst-schema, --dst-table.') |
145 raise BadArgsError('Cannot specify both --regex and --dst-schema, --dst-table.') |
178 else: |
179 else: |
179 # one table |
180 # one table |
180 yield (self.schema1, self.table1, self.schema2, self.table2) |
181 yield (self.schema1, self.table1, self.schema2, self.table2) |
181 finally: |
182 finally: |
182 self.pgm.put_conn(srcconn, 'src') |
183 self.pgm.put_conn(srcconn, 'src') |
183 |
184 |
184 def _iter_schemas_regex(self, browser, regex): |
185 def _iter_schemas_regex(self, browser, regex): |
185 for schema in browser.list_schemas(): |
186 for schema in browser.list_schemas(): |
186 if schema['system']: |
187 if schema['system']: |
187 continue |
188 continue |
188 schemaname = schema['name'] |
189 schemaname = schema['name'] |
189 if re.match(regex, schemaname): |
190 if re.match(regex, schemaname): |
190 for item in self._iter_tables_regex(browser, schemaname, schemaname, ''): |
191 for item in self._iter_tables_regex(browser, schemaname, schemaname, ''): |
191 yield item |
192 yield item |
192 |
193 |
193 def _iter_tables_regex(self, browser, schema1, schema2, regex): |
194 def _iter_tables_regex(self, browser, schema1, schema2, regex): |
194 for table in browser.list_tables(schema1): |
195 for table in browser.list_tables(schema1): |
195 tablename = table['name'] |
196 tablename = table['name'] |
196 if re.match(regex, tablename): |
197 if re.match(regex, tablename): |
197 yield (schema1, tablename, schema2, tablename) |
198 yield (schema1, tablename, schema2, tablename) |