client.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. import os
  2. import signal
  3. import subprocess
  4. from django.db.backends.base.client import BaseDatabaseClient
  5. class DatabaseClient(BaseDatabaseClient):
  6. executable_name = 'psql'
  7. @classmethod
  8. def runshell_db(cls, conn_params):
  9. args = [cls.executable_name]
  10. host = conn_params.get('host', '')
  11. port = conn_params.get('port', '')
  12. dbname = conn_params.get('database', '')
  13. user = conn_params.get('user', '')
  14. passwd = conn_params.get('password', '')
  15. sslmode = conn_params.get('sslmode', '')
  16. sslrootcert = conn_params.get('sslrootcert', '')
  17. sslcert = conn_params.get('sslcert', '')
  18. sslkey = conn_params.get('sslkey', '')
  19. if user:
  20. args += ['-U', user]
  21. if host:
  22. args += ['-h', host]
  23. if port:
  24. args += ['-p', str(port)]
  25. args += [dbname]
  26. sigint_handler = signal.getsignal(signal.SIGINT)
  27. subprocess_env = os.environ.copy()
  28. if passwd:
  29. subprocess_env['PGPASSWORD'] = str(passwd)
  30. if sslmode:
  31. subprocess_env['PGSSLMODE'] = str(sslmode)
  32. if sslrootcert:
  33. subprocess_env['PGSSLROOTCERT'] = str(sslrootcert)
  34. if sslcert:
  35. subprocess_env['PGSSLCERT'] = str(sslcert)
  36. if sslkey:
  37. subprocess_env['PGSSLKEY'] = str(sslkey)
  38. try:
  39. # Allow SIGINT to pass to psql to abort queries.
  40. signal.signal(signal.SIGINT, signal.SIG_IGN)
  41. subprocess.run(args, check=True, env=subprocess_env)
  42. finally:
  43. # Restore the original SIGINT handler.
  44. signal.signal(signal.SIGINT, sigint_handler)
  45. def runshell(self):
  46. DatabaseClient.runshell_db(self.connection.get_connection_params())