Creating a new Postgres database from a template DB using Python psycopg2

import psycopg2
#from psycopg2 import sql
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

def create_db(db_name, db_username, db_password):
      con = psycopg2.connect(
            dbname='postgres',
            user='postgres',
            host='localhost',
            password='secret', 
            port='5432')

      con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
      cur = con.cursor()
      DB_NAME = db_name
      DB_USERNAME = db_username
      DB_PASSWORD = db_password
      
      #terminate any process on the template1 db
      cur.execute(f"""SELECT pg_terminate_backend(pg_stat_activity.pid)
                  FROM pg_stat_activity
                  WHERE pg_stat_activity.datname = 'template1'
                  AND pid <> pg_backend_pid();""")

      cur.execute(f"""CREATE DATABASE {DB_NAME} WITH TEMPLATE template1 OWNER postgres;""")
      cur.execute(f"""CREATE USER {DB_USERNAME} WITH PASSWORD '{DB_PASSWORD}';""")
      cur.execute(f"""GRANT ALL PRIVILEGES ON DATABASE {DB_NAME} TO {DB_USERNAME};""")

      cur.close()
      con.close()

      con = psycopg2.connect(
            dbname=DB_NAME,
            user='postgres',
            host='localhost',
            password='secret',
            port='5432')

      con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
      cur = con.cursor()

      cur.execute(f"""REVOKE ALL ON DATABASE {DB_NAME} FROM {DB_USERNAME};
                  REVOKE CONNECT ON DATABASE {DB_NAME} FROM PUBLIC; 
                  GRANT ALL PRIVILEGES ON DATABASE {DB_NAME} TO {DB_USERNAME};
                  GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO {DB_USERNAME};
                  GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO {DB_USERNAME};
                  ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO {DB_USERNAME};
                  ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO {DB_USERNAME};""")

      cur.close()
      con.close()

Leave a Reply

Your email address will not be published. Required fields are marked *