ottawa/test_db.py at master · GovLab/ottawa · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
#!/usr/bin/python
import psycopg2
import sys


def test_query(cursor):
    print >> sys.stdout, "Testing query."
    cursor.execute("SELECT * FROM users")
    records = cursor.fetchall()
    for record in records:
        print record


def test_insert(cursor, user):
    cursor.execute("""INSERT INTO users(id, fname, lname, location, skills, lang) VALUES
        ('%(id)s', '%(fname)s', '%(lname)s', '%(location)s', '%(skills)s', '%(lang)s')""" % user)
    cursor.connection.commit()
    print "INSERT done."


def getCursor():
    conn_string = "host='localhost' dbname='noi' user='postgres' password='ottawa6491' port='5433'"
    print "Connecting to database\n ->%s" % (conn_string)
    conn = psycopg2.connect(conn_string)
    print "Connected!\n"
    cursor = conn.cursor()
    return cursor


def findExperts(cursor, location, langs, skills):
    SQL = """SELECT * FROM
        (
        SELECT fname, lname, lang, plv8_score(skills, %s) AS score, id, location
        FROM users
        WHERE ( (location=%s) OR (%s='') ) AND
        ( (lang::jsonb ?| %s) OR (%s='{}') )
        ) AS T1 WHERE score >= 0
    ORDER BY score DESC LIMIT 10"""
    data = (skills, location, location, langs, langs)
    print cursor.mogrify(SQL, data)
    cursor.execute(SQL, data)
    records = cursor.fetchall()
    return records


def getUser(cursor, userid):
    SQL = """SELECT row_to_json(users) FROM users WHERE id = %s"""
    data = (userid,)
    print cursor.mogrify(SQL, data)
    cursor.execute(SQL, data)
    record = cursor.fetchone()
    if record:
        return record[0]
    else:
        record


def findExpertsAsJSON(cursor, location, langs, skills):
    SQL = """SELECT row_to_json(T1) FROM
        (
        SELECT *, plv8_score(skills, %s) AS score
        FROM users
        WHERE ( (location=%s) OR (%s='') ) AND
        ( (lang::jsonb ?| %s) OR (%s='{}') )
        ) AS T1 WHERE score >= 0
    ORDER BY score DESC LIMIT 10"""
    data = (skills, location, location, langs, langs)
    print cursor.mogrify(SQL, data)
    cursor.execute(SQL, data)
    records = cursor.fetchall()
    return map(lambda x:x[0], records)

def main():
    cursor = getCursor()

    #test_query(cursor)
    user = { 'id': 'google:sahuguet',
        'fname': 'Arnaud',
        'lname': 'Sahuguet',
        'location': 'US',
        'skills': '{  "opendata": {"scraping":3, "machine-readable":3, "standards": 3, "publishing-infrastructure": 3, "pii": 2, "quality": 0}, "prizes": { "legal": 0, "incentives": 3 } }',
        'lang': '["fr", "en", "es"]'
           }
    #test_insert(cursor, user)
    records = findExperts(cursor, 'United States', ['fr', 'ar'], ['scraping', 'standards'])
    for record in records:
        print record

if __name__ == "__main__":
    main()