├── .gitignore ├── LICENSE.txt ├── README.md ├── dump ├── __init__.py ├── __main__.py ├── interface │ ├── __init__.py │ └── postgres.py └── postgres.py ├── dump_test.py ├── setup.py └── tox.ini /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | *.tmproj 3 | *.pyc 4 | .vagrant 5 | 6 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2014+ Jay Marcyes 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Dump 2 | 3 | Python wrapper around psql and pg_dump to make it easier to backup/restore a PostgreSQL database. 4 | 5 | 6 | ## Backup 7 | 8 | You backup per table: 9 | 10 | $ dump backup --dbname=... --username=... --password=... --dir=/some/base/path table1 table2 ... 11 | 12 | 13 | ## Restore 14 | 15 | You can restore the entire backup directory: 16 | 17 | $ dump restore --dbname=... --username=... --password=... --dir=/some/base/path 18 | 19 | 20 | ## Install 21 | 22 | Use pip: 23 | 24 | $ pip install dump 25 | 26 | to install the latest and greatest: 27 | 28 | $ pip install --upgrade git+https://github.com/Jaymon/dump#egg=dump 29 | 30 | -------------------------------------------------------------------------------- /dump/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import unicode_literals, division, print_function, absolute_import 3 | import logging 4 | 5 | 6 | logging.getLogger(__name__).addHandler(logging.NullHandler()) 7 | 8 | 9 | __version__ = '0.0.5' 10 | 11 | -------------------------------------------------------------------------------- /dump/__main__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import unicode_literals, division, print_function, absolute_import 3 | import argparse 4 | import sys 5 | import logging 6 | 7 | from dump import __version__ 8 | from dump.interface import postgres 9 | 10 | 11 | def console_backup(args): 12 | kwargs = vars(args) 13 | tables = kwargs.pop("tables") 14 | 15 | db = postgres.Postgres(**kwargs) 16 | for table in tables: 17 | db.table_dump(table) 18 | 19 | return 0 20 | 21 | 22 | def console_restore(args): 23 | kwargs = vars(args) 24 | db = postgres.Postgres(**kwargs) 25 | db.restore() 26 | return 0 27 | 28 | 29 | def console(): 30 | ''' 31 | cli hook 32 | 33 | return -- integer -- the exit code 34 | ''' 35 | parser = argparse.ArgumentParser(description="backup/restore PostgreSQL databases", add_help=False) 36 | parser.add_argument("--version", action='version', version="%(prog)s {}".format(__version__)) 37 | 38 | parent_parser = argparse.ArgumentParser(add_help=False) 39 | parent_parser.add_argument("-d", "--dbname", dest="dbname", help="database name to connect to") 40 | parent_parser.add_argument("-U", "--username", "--user", dest="username", help="database user name") 41 | parent_parser.add_argument("-W", "--password", dest="password", help="database password") 42 | parent_parser.add_argument( 43 | "-h", "--host", "--hostname", 44 | dest="host", 45 | default="localhost", 46 | help="database server host or socket directory" 47 | ) 48 | parent_parser.add_argument("-p", "--port", type=int, default=5432, dest="port", help="database server post") 49 | parent_parser.add_argument( 50 | "--help", 51 | action="help", 52 | default=argparse.SUPPRESS, 53 | help="show this help message and exit" 54 | ) 55 | 56 | subparsers = parser.add_subparsers(dest="command", help="a sub command") 57 | subparsers.required = True # https://bugs.python.org/issue9253#msg186387 58 | backup_parser = subparsers.add_parser( 59 | "backup", 60 | parents=[parent_parser], 61 | help="backup a PostgreSQL database", 62 | add_help=False 63 | ) 64 | backup_parser.add_argument( 65 | "-D", "--dir", "--directory", 66 | dest="directory", 67 | help="directory where the backup files should go" 68 | ) 69 | backup_parser.add_argument( 70 | "--debug", 71 | dest="debug", 72 | action="store_true", 73 | help="Turn on debugging output" 74 | ) 75 | backup_parser.add_argument("tables", nargs="+") 76 | backup_parser.set_defaults(func=console_backup) 77 | 78 | restore_parser = subparsers.add_parser( 79 | "restore", 80 | parents=[parent_parser], 81 | help="restore a PostgreSQL database", 82 | add_help=False 83 | ) 84 | restore_parser.add_argument( 85 | "-D", "--dir", "--directory", 86 | dest="directory", 87 | help="directory where the backup files are located" 88 | ) 89 | restore_parser.add_argument( 90 | "--debug", 91 | dest="debug", 92 | action="store_true", 93 | help="Turn on debugging output" 94 | ) 95 | restore_parser.set_defaults(func=console_restore) 96 | 97 | args = parser.parse_args() 98 | 99 | if args.debug: 100 | logging.basicConfig(format="[%(levelname).1s] %(message)s", level=logging.DEBUG, stream=sys.stdout) 101 | else: 102 | logging.basicConfig(format="[%(levelname).1s] %(message)s", level=logging.INFO, stream=sys.stdout) 103 | 104 | ret_code = args.func(args) 105 | return ret_code 106 | 107 | 108 | if __name__ == "__main__": 109 | sys.exit(console()) 110 | 111 | -------------------------------------------------------------------------------- /dump/interface/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import unicode_literals, division, print_function, absolute_import 3 | 4 | -------------------------------------------------------------------------------- /dump/interface/postgres.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Make backing up a postgres db a little easier to manage 4 | 5 | 6 | tuning for restoring the db: 7 | http://stackoverflow.com/questions/2094963/postgresql-improving-pg-dump-pg-restore-performance 8 | 9 | other links: 10 | http://blog.kimiensoftware.com/2011/05/compare-pg_dump-and-gzip-compression-241 11 | """ 12 | from __future__ import unicode_literals, division, print_function, absolute_import 13 | import re 14 | import subprocess 15 | import os, time 16 | import tempfile 17 | import logging 18 | 19 | 20 | logger = logging.getLogger(__name__) 21 | 22 | 23 | class Postgres(object): 24 | """wrapper to dump postgres tables""" 25 | 26 | # http://dbaspot.com/postgresql/348627-pg_dump-t-give-where-condition.html 27 | # http://docs.python.org/2/library/tempfile.html 28 | 29 | def __init__(self, dbname, username, password, host=None, port=5432, directory=None, **kwargs): 30 | self.tmp_files = set() 31 | self.outfile_count = 0 32 | 33 | if directory: 34 | if not os.path.exists(directory): 35 | os.makedirs(directory) 36 | 37 | else: 38 | directory = tempfile.mkdtemp(prefix="postgres-{}".format(time.strftime("%y%m%d%S"))) 39 | 40 | self.directory = directory 41 | self.dbname = dbname 42 | self.username = username 43 | self.password = password 44 | self.host = host 45 | self.port = port 46 | 47 | # make sure we have the needed stuff 48 | self._run_cmd(["which", "psql"]) 49 | self._run_cmd(["which", "pg_dump"]) 50 | self._run_cmd(["which", "gzip"]) 51 | 52 | 53 | def __del__(self): 54 | # cleanup by getting rid of all the temporary files 55 | for tf in self.tmp_files: 56 | os.unlink(tf) 57 | 58 | def restore(self): 59 | """use the self.directory to restore a db 60 | 61 | NOTE -- this will only restore a database dumped with one of the methods 62 | of this class 63 | """ 64 | sql_files = [] 65 | for root, dirs, files in os.walk(self.directory): 66 | for f in files: 67 | if f.endswith(".sql.gz"): 68 | path = os.path.join(self.directory, f) 69 | self._run_cmd(["gunzip", path]) 70 | sql_files.append(f.rstrip(".gz")) 71 | 72 | elif f.endswith('.sql'): 73 | sql_files.append(f) 74 | 75 | sql_files.sort() # we want to go in the order the tables were dumped 76 | r = re.compile('\d{3,}_([^\.]+)') 77 | for f in sql_files: 78 | path = os.path.join(self.directory, f) 79 | m = r.match(f) 80 | if m: 81 | table = m.group(1) 82 | logger.info('------- restoring table {}'.format(table)) 83 | 84 | #psql_args = self._get_args('psql', '-X', '--echo-queries', '-f {}'.format(path)) 85 | psql_args = self._get_args('psql', '-X', '--quiet', '--file={}'.format(path)) 86 | self._run_cmd(psql_args) 87 | 88 | logger.info('------- restored table {}'.format(table)) 89 | 90 | return True 91 | 92 | def table_dump(self, table): 93 | """dump all the rows of the given table name""" 94 | if not table: raise ValueError("no table") 95 | 96 | cmds = [] 97 | logger.info('------- dumping table {}'.format(table)) 98 | cmd = self._get_args( 99 | "pg_dump", 100 | "--table={}".format(table), 101 | #"--data-only", 102 | "--clean", 103 | "--no-owner", 104 | "--column-inserts", 105 | ) 106 | cmds.append((cmd, {})) 107 | 108 | outfile_path = self._get_outfile_path(table) 109 | cmds.append(('gzip > "{}"'.format(outfile_path), {"shell": True})) 110 | 111 | #cmd += ' | {}'.format(' | '.join(pipes)) 112 | #cmd += ' > {}'.format(outfile_path) 113 | 114 | self._run_cmds(cmds) 115 | logger.info('------- dumped table {}'.format(table)) 116 | return True 117 | 118 | def _get_file(self): 119 | ''' 120 | return an opened tempfile pointer that can be used 121 | 122 | http://docs.python.org/2/library/tempfile.html 123 | ''' 124 | f = tempfile.NamedTemporaryFile(delete=False) 125 | self.tmp_files.add(f.name) 126 | return f 127 | 128 | def _get_args(self, executable, *args): 129 | """compile all the executable and the arguments, combining with common arguments 130 | to create a full batch of command args""" 131 | args = list(args) 132 | args.insert(0, executable) 133 | if self.username: 134 | args.append("--username={}".format(self.username)) 135 | if self.host: 136 | args.append("--host={}".format(self.host)) 137 | if self.port: 138 | args.append("--port={}".format(self.port)) 139 | 140 | args.append(self.dbname) 141 | #args.extend(other_args) 142 | return args 143 | 144 | def _get_env(self): 145 | """this returns an environment dictionary we want to use to run the command 146 | 147 | this will also create a fake pgpass file in order to make it possible for 148 | the script to be passwordless""" 149 | if hasattr(self, 'env'): return self.env 150 | 151 | # create a temporary pgpass file 152 | pgpass = self._get_file() 153 | # format: http://www.postgresql.org/docs/9.2/static/libpq-pgpass.html 154 | pgpass.write('*:*:*:{}:{}\n'.format(self.username, self.password).encode("utf-8")) 155 | pgpass.close() 156 | self.env = dict(os.environ) 157 | self.env['PGPASSFILE'] = pgpass.name 158 | 159 | # we want to assure a consistent environment 160 | if 'PGOPTIONS' in self.env: del self.env['PGOPTIONS'] 161 | return self.env 162 | 163 | def _run_cmds(self, cmds): 164 | 165 | pipes = [] 166 | env = self._get_env() 167 | for args, kwargs in cmds: 168 | logger.debug("Running: {}".format(args)) 169 | 170 | if "env" not in kwargs: 171 | kwargs["env"] = env 172 | 173 | if pipes: 174 | kwargs["stdin"] = pipes[-1].stdout 175 | 176 | kwargs["stdout"] = subprocess.PIPE 177 | 178 | pipes.append(subprocess.Popen(args, **kwargs)) 179 | 180 | for i, pipe in enumerate(pipes): 181 | ret_code = pipe.wait() 182 | if ret_code > 0: 183 | raise IOError("Command {} exited with {}".format(cmds[i][0], ret_code)) 184 | 185 | def _run_cmd(self, cmd, **kwargs): 186 | 187 | ignore_ret_code = kwargs.pop("ignore_ret_code", False) 188 | 189 | try: 190 | self._run_cmds([(cmd, kwargs)]) 191 | 192 | except IOError as e: 193 | if ignore_ret_code: 194 | logger.warn(e, exc_info=True) 195 | pass 196 | else: 197 | raise 198 | 199 | def _get_outfile_path(self, table): 200 | """return the path for a file we can use to back up the table""" 201 | self.outfile_count += 1 202 | outfile = os.path.join(self.directory, '{:03d}_{}.sql.gz'.format(self.outfile_count, table)) 203 | return outfile 204 | 205 | -------------------------------------------------------------------------------- /dump/postgres.py: -------------------------------------------------------------------------------- 1 | """ 2 | Backup only the most important parts of a Postgres db 3 | 4 | example -- 5 | # backup only guys named foo in the user table: 6 | import postgres 7 | pg = postgres.Dump(db, user, passwd, host, port) 8 | pg.select_dump('user', "SELECT * FROM user WHERE username='foo'") 9 | 10 | tuning for restoring the db: 11 | http://stackoverflow.com/questions/2094963/postgresql-improving-pg-dump-pg-restore-performance 12 | 13 | other links: 14 | http://blog.kimiensoftware.com/2011/05/compare-pg_dump-and-gzip-compression-241 15 | """ 16 | from __future__ import print_function 17 | import re 18 | import subprocess 19 | import os, time 20 | import tempfile 21 | 22 | 23 | class Postgres(object): 24 | """wrapper to dump postgres tables""" 25 | 26 | # http://dbaspot.com/postgresql/348627-pg_dump-t-give-where-condition.html 27 | # http://docs.python.org/2/library/tempfile.html 28 | 29 | def __init__(self, dbname, username, password, host=None, port=5432, directory=None, **kwargs): 30 | self.tmp_files = set() 31 | self.outfile_count = 0 32 | 33 | if directory: 34 | if not os.path.exists(directory): 35 | os.makedirs(directory) 36 | 37 | else: 38 | directory = tempfile.mkdtemp(prefix="postgres-{}".format(time.strftime("%y%m%d%S"))) 39 | 40 | self.directory = directory 41 | self.dbname = dbname 42 | self.username = username 43 | self.password = password 44 | self.host = host 45 | self.port = port 46 | 47 | # make sure we have the needed stuff 48 | self._run_cmd("which psql") 49 | self._run_cmd("which pg_dump") 50 | self._run_cmd("which gzip") 51 | 52 | 53 | def __del__(self): 54 | # cleanup by getting rid of all the temporary files 55 | for tf in self.tmp_files: 56 | os.unlink(tf) 57 | 58 | def restore(self): 59 | """use the self.directory to restore a db 60 | 61 | NOTE -- this will only restore a database dumped with one of the methods 62 | of this class 63 | """ 64 | 65 | # let's unzip all the files in the dir 66 | path = os.path.join(self.directory, '*') 67 | self._run_cmd('gunzip {}'.format(path), ignore_ret_code=True) 68 | 69 | sql_files = [] 70 | for root, dirs, files in os.walk(self.directory): 71 | for f in files: 72 | if f.endswith('.sql'): 73 | sql_files.append(f) 74 | 75 | sql_files.sort() # we want to go in the order the tables were dumped 76 | r = re.compile('\d{3,}_([^\.]+)') 77 | for f in sql_files: 78 | path = os.path.join(self.directory, f) 79 | m = r.match(f) 80 | if m: 81 | table = m.group(1) 82 | print('------- restoring table {}'.format(table)) 83 | 84 | #psql_args = self._get_args('psql', '-X', '--echo-queries', '-f {}'.format(path)) 85 | psql_args = self._get_args('psql', '-X', '--quiet', '-f {}'.format(path)) 86 | self._run_cmd(' '.join(psql_args)) 87 | 88 | # restore the sequence 89 | #self._restore_auto_increment(table) 90 | print('------- restored table {}'.format(table)) 91 | 92 | return True 93 | 94 | def table_dump(self, table): 95 | """dump all the rows of the given table name""" 96 | if not table: raise ValueError("no table") 97 | 98 | print('------- dumping table {}'.format(table)) 99 | pipes = ["gzip"] 100 | outfile_path = self._get_outfile_path(table) 101 | cmd = self._get_args( 102 | "pg_dump", 103 | "--table={}".format(table), 104 | #"--data-only", 105 | "--clean", 106 | "--no-owner", 107 | "--column-inserts", 108 | ) 109 | cmd = ' '.join(cmd) 110 | cmd += ' | {}'.format(' | '.join(pipes)) 111 | cmd += ' > {}'.format(outfile_path) 112 | 113 | self._run_cmd(cmd) 114 | print('------- dumped table {}'.format(table)) 115 | return True 116 | 117 | def _get_file(self): 118 | ''' 119 | return an opened tempfile pointer that can be used 120 | 121 | http://docs.python.org/2/library/tempfile.html 122 | ''' 123 | f = tempfile.NamedTemporaryFile(delete=False) 124 | self.tmp_files.add(f.name) 125 | return f 126 | 127 | def _get_args(self, executable, *args): 128 | """compile all the executable and the arguments, combining with common arguments 129 | to create a full batch of command args""" 130 | args = list(args) 131 | args.insert(0, executable) 132 | if self.username: 133 | args.append("--username={}".format(self.username)) 134 | if self.host: 135 | args.append("--host={}".format(self.host)) 136 | if self.port: 137 | args.append("--port={}".format(self.port)) 138 | 139 | args.append(self.dbname) 140 | #args.extend(other_args) 141 | return args 142 | 143 | def _get_env(self): 144 | """this returns an environment dictionary we want to use to run the command 145 | 146 | this will also create a fake pgpass file in order to make it possible for 147 | the script to be passwordless""" 148 | if hasattr(self, 'env'): return self.env 149 | 150 | # create a temporary pgpass file 151 | pgpass = self._get_file() 152 | # format: http://www.postgresql.org/docs/9.2/static/libpq-pgpass.html 153 | pgpass.write('*:*:*:{}:{}\n'.format(self.username, self.password)) 154 | pgpass.close() 155 | self.env = dict(os.environ) 156 | self.env['PGPASSFILE'] = pgpass.name 157 | 158 | # we want to assure a consistent environment 159 | if 'PGOPTIONS' in self.env: del self.env['PGOPTIONS'] 160 | return self.env 161 | 162 | def _run_cmd(self, cmd, ignore_ret_code=False, popen_kwargs=None): 163 | print(cmd) 164 | 165 | env = self._get_env() 166 | kwargs = { 167 | 'shell': True, 168 | 'env': env 169 | } 170 | if popen_kwargs: kwargs.update(popen_kwargs) 171 | pipe = subprocess.Popen( 172 | cmd, 173 | **kwargs 174 | ) 175 | ret_code = pipe.wait() 176 | if not ignore_ret_code and ret_code > 0: 177 | raise RuntimeError('command {} did not execute correctly'.format(cmd)) 178 | 179 | return pipe 180 | 181 | def _get_outfile_path(self, table): 182 | """return the path for a file we can use to back up the table""" 183 | self.outfile_count += 1 184 | outfile = os.path.join(self.directory, '{:03d}_{}.sql.gz'.format(self.outfile_count, table)) 185 | return outfile 186 | 187 | def _run_queries(self, queries, *args, **kwargs): 188 | """run the queries 189 | 190 | queries -- list -- the queries to run 191 | return -- string -- the results of the query? 192 | """ 193 | # write out all the commands to a temp file and then have psql run that file 194 | f = self._get_file() 195 | for q in queries: 196 | f.write("{};\n".format(q)) 197 | f.close() 198 | 199 | psql_args = self._get_args('psql', '-X', '-f {}'.format(f.name)) 200 | return self._run_cmd(' '.join(psql_args), *args, **kwargs) 201 | 202 | def _restore_auto_increment(self, table): 203 | """restore the auto increment value for the table to what it was previously""" 204 | query, seq_table, seq_column, seq_name = self._get_auto_increment_info(table) 205 | if query: 206 | queries = [query, "select nextval('{}')".format(seq_name)] 207 | return self._run_queries(queries) 208 | 209 | def _get_auto_increment_info(self, table): 210 | """figure out the the autoincrement value for the given table""" 211 | query = '' 212 | seq_table = '' 213 | seq_column = '' 214 | seq_name = '' 215 | find_query = "\n".join([ 216 | "SELECT", 217 | " t.relname as related_table,", 218 | " a.attname as related_column,", 219 | " s.relname as sequence_name", 220 | "FROM pg_class s", 221 | "JOIN pg_depend d ON d.objid = s.oid", 222 | "JOIN pg_class t ON d.objid = s.oid AND d.refobjid = t.oid", 223 | "JOIN pg_attribute a ON (d.refobjid, d.refobjsubid) = (a.attrelid, a.attnum)", 224 | "JOIN pg_namespace n ON n.oid = s.relnamespace", 225 | "WHERE", 226 | " s.relkind = 'S'", 227 | "AND", 228 | " n.nspname = 'public'", 229 | "AND", 230 | " t.relname = '{}'".format(table) 231 | ]) 232 | 233 | pipe = self._run_queries([find_query], popen_kwargs={'stdout': subprocess.PIPE}) 234 | stdout, stderr = pipe.communicate() 235 | if stdout: 236 | try: 237 | m = re.findall('^\s*(\S+)\s*\|\s*(\S+)\s*\|\s*(\S+)\s*$', stdout, flags=re.MULTILINE) 238 | seq_table, seq_column, seq_name = m[1] 239 | # http://www.postgresql.org/docs/9.2/static/functions-sequence.html 240 | # http://www.postgresql.org/docs/9.2/static/functions-conditional.html 241 | query = "\n".join([ 242 | "SELECT", 243 | " setval('{}',".format(seq_name.strip()), 244 | " coalesce(max({}), 1),".format(seq_column.strip()), 245 | " max({}) IS NOT null)".format(seq_column.strip()), 246 | "FROM \"{}\"".format(seq_table.strip()) 247 | ]) 248 | 249 | except IndexError: 250 | query = '' 251 | 252 | return query, seq_table, seq_column, seq_name 253 | 254 | 255 | -------------------------------------------------------------------------------- /dump_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | test dump 4 | 5 | to run on the command line: 6 | python -m unittest test_pout[.ClassTest[.test_method]] 7 | """ 8 | from __future__ import unicode_literals, division, print_function, absolute_import 9 | import unittest 10 | import subprocess 11 | import os 12 | import random 13 | 14 | import testdata 15 | import dsnparse 16 | import psycopg2 17 | import psycopg2.extras 18 | 19 | 20 | class Connection(object): 21 | instance = None 22 | 23 | @classmethod 24 | def get_instance(cls): 25 | if not cls.instance: 26 | cls.instance = cls(os.environ["DUMP_DSN"]) 27 | return cls.instance 28 | 29 | def __init__(self, dsn): 30 | self.dsn = dsnparse.parse(dsn) 31 | 32 | self.dbname = self.dsn.dbname 33 | self.user = self.dsn.username 34 | self.password = self.dsn.password 35 | self.host = self.dsn.hostname 36 | self.port = self.dsn.port 37 | 38 | if self.dsn.scheme.startswith("postgres"): 39 | self.conn = psycopg2.connect( 40 | dbname=self.dbname, 41 | user=self.user, 42 | password=self.password, 43 | host=self.host, 44 | port=self.port, 45 | cursor_factory=psycopg2.extras.RealDictCursor, 46 | ) 47 | self.conn.autocommit = True 48 | else: 49 | raise ValueError("{} is unsupported".format(self.dsn.scheme)) 50 | 51 | def close(self): 52 | self.conn.close() 53 | type(self).instance = None 54 | 55 | 56 | class Foo(object): 57 | table_name = "foo" 58 | 59 | fields = [ 60 | ("_id", "BIGSERIAL PRIMARY KEY"), 61 | ("bar", "INTEGER"), 62 | ] 63 | 64 | def __init__(self, **fields): 65 | for k, v in fields.items(): 66 | setattr(self, k, v) 67 | 68 | def delete(self): 69 | query_str = 'DROP TABLE IF EXISTS "{}" CASCADE'.format(self.table_name) 70 | ret = self.query(query_str, ignore_result=True) 71 | 72 | def install(self): 73 | self.delete() 74 | query_str = [] 75 | query_str.append('CREATE TABLE "{}" ('.format(self.table_name)) 76 | 77 | query_fields = [] 78 | for field_name, field in self.fields: 79 | query_fields.append(' {} {}'.format(field_name, field)) 80 | 81 | query_str.append(",\n".join(query_fields)) 82 | query_str.append(')') 83 | query_str = "\n".join(query_str) 84 | self.query(query_str, ignore_result=True) 85 | 86 | def count(self): 87 | ret = self.query('SELECT COUNT(*) FROM "{}"'.format(self.table_name)) 88 | return int(ret[0]['count']) 89 | 90 | def save(self): 91 | field_names = [] 92 | field_formats = [] 93 | query_vals = [] 94 | for k, _ in self.fields: 95 | try: 96 | v = getattr(self, k) 97 | query_vals.append(v) 98 | field_names.append(k) 99 | except AttributeError: pass 100 | 101 | #query_str = 'INSERT INTO {} ({}) VALUES ({}) RETURNING {}'.format( 102 | query_str = 'INSERT INTO "{}" ("{}") VALUES ({}) RETURNING "_id"'.format( 103 | self.table_name, 104 | '", "'.join(field_names), 105 | ', '.join(['%s'] * len(field_names)), 106 | ) 107 | 108 | return self.query(query_str, query_vals)[0]["_id"] 109 | 110 | def query(self, query_str, query_vals=[], **kwargs): 111 | ret = None 112 | cur = Connection.get_instance().conn.cursor() 113 | if query_vals: 114 | cur.execute(query_str, query_vals) 115 | else: 116 | cur.execute(query_str) 117 | 118 | q = query_str.lower() 119 | if not kwargs.get("ignore_result", False): 120 | ret = cur.fetchall() 121 | return ret 122 | 123 | def close(self): 124 | Connection.get_instance().close() 125 | 126 | 127 | class Bar(Foo): 128 | table_name = "bar" 129 | 130 | fields = [ 131 | ("_id", "BIGSERIAL PRIMARY KEY"), 132 | ("foo", "INTEGER"), 133 | ] 134 | 135 | 136 | class Client(object): 137 | """makes running a command nice and easy for easy peasy testing""" 138 | @property 139 | def files(self): 140 | for path, dirs, files in os.walk(self.directory): 141 | return [os.path.join(path, f) for f in files] 142 | #return files 143 | 144 | def __init__(self): 145 | self.code = 0 146 | self.output = "" 147 | self.directory = testdata.create_dir() 148 | 149 | conn = Connection.get_instance() 150 | self.arg_str = " ".join([ 151 | "--dbname={}".format(conn.dbname), 152 | "--username={}".format(conn.user), 153 | "--password={}".format(conn.password), 154 | "--host={}".format(conn.host), 155 | "--port={}".format(conn.port), 156 | "--dir={}".format(self.directory), 157 | "--debug", 158 | ]) 159 | 160 | def run(self, arg_str): 161 | cmd = "python -m dump {}".format(arg_str) 162 | 163 | try: 164 | self.output = subprocess.check_output( 165 | cmd, 166 | shell=True, 167 | stderr=subprocess.STDOUT, 168 | cwd=os.curdir 169 | ) 170 | 171 | except subprocess.CalledProcessError as e: 172 | self.code = e.returncode 173 | self.output = e.output 174 | 175 | def backup(self, *tables): 176 | subcommand = "backup" 177 | arg_str = "{} {} {}".format(subcommand, self.arg_str, " ".join(tables)) 178 | return self.run(arg_str) 179 | 180 | def restore(self): 181 | subcommand = "restore" 182 | arg_str = "{} {}".format(subcommand, self.arg_str) 183 | return self.run(arg_str) 184 | 185 | 186 | class DumpTest(unittest.TestCase): 187 | def setUp(self): 188 | Foo().install() 189 | Bar().install() 190 | 191 | def test_default_help(self): 192 | c = Client() 193 | c.run("") 194 | self.assertLess(0, c.code) 195 | 196 | def test_table_not_exist(self): 197 | Foo().delete() 198 | c = Client() 199 | c.backup(Foo.table_name) 200 | self.assertEqual(1, c.code, c.output) 201 | 202 | def test_full_table_backup_and_restore(self): 203 | for x in range(100): 204 | f = Foo(bar=x) 205 | _id = f.save() 206 | self.assertLess(0, _id) 207 | 208 | self.assertEqual(100, Foo().count()) 209 | 210 | c = Client() 211 | c.backup(Foo.table_name) 212 | self.assertEqual(1, len(c.files)) 213 | for path in c.files: 214 | self.assertLess(0, os.path.getsize(path)) 215 | 216 | c.restore() 217 | self.assertEqual(100, Foo().count()) 218 | 219 | f = Foo(bar=101) 220 | pk = f.save() 221 | self.assertLess(100, pk) 222 | 223 | 224 | self.setUp() 225 | c.restore() 226 | self.assertEqual(100, Foo().count()) 227 | 228 | def test_multi_table_backup(self): 229 | count = 10 230 | for x in range(count): 231 | f = Foo(bar=x) 232 | f.save() 233 | b = Bar(foo=x) 234 | b.save() 235 | 236 | self.assertEqual(count, Foo().count()) 237 | self.assertEqual(count, Bar().count()) 238 | 239 | c = Client() 240 | c.backup(Foo.table_name, Bar.table_name) 241 | 242 | self.setUp() 243 | c.restore() 244 | self.assertEqual(count, Foo().count()) 245 | self.assertEqual(count, Bar().count()) 246 | 247 | 248 | if __name__ == '__main__': 249 | unittest.main() 250 | 251 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # http://docs.python.org/distutils/setupscript.html 3 | # http://docs.python.org/2/distutils/examples.html 4 | 5 | from setuptools import setup, find_packages 6 | import re 7 | import os 8 | from codecs import open 9 | 10 | 11 | name = 'dump' 12 | with open(os.path.join(name, "__init__.py"), encoding='utf-8') as f: 13 | version = re.search(r"^__version__\s*=\s*[\'\"]([^\'\"]+)", f.read(), flags=re.I | re.M).group(1) 14 | 15 | long_description = "" 16 | if os.path.isfile('README.rst'): 17 | with open('README.rst', encoding='utf-8') as f: 18 | long_description = f.read() 19 | 20 | setup( 21 | name=name, 22 | version=version, 23 | description='Wrapper around psql and pg_dump to make it easier to backup/restore a PostgreSQL database', 24 | long_description=long_description, 25 | author='Jay Marcyes', 26 | author_email='jay@marcyes.com', 27 | url='http://github.com/jaymon/{}'.format(name), 28 | packages=find_packages(), 29 | #py_modules=[name], 30 | license="MIT", 31 | #install_requires=[], 32 | tests_require=["dsnparse", "psycopg2", "testdata"], 33 | classifiers=[ # https://pypi.python.org/pypi?:action=list_classifiers 34 | 'Development Status :: 4 - Beta', 35 | 'Environment :: Web Environment', 36 | 'Intended Audience :: Developers', 37 | 'License :: OSI Approved :: MIT License', 38 | 'Operating System :: OS Independent', 39 | 'Topic :: Database', 40 | 'Topic :: Software Development :: Libraries', 41 | 'Topic :: Utilities', 42 | 'Programming Language :: Python :: 2.7', 43 | 'Programming Language :: Python :: 3.6', 44 | ], 45 | entry_points = { 46 | 'console_scripts': ['{} = {}.__main__:console'.format(name, name)] 47 | } 48 | ) 49 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | # pyenv install 2.7.14 2 | # pyenv install 3.6.3 3 | # pyenv global 2.7.14 3.6.3 4 | # tox 5 | [tox] 6 | envlist=py27,py36 7 | [testenv] 8 | passenv= 9 | DUMP_DSN 10 | deps= 11 | testdata 12 | pyt 13 | psycopg2 14 | dsnparse 15 | commands = pyt -ad 16 | 17 | --------------------------------------------------------------------------------