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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
|
import os
from sqlplain.util import openclose, getoutput
GET_PKEYS = '''\
SELECT attname FROM pg_attribute
WHERE attrelid = (
SELECT indexrelid FROM pg_index AS i
WHERE i.indrelid = (SELECT oid FROM pg_class WHERE relname=:table)
AND i.indisprimary = 't')
ORDER BY attnum
'''
def get_kfields_postgres(conn, tname):
return [x.attname for x in conn.execute(GET_PKEYS, (tname,))]
def create_db_postgres(uri):
openclose(uri.copy(database='template1'),
'CREATE DATABASE %(database)s' % uri)
def drop_db_postgres(uri):
openclose(uri.copy(database='template1'),
'DROP DATABASE %(database)s' % uri)
def get_sizeK_postgres(conn, table):
return conn.execute('SELECT relpages*8 FROM pg_class WHERE relname=:table',
(table,), scalar=True)
def get_tables_postgres(conn, schema):
query = 'SELECT tablename FROM pg_tables'
if schema:
res = conn.execute(query + ' WHERE schemaname=:a', (schema,))
else:
res = conn.execute(query)
return [r[0] for r in res]
def exists_table_postgres(conn, tname):
return conn.execute('SELECT count(*) FROM pg_tables WHERE tablename=:table',
(tname,), scalar=True)
def exists_db_postgres(uri):
dbname = uri['database']
for row in openclose(
uri.copy(database='template1'), 'SELECT datname FROM pg_database'):
if row[0] == dbname:
return True
return False
def dump_file_postgres(uri, query, filename, mode, sep='\t', null='',
schema=None): # or null='\N'
"""
Save the result of a query on a local file by using COPY TO and psql
"""
if not ' ' in query: # assumes a table name was given
query = '(select * from %s)' % query
else:
query = '(%s)' % query
if mode == 'b':
return psql(uri, "COPY %s TO STDOUT BINARY" % query, filename, schema)
else:
return psql(
uri, "COPY %s TO STDOUT WITH DELIMITER '%s' NULL '%s'" %
(query, sep, null), filename, schema)
# apparently copy_from from psycopg2 is buggy for large files
def load_file_postgres(uri, tname, filename, mode, sep='\t', null='',
schema=None): # or null='\N'
"""
Load a file into a table by using COPY FROM and psql
"""
stdin = file(filename)
if mode == 'b':
return psql(uri, "COPY %s FROM STDIN BINARY" % tname, stdin=stdin,
schema=schema)
else: # csv mode
copy_from = "COPY %s FROM STDIN WITH DELIMITER E'%s' NULL E'%s'" % (
tname, sep, null)
return psql(uri, copy_from, stdin=stdin, schema=schema)
###############################################################################
## pg_dump and pg_restore should be used for multiple tables or whole databases
def pg_dump(uri, filename, *args):
"""
A small wrapper over pg_dump. Example:
>> pg_dump(uri, thetable, thefile)
"""
cmd = ['pg_dump', '-h', uri['host'], '-U', uri['user'],
'-d', uri['database'], '-f', filename] + list(args)
print ' '.join(cmd)
return getoutput(cmd)
def pg_restore(uri, filename, *args):
"""
A small wrapper over pg_restore. Example:
>> pg_restore(uri, thetable, thefile)
"""
cmd = ['pg_restore', '-h', uri['host'], '-U', uri['user'],
'-d', uri['database'], filename] + list(args)
print ' '.join(cmd)
return getoutput(cmd)
def psql(uri, query, filename=os.devnull, stdin=None, schema=None):
"Execute a query and save its result on filename"
if not ' ' in query: # assumes a table name was given
query = 'select * from %s' % query
if schema:
query = 'SET search_path TO %s; %s' % (schema, query)
cmd = ['psql', '-h', uri['host'], '-U', uri['user'], '-d', uri['database'],
'-c', query, '-o', filename]
# print cmd
return getoutput(cmd, stdin)
|
