machine_learn/util/pgsql_util.py

#!/usr/bin/python
# -*- coding:utf-8 -*-

import sys
import os
from pgdb import Connection
import json

sys.path.append(os.path.abspath('..'))
from util.config import config


def get_conn():
    conn = Connection(database=config.get('pgsql', 'db'), host=config.get('pgsql', 'host'),
                            user=config.get('pgsql', 'user'), password=config.get('pgsql', 'password'))
    return conn


def insert(sql):
    try:
        conn = get_conn()
        cur = conn.cursor()
        cur.execute(sql)
        conn.commit()
    except Exception as e:
        print(e)
    finally:
        cur.close()
        conn.close()


def update(sql):
    try:
        conn = get_conn()
        conn.execute(sql)
        conn.commit()
    except Exception as e:
        print(e)
    finally:
        conn.close()


def get_rows(sql):
    try:
        conn = get_conn()
        return conn.query(sql)
    except Exception as e:
        print(e)
    finally:
        conn.close()


def get_total(sql):
    conn = None
    try:
        conn = get_conn()
        result = conn.query(sql)
        return int(result[0]['count'])
    except Exception as e:
        print(e)
    finally:
        conn.close()


def demo():
    sql = "select aid,other_info FROM articles where aid = '5afa6e5a-9004-4182-a743-87d77c97b451'"
    rows = get_rows(sql)
    for row in rows:
        print(row)
        other_info = row['other_info']
        print(other_info)
        other_info_json = json.loads(other_info)
        print(other_info_json.get('abstract'))


def demo1():
    sql = "select count(*) FROM articles"
    total = get_total(sql)
    print(total)


PAGE_SIZE = 50


def demo2(offset):
    sql = "select aid,other_info FROM articles order by aid LIMIT %s OFFSET %s"
    update_sql = '''
        UPDATE articles
        set other_info = %s
        WHERE aid = %s
    '''
    rows = get_rows(sql % (PAGE_SIZE, 0))

    conn = get_conn()

    for row in rows:
        aid = row['aid']
        other_info = json.loads(row['other_info'])
        if 'abstract' in other_info:
            del other_info['abstract']
        if 'notes' in other_info:
            del other_info['notes']
        if 'content' in other_info:
            del other_info['content']
        print(json.dumps(other_info), aid)

        conn.execute(update_sql, (json.dumps(other_info), aid))
    conn.close()


def demo3():
    for i in range(0, 100, 50):
        demo2(i)


if __name__ == "__main__":
    demo1()