Sabyasachi
Sabyasachi

Reputation: 1554

Pass a connection object from a process to the main process in python

I have an application which creates database connection on demand. But I wanted a functionality, that if creating connection takes more than a certain time, it should be timed out. So I used multiprocess, to spawn a process, and join it after a certain time, and check the state of the process. This process works well, but I am not able to retrieve the connection object from the spawned process because it's not serializable. Any help?

import psycopg2
import multiprocessing
from multiprocessing import Process

config = {
    "user": "xxx",
    "password": "xxx",
    "host": "xxx",
    "database": "xxx",
}


def create_postgres_connection(config, return_dict):
    # Connect to the database, we have to timeout this function if it takes too long.
    host = config['host']
    database = config['database']
    username = config['user']
    password = config['password']
    connection = psycopg2.connect(host=host, database=database, user=username, password=password)
    return_dict['connection'] = connection


def run_with_limited_time(func, args, kwargs, time):
    """Runs a function with time limit

    :param func: The function to run
    :param args: The functions args, given as tuple
    :param kwargs: The functions keywords, given as dict
    :param time: The time limit in seconds
    :return: True if the function ended successfully. False if it was terminated.
    """
    p = Process(target=func, args=args, kwargs=kwargs)
    p.start()
    p.join(5)
    if p.is_alive():
        p.terminate()
        print('Timed Out')
        return False
    return True


if __name__ == '__main__':
    manager = multiprocessing.Manager()
    return_dict = manager.dict()
    run_with_limited_time(create_postgres_connection, (config, return_dict), {}, 3)
    print(return_dict)

Upvotes: 0

Views: 1042

Answers (1)

MyNameIsCaleb
MyNameIsCaleb

Reputation: 4489

You won't be able to do this using multiprocessing because it relies on Pickling an object in order to transfer it between processes. And as you found out, connection objects cannot be pickled.

However, I don't think that you need to do this anyway because Postgres allows for a connect_timeout parameter when connecting, which is really the problem you need to solve. The connect_timeout is specified in seconds. [postgres docs]

According to the psycopg2 docs you can pass any database specific parameters as a keyword argument.

It would look something like:

def create_postgres_connection(config, connection):
    # Connect to the database
    host = config['host']
    database = config['database']
    username = config['user']
    password = config['password']
    connection = psycopg2.connect(host=host, database=database, user=username, 
                                  password=password, connect_timeout=3)
    return connection

Upvotes: 1

Related Questions