CODEWITHSUNDEEP
pythonpostgresqlpsycopg2
NicLovin
NicLovin

Reputation: 355

Creating a method to connect to postgres database in python

I'm working on a python program with functionality such as inserting and retrieving values from a postgres database using psycopg2. The issue is that every time I want to create a query I have to connect to the database so the following code snippet is present multiple times throughout the file:

# Instantiate Connection
    try:
        conn = psycopg2.connect(
            user=userName,
            password=passwrd,
            host=hostAddr,
            database=dbName
        )

        # Instantiate Cursor
        cur = conn.cursor()
        return cur

    except psycopg2.Error as e:
        print(f"Error connecting to Postgres Platform: {e}")
        sys.exit(1)

My question is:

  1. Is there a way I could just create a method to call every time I wish to connect to the database? I've tried creating one but I get a bunch of errors since variables cur and conn are not global

  2. Could I just connect to the database once at the beginning of the program and keep the connection open for the entire time that the program is running? This seems like the easiest option but I am not sure if it would be bad practice (for reference the program will be running 24/7 so I assumed it would be better to only connect when a query is being made).

Thanks for the help.

Upvotes: 1

Views: 3209

Answers (4)

Maurice Meyer
Maurice Meyer

Reputation: 18136

You could wrap your own database handling class in a context manager, so you can manage the connections in a single place:

import psycopg2
import traceback
from psycopg2.extras import RealDictCursor


class Postgres(object):
    def __init__(self, *args, **kwargs):
        self.dbName = args[0] if len(args) > 0 else 'prod'
        self.args = args

    def _connect(self, msg=None):
        if self.dbName == 'dev':
            dsn = 'host=127.0.0.1 port=5556 user=xyz password=xyz dbname=development'
        else:
            dsn = 'host=127.0.0.1 port=5557 user=xyz password=xyz dbname=production'

        try:
            self.con = psycopg2.connect(dsn)
            self.cur = self.con.cursor(cursor_factory=RealDictCursor)
        except:
            traceback.print_exc()

    def __enter__(self, *args, **kwargs):
        self._connect()
        return (self.con, self.cur)

    def __exit__(self, *args):
        for c in ('cur', 'con'):
            try:
                obj = getattr(self, c)
                obj.close()
            except:
                pass # handle it silently!?
        self.args, self.dbName = None, None

Usage:

with Postgres('dev') as (con, cur):
    print(con)
    print(cur.execute('select 1+1'))
print(con) # verify connection gets closed!

Out:

<connection object at 0x109c665d0; dsn: '...', closed: 0>
[RealDictRow([('sum', 2)])]
<connection object at 0x109c665d0; dsn: '...', closed: 1>

Upvotes: 4

jjanes
jjanes

Reputation: 44363

If I wrap your code-fragment into a function definition, I don't get "a bunch of errors since variables cur and conn are not global". Why would they need to be global? Whatever the error was, you removed it from your code fragment before posting it.

Your try-catch doesn't make any sense to me. Catching an error just to hide the calling site and then bail out seems like the opposite of helpful.

When to connect depends on how you structure your transactions, how often you do them, and what you want to do if your database ever restarts in the middle of a program execution.

Upvotes: 0

spencer.pinegar
spencer.pinegar

Reputation: 460

I have been down this road lots before and you may be reinventing the wheel. I would highly recommend you use a ORM like [Django][1] or if you need to interact with a database - it handles all this stuff for you using best practices. It is some learning up front but I promise it pays off.

If you don't want to use Django, you can use this code to get or create the connection and the context manager of cursors to avoid errors with

import pscyopg2

CONN = None

def create_or_get_connection():
    global CONN
    if CONN is None or CONN.closed:
        CONN = psycopg2.connect(...)
    return CONN

def run_sql(sql):
    con = create_or_get_connection()
    with conn.cursor() as curs:
        return curs.execute(sql)

This will allow you simply to run sql statements directly to the DB without worrying about connection or cursor issues.

Upvotes: 0

Kyle Parsons
Kyle Parsons

Reputation: 1525

It shouldn't be too bad to keep a connection open. The server itself should be responsible for closing connections it thinks have been around for too long or that are too inactive. We then just need to make our code resilient in case the server has closed the connection:

import pscyopg2

CONN = None

def create_or_get_connection():
    global CONN
    if CONN is None or CONN.closed:
        CONN = psycopg2.connect(...)
    return CONN

Upvotes: 0

Related Questions