database.py 文件源码

def get_pkey(table, conn=None, **kwargs):
    '''
    Return the primary key column for a table as a named tuple with fields
    "column" and "type"

    If no primary key, return None

    Ref: https://wiki.postgresql.org/wiki/Retrieve_primary_key_columns
    '''

    p_key = namedtuple('PrimaryKey', ['column', 'type'])
    cur = conn.cursor()

    try:
        cur.execute(sql.SQL('''
            SELECT a.attname, format_type(a.atttypid, a.atttypmod) AS data_type
            FROM   pg_index i
            JOIN   pg_attribute a ON a.attrelid = i.indrelid
                                 AND a.attnum = ANY(i.indkey)
            WHERE  i.indrelid = {}::regclass
            AND    i.indisprimary;
        ''').format(
            sql.Literal(table)))

        data = cur.fetchall()[0]
        return p_key(column=data[0], type=data[1])
    except IndexError:
        return None
    except (psycopg2.ProgrammingError, psycopg2.InternalError) as e:
        conn.rollback()
        return None