Передача и возврат соединения и курсора для Postgresql in Python

Question

Я написал файл connection.py следующим образом:

from configparser import ConfigParser
import psycopg2
import pandas as pd


def get_section(a):
    section_dict = {
            1:'one_db',
            2:'two_db',
            3:'three_db'
            }
    return section_dict.get(a)

def config(a):

    filename='config.ini',
    section=get_section(a)
    # create a parser
    parser = ConfigParser()
    # read config file
    parser.read(filename)

    # get section, default to postgresql
    db = {}
    if parser.has_section(section):
        params = parser.items(section)
        for param in params:
            db[param[0]] = param[1]
    else:
        raise Exception('Section {0} not found in the {1} file'.format(section, filename))

    return db

def get_data(query,a):
    """ Connect to the PostgreSQL database server """

    connection = None
    try:
        # read connection parameters
        params = config(a)

        # connect to the PostgreSQL server
        print('Connecting to '+ get_section(a) )
        connection = psycopg2.connect(**params)
        print('Connected')

        # create a cursor
        cursor = connection.cursor()

        cursor.execute(query)
        df = pd.DataFrame(cursor.fetchall())
        colnames = [desc[0] for desc in cursor.description]
        df.columns = colnames
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
    finally:
        print('Closing connection...')
        cursor.close()
        connection.close()
        print('Connection closed')

    return df

Я использую функцию get_data в разных местах для получения данных из базы данных следующим образом

from connection import get_data
query = "some_query"
df = get_data(query,1)

Работает хорошо. Но в некоторых местах мне приходится постоянно использовать множество запросов один за другим и получать данные. Проблема с этой структурой - каждый раз, когда она подключается и отключается. Поэтому, если у меня есть десять запросов в одном файле, подключение и отключение происходит десять раз. Как изменить код так, чтобы я мог иметь следующий формат:

connect(1)
df1 = get_data(query1)
df2 = get_data(query2)
df3 = get_data(query3)
.
.
.
.
.
.
disconnect()

Число 1 или 2 или 3 в Connect (X), что у меня есть три базы данных на выбор. Я надеюсь, что я в порядке с проблемой. Готов также поделиться любыми другими деталями.

Я попробовал следующее, но не уверен, завершено ли это:
Используйте функцию get_data, где вы возвращаете cur и con. Я внес следующие изменения в файл connection.py

def get_connection(a):
    """ Connect to the PostgreSQL database server """

    connection = None
    try:
        # read connection parameters
        params = config(a)

        # connect to the PostgreSQL server
        print('Connecting to '+ get_section(a) )
        connection = psycopg2.connect(**params)
        print('Connected')

        # create a cursor
        cursor = connection.cursor()

        return (connection,cursor)
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)

def get_data(cursor,query):
    xcurs.execute(query)
    df = pd.DataFrame(xcurs.fetchall())
    colnames = [desc[0] for desc in xcurs.description]
    df.columns = colnames
    return df

Теперь в любом файле я делаю это:

conn,curs = get_connection(1)
df1 = get_data(curs,query1)
df2= get_data(curs,query2)
df3 = get_data(curs,query3)
df4 = get_data(curs,query4)
df5 = get_data(curs,query5)
.
.
.
.
.
curs.close()
conn.close()

Правильно ли этот метод?

Answer 1

Есть несколько решений в psycopg2 из коробки. Вам нужно получить любую реализацию абстракции пула соединений (т.е. psycopg2.pool.SimpleConnectionPool) и использовать метод getconn() для получения соединения.

pool = SimpleConnectionPool(**params)

def get_data():
    connection = pool.getconn()