Skip to content

Commit

Permalink
add post build hook to swap out psycopg2-binary for psycopg2 based on…
Browse files Browse the repository at this point in the history
… DBT_PSYCOPG2_NAME
  • Loading branch information
mikealfare committed Jun 6, 2024
1 parent 97f7ef7 commit 54895eb
Showing 1 changed file with 35 additions and 0 deletions.
35 changes: 35 additions & 0 deletions hatch_hooks.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 54895eb

Please sign in to comment.