108 |
113 |
109 def init(self): |
114 def init(self): |
110 ToolBase.init(self) |
115 ToolBase.init(self) |
111 self.prepare_conns_from_cmdline_args('src', 'dst') |
116 self.prepare_conns_from_cmdline_args('src', 'dst') |
112 |
117 |
|
118 |
|
119 class SrcDstTablesTool(SrcDstTool): |
|
120 def __init__(self, name, desc): |
|
121 SrcDstTool.__init__(self, name, desc) |
|
122 self.parser.add_argument('-t', '--src-table', metavar='source_table', |
|
123 dest='srctable', type=str, default='', help='Source table name.') |
|
124 self.parser.add_argument('-s', '--src-schema', metavar='source_schema', |
|
125 dest='srcschema', type=str, default='', help='Source schema name (default=public).') |
|
126 self.parser.add_argument('--dst-table', metavar='destination_table', |
|
127 dest='dsttable', type=str, default='', help='Destination table name (default=source_table).') |
|
128 self.parser.add_argument('--dst-schema', metavar='destination_schema', |
|
129 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 |
|
132 def init(self): |
|
133 SrcDstTool.init(self) |
|
134 |
|
135 self.schema1 = self.args.srcschema |
|
136 self.table1 = self.args.srctable |
|
137 self.schema2 = self.args.dstschema |
|
138 self.table2 = self.args.dsttable |
|
139 |
|
140 # check regex - it applies to source name, dest name must not be specified |
|
141 # applies to only one - schema or table name |
|
142 if self.args.regex: |
|
143 if self.table2 or (self.schema2 and not self.table1): |
|
144 raise BadArgsError('Cannot specify both --regex and --dst-schema, --dst-table.') |
|
145 # schema defaults to public |
|
146 if self.table1 and not self.schema1: |
|
147 self.schema1 = 'public' |
|
148 # dest defaults to source |
|
149 if not self.schema2: |
|
150 self.schema2 = self.schema1 |
|
151 if not self.table2: |
|
152 self.table2 = self.table1 |
|
153 |
|
154 def tables(self): |
|
155 '''Generator. Yields schema1, table1, schema2, table2.''' |
|
156 srcconn = self.pgm.get_conn('src') |
|
157 try: |
|
158 srcbrowser = pgbrowser.PgBrowser(srcconn) |
|
159 if self.args.regex: |
|
160 if not self.table1: |
|
161 # all tables from schemas defined by regex |
|
162 for item in self._iter_schemas_regex(srcbrowser, self.schema1): |
|
163 yield item |
|
164 else: |
|
165 # all tables defined by regex |
|
166 for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1): |
|
167 yield item |
|
168 else: |
|
169 if not self.table1: |
|
170 if not self.schema1: |
|
171 # all tables from all schemas |
|
172 for item in self._iter_schemas_regex(srcbrowser, self.schema1): |
|
173 yield item |
|
174 else: |
|
175 # all tables from specified schema |
|
176 for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1): |
|
177 yield item |
|
178 else: |
|
179 # one table |
|
180 yield (self.schema1, self.table1, self.schema2, self.table2) |
|
181 finally: |
|
182 self.pgm.put_conn(srcconn, 'src') |
|
183 |
|
184 def _iter_schemas_regex(self, browser, regex): |
|
185 for schema in browser.list_schemas(): |
|
186 if schema['system']: |
|
187 continue |
|
188 schemaname = schema['name'] |
|
189 if re.match(regex, schemaname): |
|
190 for item in self._iter_tables_regex(browser, schemaname, schemaname, ''): |
|
191 yield item |
|
192 |
|
193 def _iter_tables_regex(self, browser, schema1, schema2, regex): |
|
194 for table in browser.list_tables(schema1): |
|
195 tablename = table['name'] |
|
196 if re.match(regex, tablename): |
|
197 yield (schema1, tablename, schema2, tablename) |
|
198 |