Another Nix Success Story

Thursday, August 5, 2021

•

496 words

•

3 minute read

Yesterday I found myself in a position where I needed to know the dependency graph of tables within a Postgres database.

I figured that someone else had written a program to solve this problem, and some quick Google-fu revealed that I was correct.

I happened upon this post that contained a Python script that generated the dependency graph I was looking for and dumped it to Graphviz DOT. Beautiful!

Here's the script:

from optparse import OptionParser, OptionGroup

import psycopg2
import sys


def writedeps(cursor, tbl):
    sql = """SELECT
        tc.constraint_name, tc.table_name, kcu.column_name,
        ccu.table_name AS foreign_table_name,
        ccu.column_name AS foreign_column_name
    FROM
        information_schema.table_constraints AS tc
    JOIN information_schema.key_column_usage AS kcu ON
        tc.constraint_name = kcu.constraint_name
    JOIN information_schema.constraint_column_usage AS ccu ON
        ccu.constraint_name = tc.constraint_name
    WHERE constraint_type = 'FOREIGN KEY' AND tc.table_name = '%s'"""
    cursor.execute(sql % tbl)
    for row in cursor.fetchall():
        constraint, table, column, foreign_table, foreign_column = row
        print '"%s" -> "%s" [label="%s"];' % (tbl, foreign_table, constraint)


def get_tables(cursor):
    cursor.execute("SELECT tablename FROM pg_tables WHERE schemaname='public'")
    for row in cursor.fetchall():
        yield row[0]


def main():
    parser = OptionParser()

    group = OptionGroup(parser, "Database Options")
    group.add_option("--dbname", action="store", dest="dbname",
            help="The database name.")
    group.add_option("--dbhost", action="store", dest="dbhost",
            default="localhost",  help="The database host.")
    group.add_option("--dbuser", action="store", dest="dbuser",
            help="The database username.")
    group.add_option("--dbpass", action="store", dest="dbpass",
            help="The database password.")
    parser.add_option_group(group)

    (options, args) = parser.parse_args()

    if not options.dbname:
        print "Please supply a database name, see --help for more info."
        sys.exit(1)

    try:
        conn = psycopg2.connect("dbname='%s' user='%s' host='%s' password='%s'"
            % (options.dbname, options.dbuser, options.dbhost, options.dbpass))
    except psycopg2.OperationalError, e:
        print "Failed to connect to database,",
        print "perhaps you need to supply auth details:\n %s" % str(e)
        print "Use --help for more info."
        sys.exit(1)

    cursor = conn.cursor()

    print "Digraph F {\n"
    print 'ranksep=1.0; size="18.5, 15.5"; rankdir=LR;'
    for i in get_tables(cursor):
        writedeps(cursor, i)
    print "}"

    sys.exit(0)


if __name__ == "__main__":
    main()

It was herein that I discovered two problems:

The script depends on psycopg2, which means I need to muck about with Python dependency management
I can tell from the print statements that the script is written in Python 2, and my work MacBook has Python 3

My immediate thought was to use Nix to solve this. I grabbed a shell.nix template from the Packaging/Python section of the NixOS wiki and added the psycopg2 dependency along with graphviz:

 with import <nixpkgs> {};

 stdenv.mkDerivation {
   name = "pip-env";
   buildInputs = [
+    graphviz
+
     # System requirements.
     readline

     # Python requirements (enough to get a virtualenv going).
     python27Full
     python27Packages.virtualenv
     python27Packages.pip
+    python27Packages.psycopg2
   ];
   src = null;
   shellHook = ''
     # Allow the use of wheels.
     SOURCE_DATE_EPOCH=$(date +%s)

     # Augment the dynamic linker path
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${R}/lib/R/lib:$ {readline}/lib
   '';
 }

I was then able to drop into a nix-shell and run the script:

python database_dependency_order.py --dbname <NAME> --dbhost <HOST> --dbuser <USER> --dbpass <PASSWORD> | dot -Tpng > deps.png

Easy as Py Nix.

Marshall Bowers

Conjurer of code. Devourer of art. Pursuer of æsthetics.

Another Nix Success Story