Skip to content

Commit

Permalink
Use PG* environment rather than command-line for PostgreSQL
Browse files Browse the repository at this point in the history
Closes jazzband#384.
  • Loading branch information
Natureshadow committed Feb 17, 2021
1 parent 9d1909c commit 2ed8914
Showing 1 changed file with 14 additions and 24 deletions.
38 changes: 14 additions & 24 deletions dbbackup/db/postgresql.py
Original file line number Diff line number Diff line change
@@ -1,31 +1,13 @@
from urllib.parse import quote
import logging
import os

from .base import BaseCommandDBConnector
from .exceptions import DumpError

logger = logging.getLogger('dbbackup.command')


def create_postgres_uri(self):
host = self.settings.get('HOST')
if not host:
raise DumpError('A host name is required')

dbname = self.settings.get('NAME') or ''
user = quote(self.settings.get('USER') or '')
password = self.settings.get('PASSWORD') or ''
password = ':{}'.format(quote(password)) if password else ''
if not user:
password = ''
else:
host = '@' + host

port = ':{}'.format(self.settings.get('PORT')) if self.settings.get('PORT') else ''
dbname = f'--dbname=postgresql://{user}{password}{host}{port}/{dbname}'
return dbname


class PgDumpConnector(BaseCommandDBConnector):
"""
PostgreSQL connector, it uses pg_dump`` to create an SQL text file
Expand All @@ -37,9 +19,20 @@ class PgDumpConnector(BaseCommandDBConnector):
single_transaction = True
drop = True

@property
def dump_env(self):
env = os.environ.copy()
for var in ('HOST', 'PORT', 'USER', 'PASSWORD'):
if self.settings.get(var, None):
env['PG%s' % var] = self.settings.get(var)
env['PGDATABASE'] = self.settings.get(var, 'postgres')

return env

restore_env = dump_env

def _create_dump(self):
cmd = '{} '.format(self.dump_cmd)
cmd = cmd + create_postgres_uri(self)

for table in self.exclude:
cmd += ' --exclude-table-data={}'.format(table)
Expand All @@ -52,7 +45,6 @@ def _create_dump(self):

def _restore_dump(self, dump):
cmd = '{} '.format(self.restore_cmd)
cmd = cmd + create_postgres_uri(self)

# without this, psql terminates with an exit value of 0 regardless of errors
cmd += ' --set ON_ERROR_STOP=on'
Expand Down Expand Up @@ -101,7 +93,6 @@ class PgDumpBinaryConnector(PgDumpConnector):

def _create_dump(self):
cmd = '{} '.format(self.dump_cmd)
cmd = cmd + create_postgres_uri(self)

cmd += ' --format=custom'
for table in self.exclude:
Expand All @@ -111,8 +102,7 @@ def _create_dump(self):
return stdout

def _restore_dump(self, dump):
dbname = create_postgres_uri(self)
cmd = '{} {}'.format(self.restore_cmd, dbname)
cmd = '{} {}'.format(self.restore_cmd, self.settings['NAME'])

if self.single_transaction:
cmd += ' --single-transaction'
Expand Down

0 comments on commit 2ed8914

Please sign in to comment.