Add psycopg2 extra, patch psycogreen in entry points. (#915)

bin/baseplate-script

@@ -12,6 +12,15 @@ from baseplate.server.monkey import patch_stdlib_queues
 patch_all()
 patch_stdlib_queues()
 
+try:
+    import psycopg2
+    from psycogreen.gevent import patch_psycopg
+
+    patch_psycopg()
+except ImportError:
+    pass
+
+
 from baseplate.server import load_and_run_script
 
 load_and_run_script()

bin/baseplate-serve

@@ -12,6 +12,15 @@ from baseplate.server.monkey import patch_stdlib_queues
 patch_all()
 patch_stdlib_queues()
 
+try:
+    import psycopg2
+    from psycogreen.gevent import patch_psycopg
+
+    patch_psycopg()
+except ImportError:
+    pass
+
+
 from baseplate.server import load_app_and_run_server
 
 load_app_and_run_server()

bin/baseplate-shell

@@ -12,6 +12,15 @@ from baseplate.server.monkey import patch_stdlib_queues
 patch_all()
 patch_stdlib_queues()
 
+try:
+    import psycopg2
+    from psycogreen.gevent import patch_psycopg
+
+    patch_psycopg()
+except ImportError:
+    pass
+
+
 from baseplate.server import load_and_run_shell
 
 load_and_run_shell()

bin/baseplate-tshell

@@ -12,6 +12,15 @@ from baseplate.server.monkey import patch_stdlib_queues
 patch_all()
 patch_stdlib_queues()
 
+try:
+    import psycopg2
+    from psycogreen.gevent import patch_psycopg
+
+    patch_psycopg()
+except ImportError:
+    pass
+
+
 from baseplate.server import load_and_run_shell
 
 print("baseplate-tshell has been renamed to baseplate-shell. Please use that in the future!")

pyproject.toml

@@ -39,6 +39,8 @@ sentry-sdk = { version = ">=1.35.0,<2.0", optional = true }
 sqlalchemy = { version = ">=1.4.49,<2", optional = true }
 thrift-unofficial = ">=0.19.0,<1.0"
 typing-extensions = "^4.11.0"
+psycopg2 = "^2.0.0"
+psycogreen = "^1.0.0"
 
 [tool.poetry.extras]
 amqp = ["kombu"]
@@ -55,6 +57,7 @@ requests = ["advocate"]
 s3fetcher = []  # Kept for backwards compatibility, these are now main requirements
 sentry = ["sentry-sdk"]
 sql = ["sqlalchemy"]
+psycopg2 = ["psycopg2", "psycogreen"]
 zookeeper = ["kazoo"]
 
 [tool.poetry.group.dev.dependencies]