client.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  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")
  12. user = settings_dict.get("USER")
  13. passwd = settings_dict.get("PASSWORD")
  14. passfile = options.get("passfile")
  15. service = options.get("service")
  16. sslmode = options.get("sslmode")
  17. sslrootcert = options.get("sslrootcert")
  18. sslcert = options.get("sslcert")
  19. sslkey = options.get("sslkey")
  20. if not dbname and not service:
  21. # Connect to the default 'postgres' db.
  22. dbname = "postgres"
  23. if user:
  24. args += ["-U", user]
  25. if host:
  26. args += ["-h", host]
  27. if port:
  28. args += ["-p", str(port)]
  29. args.extend(parameters)
  30. if dbname:
  31. args += [dbname]
  32. env = {}
  33. if passwd:
  34. env["PGPASSWORD"] = str(passwd)
  35. if service:
  36. env["PGSERVICE"] = str(service)
  37. if sslmode:
  38. env["PGSSLMODE"] = str(sslmode)
  39. if sslrootcert:
  40. env["PGSSLROOTCERT"] = str(sslrootcert)
  41. if sslcert:
  42. env["PGSSLCERT"] = str(sslcert)
  43. if sslkey:
  44. env["PGSSLKEY"] = str(sslkey)
  45. if passfile:
  46. env["PGPASSFILE"] = str(passfile)
  47. return args, (env or None)
  48. def runshell(self, parameters):
  49. sigint_handler = signal.getsignal(signal.SIGINT)
  50. try:
  51. # Allow SIGINT to pass to psql to abort queries.
  52. signal.signal(signal.SIGINT, signal.SIG_IGN)
  53. super().runshell(parameters)
  54. finally:
  55. # Restore the original SIGINT handler.
  56. signal.signal(signal.SIGINT, sigint_handler)