123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354 |
- import os
- import signal
- import subprocess
- from django.db.backends.base.client import BaseDatabaseClient
- class DatabaseClient(BaseDatabaseClient):
- executable_name = 'psql'
- @classmethod
- def runshell_db(cls, conn_params):
- args = [cls.executable_name]
- host = conn_params.get('host', '')
- port = conn_params.get('port', '')
- dbname = conn_params.get('database', '')
- user = conn_params.get('user', '')
- passwd = conn_params.get('password', '')
- sslmode = conn_params.get('sslmode', '')
- sslrootcert = conn_params.get('sslrootcert', '')
- sslcert = conn_params.get('sslcert', '')
- sslkey = conn_params.get('sslkey', '')
- if user:
- args += ['-U', user]
- if host:
- args += ['-h', host]
- if port:
- args += ['-p', str(port)]
- args += [dbname]
- sigint_handler = signal.getsignal(signal.SIGINT)
- subprocess_env = os.environ.copy()
- if passwd:
- subprocess_env['PGPASSWORD'] = str(passwd)
- if sslmode:
- subprocess_env['PGSSLMODE'] = str(sslmode)
- if sslrootcert:
- subprocess_env['PGSSLROOTCERT'] = str(sslrootcert)
- if sslcert:
- subprocess_env['PGSSLCERT'] = str(sslcert)
- if sslkey:
- subprocess_env['PGSSLKEY'] = str(sslkey)
- try:
- # Allow SIGINT to pass to psql to abort queries.
- signal.signal(signal.SIGINT, signal.SIG_IGN)
- subprocess.run(args, check=True, env=subprocess_env)
- finally:
- # Restore the original SIGINT handler.
- signal.signal(signal.SIGINT, sigint_handler)
- def runshell(self):
- DatabaseClient.runshell_db(self.connection.get_connection_params())
|