client.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. import signal
  2. from django.db.backends.base.client import BaseDatabaseClient
  3. class DatabaseClient(BaseDatabaseClient):
  4. executable_name = 'psql'
  5. @classmethod
  6. def settings_to_cmd_args_env(cls, settings_dict, parameters):
  7. args = [cls.executable_name]
  8. options = settings_dict.get('OPTIONS', {})
  9. host = settings_dict.get('HOST')
  10. port = settings_dict.get('PORT')
  11. dbname = settings_dict.get('NAME') or 'postgres'
  12. user = settings_dict.get('USER')
  13. passwd = settings_dict.get('PASSWORD')
  14. sslmode = options.get('sslmode')
  15. sslrootcert = options.get('sslrootcert')
  16. sslcert = options.get('sslcert')
  17. sslkey = options.get('sslkey')
  18. if user:
  19. args += ['-U', user]
  20. if host:
  21. args += ['-h', host]
  22. if port:
  23. args += ['-p', str(port)]
  24. args += [dbname]
  25. args.extend(parameters)
  26. env = {}
  27. if passwd:
  28. env['PGPASSWORD'] = str(passwd)
  29. if sslmode:
  30. env['PGSSLMODE'] = str(sslmode)
  31. if sslrootcert:
  32. env['PGSSLROOTCERT'] = str(sslrootcert)
  33. if sslcert:
  34. env['PGSSLCERT'] = str(sslcert)
  35. if sslkey:
  36. env['PGSSLKEY'] = str(sslkey)
  37. return args, (env or None)
  38. def runshell(self, parameters):
  39. sigint_handler = signal.getsignal(signal.SIGINT)
  40. try:
  41. # Allow SIGINT to pass to psql to abort queries.
  42. signal.signal(signal.SIGINT, signal.SIG_IGN)
  43. super().runshell(parameters)
  44. finally:
  45. # Restore the original SIGINT handler.
  46. signal.signal(signal.SIGINT, sigint_handler)