forked from coleifer/peewee
/
pwiz.py
executable file
·120 lines (92 loc) · 3.66 KB
/
pwiz.py
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
114
115
116
117
118
119
120
#!/usr/bin/env python
from optparse import OptionParser
import sys
from peewee import *
from peewee import print_
from playhouse.reflection import *
TEMPLATE = """from peewee import *
database = %s('%s', **%s)
class UnknownField(object):
pass
class BaseModel(Model):
class Meta:
database = database
"""
DATABASE_ALIASES = {
MySQLDatabase: ['mysql', 'mysqldb'],
PostgresqlDatabase: ['postgres', 'postgresql'],
SqliteDatabase: ['sqlite', 'sqlite3'],
}
DATABASE_MAP = dict((value, key)
for key in DATABASE_ALIASES
for value in DATABASE_ALIASES[key])
def make_introspector(database_type, database_name, **kwargs):
if database_type not in DATABASE_MAP:
err('Unrecognized database, must be one of: %s' %
', '.join(DATABASE_MAP.keys()))
sys.exit(1)
schema = kwargs.pop('schema', None)
DatabaseClass = DATABASE_MAP[database_type]
db = DatabaseClass(database_name, **kwargs)
return Introspector.from_database(db, schema=schema)
def print_models(introspector, tables=None):
columns, foreign_keys, model_names = introspector.introspect()
print_(TEMPLATE % (
introspector.get_database_class().__name__,
introspector.get_database_name(),
repr(introspector.get_database_kwargs())))
def _print_table(table, seen, accum=None):
accum = accum or []
for foreign_key in foreign_keys[table]:
dest = foreign_key.dest_table
# In the event the destination table has already been pushed
# for printing, then we have a reference cycle.
if dest in accum and table not in accum:
print_('# Possible reference cycle: %s' % foreign_key)
# If this is not a self-referential foreign key, and we have
# not already processed the destination table, do so now.
if dest not in seen and dest not in accum:
seen.add(dest)
if dest != table:
_print_table(dest, seen, accum + [table])
print_('class %s(BaseModel):' % model_names[table])
for name, column in sorted(columns[table].items()):
if name == 'id' and column.field_class in introspector.pk_classes:
continue
print_(' %s' % column.get_field())
print_('')
print_(' class Meta:')
print_(' db_table = \'%s\'' % table)
print_('')
seen.add(table)
seen = set()
for table in sorted(model_names.keys()):
if table not in seen:
if not tables or table in tables:
_print_table(table, seen)
def err(msg):
sys.stderr.write('\033[91m%s\033[0m\n' % msg)
sys.stderr.flush()
if __name__ == '__main__':
parser = OptionParser(usage='usage: %prog [options] database_name')
ao = parser.add_option
ao('-H', '--host', dest='host')
ao('-p', '--port', dest='port', type='int')
ao('-u', '--user', dest='user')
ao('-P', '--password', dest='password')
ao('-e', '--engine', dest='engine', default='postgresql')
ao('-s', '--schema', dest='schema')
ao('-t', '--tables', dest='tables')
options, args = parser.parse_args()
ops = ('host', 'port', 'user', 'password', 'schema')
connect = dict((o, getattr(options, o)) for o in ops if getattr(options, o))
if len(args) < 1:
err('Missing required parameter "database"')
parser.print_help()
sys.exit(1)
database = args[-1]
tables = None
if options.tables:
tables = [x for x in options.tables.split(',') if x]
introspector = make_introspector(options.engine, database, **connect)
print_models(introspector, tables)