From 54895ebed4587fe8d4ca60266516c7c064b7fb06 Mon Sep 17 00:00:00 2001
From: Mike Alfare <mike.alfare@dbtlabs.com>
Date: Wed, 5 Jun 2024 20:04:05 -0400
Subject: [PATCH] add post build hook to swap out psycopg2-binary for psycopg2
 based on DBT_PSYCOPG2_NAME

---
 hatch_hooks.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)
 create mode 100644 hatch_hooks.py

diff --git a/hatch_hooks.py b/hatch_hooks.py
new file mode 100644
index 00000000..fc036128
--- /dev/null
+++ b/hatch_hooks.py
@@ -0,0 +1,35 @@
+import os
+import subprocess
+import sys
+
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+from hatchling.plugin import hookimpl
+
+
+class Psycopg2NoBinary(BuildHookInterface):
+    """
+    Custom build hook to install psycopg2 instead of psycopg2-binary based on the presence of `DBT_PSYCOPG2_NAME`.
+    This is necessary as psycopg2-binary is better for local development, but psycopg2 is better for production.
+    """
+
+    PLUGIN_NAME = "psycopg2"
+
+    def finalize(self, version, build_data, artifact_path) -> None:
+        if os.getenv("DBT_PSYCOPG2_NAME", "") == "psycopg2":
+            psycopg2_binary_pinned = [
+                package
+                for package in build_data["dependencies"]
+                if package.startswith("psycopg2-binary")
+            ].pop()
+            psycopg2_pinned = psycopg2_binary_pinned.replace("-binary", "")
+            subprocess.check_call(
+                [sys.executable, "-m", "pip", "-y", "uninstall", "psycopg2-binary"]
+            )
+            subprocess.check_call(
+                [sys.executable, "-m", "pip", "-y", "install", f'"{psycopg2_pinned}"']
+            )
+
+
+@hookimpl
+def hatch_register_build_hook():
+    return Psycopg2NoBinary