Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Newer
Older
100755 353 lines (291 sloc) 10.471 kb
6cf9de8 @coleifer oh, derp
authored
1 #!/usr/bin/env python
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
2 # .----.
3 # ===(_)== THIS WONT HURT A BIT...
4 # // 6 6 \\ /
5 # ( 7 )
6 # \ '--' /
7 # \_ ._/
8 # __) (__
9 # /"`/`\`V/`\`\
10 # / \ `Y _/_ \
11 # / [DR]\_ |/ / /\
12 # | ( \/ / / /
13 # \ \ \ /
14 # \ `-/` _.`
15 # `=. `=./
16 # `"`
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
17 from optparse import OptionParser
18 import re
19 import sys
20
5efdc6e @coleifer Experimenting with mysql introspection
authored
21 try:
22 from MySQLdb.constants import FIELD_TYPE
23 MYSQL_MAP = {
24 FIELD_TYPE.BLOB: 'TextField',
25 FIELD_TYPE.CHAR: 'CharField',
26 FIELD_TYPE.DECIMAL: 'DecimalField',
27 FIELD_TYPE.NEWDECIMAL: 'DecimalField',
f9736ff @coleifer Adding introspection bits for date and time fields (thnx django)
authored
28 FIELD_TYPE.DATE: 'DateField',
5efdc6e @coleifer Experimenting with mysql introspection
authored
29 FIELD_TYPE.DATETIME: 'DateTimeField',
30 FIELD_TYPE.DOUBLE: 'FloatField',
31 FIELD_TYPE.FLOAT: 'FloatField',
32 FIELD_TYPE.INT24: 'IntegerField',
33 FIELD_TYPE.LONG: 'IntegerField',
34 FIELD_TYPE.LONGLONG: 'BigIntegerField',
35 FIELD_TYPE.SHORT: 'IntegerField',
36 FIELD_TYPE.STRING: 'CharField',
f9736ff @coleifer Adding introspection bits for date and time fields (thnx django)
authored
37 FIELD_TYPE.TIME: 'TimeField',
5efdc6e @coleifer Experimenting with mysql introspection
authored
38 FIELD_TYPE.TIMESTAMP: 'DateTimeField',
39 FIELD_TYPE.TINY: 'IntegerField',
40 FIELD_TYPE.TINY_BLOB: 'TextField',
41 FIELD_TYPE.MEDIUM_BLOB: 'TextField',
42 FIELD_TYPE.LONG_BLOB: 'TextField',
43 FIELD_TYPE.VAR_STRING: 'CharField',
44 }
056e5a1 @coleifer mysql introspection hand checked and working
authored
45 except ImportError:
5efdc6e @coleifer Experimenting with mysql introspection
authored
46 MYSQL_MAP = {}
47
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
48 from peewee import *
49
50
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
51 class DB(object):
52 conn = None
53
54 def get_conn_class(self):
55 raise NotImplementedError
56
57 def get_columns(self, table):
58 """
59 get_columns('some_table')
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
60
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
61 {
62 'name': 'CharField',
63 'age': 'IntegerField',
64 }
65 """
66 raise NotImplementedError
67
68 def get_foreign_keys(self, table):
69 """
70 get_foreign_keys('some_table')
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
71
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
72 [
73 # column, rel table, rel pk
74 ('blog_id', 'blog', 'id'),
75 ('user_id', 'users', 'id'),
76 ]
77 """
78 raise NotImplementedError
79
80 def get_tables(self):
81 return self.conn.get_tables()
82
83 def connect(self, database, **connect):
84 conn_class = self.get_conn_class()
85 self.conn = conn_class(database, **connect)
86 try:
87 self.conn.connect()
88 except:
89 err('error connecting to %s' % database)
90 raise
91
92
93 class PgDB(DB):
94 # thanks, django
95 reverse_mapping = {
96 16: 'BooleanField',
97 20: 'IntegerField',
98 21: 'IntegerField',
99 23: 'IntegerField',
100 25: 'TextField',
101 700: 'FloatField',
102 701: 'FloatField',
103 1043: 'CharField',
f9736ff @coleifer Adding introspection bits for date and time fields (thnx django)
authored
104 1082: 'DateField',
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
105 1114: 'DateTimeField',
106 1184: 'DateTimeField',
f9736ff @coleifer Adding introspection bits for date and time fields (thnx django)
authored
107 1083: 'TimeField',
108 1266: 'TimeField',
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
109 1700: 'DecimalField',
110 }
111
112 def get_conn_class(self):
113 return PostgresqlDatabase
114
115 def get_columns(self, table):
116 curs = self.conn.execute('select * from %s limit 1' % table)
117 return dict((c.name, self.reverse_mapping.get(c.type_code, 'UnknownFieldType')) for c in curs.description)
118
119 def get_foreign_keys(self, table):
120 framing = '''
121 SELECT
122 kcu.column_name, ccu.table_name, ccu.column_name
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
123 FROM information_schema.table_constraints AS tc
124 JOIN information_schema.key_column_usage AS kcu
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
125 ON tc.constraint_name = kcu.constraint_name
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
126 JOIN information_schema.constraint_column_usage AS ccu
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
127 ON ccu.constraint_name = tc.constraint_name
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
128 WHERE
129 tc.constraint_type = 'FOREIGN KEY' AND
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
130 tc.table_name = %s
131 '''
132 fks = []
133 for row in self.conn.execute(framing, (table,)):
134 fks.append(row)
135 return fks
136
137
5efdc6e @coleifer Experimenting with mysql introspection
authored
138 class MySQLDB(DB):
139 # thanks, django
140 reverse_mapping = MYSQL_MAP
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
141
5efdc6e @coleifer Experimenting with mysql introspection
authored
142 def get_conn_class(self):
143 return MySQLDatabase
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
144
5efdc6e @coleifer Experimenting with mysql introspection
authored
145 def get_columns(self, table):
146 curs = self.conn.execute('select * from %s limit 1' % table)
147 return dict((r[0], self.reverse_mapping.get(r[1], 'UnknownFieldType')) for r in curs.description)
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
148
5efdc6e @coleifer Experimenting with mysql introspection
authored
149 def get_foreign_keys(self, table):
150 framing = '''
151 SELECT column_name, referenced_table_name, referenced_column_name
152 FROM information_schema.key_column_usage
153 WHERE table_name = %s
154 AND table_schema = DATABASE()
155 AND referenced_table_name IS NOT NULL
156 AND referenced_column_name IS NOT NULL
157 '''
158 return [row for row in self.conn.execute(framing, (table,))]
159
160
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
161 class SqDB(DB):
162 # thanks, django
163 reverse_mapping = {
164 'bool': 'BooleanField',
165 'boolean': 'BooleanField',
166 'smallint': 'IntegerField',
167 'smallint unsigned': 'IntegerField',
168 'smallinteger': 'IntegerField',
169 'int': 'IntegerField',
170 'integer': 'IntegerField',
171 'bigint': 'BigIntegerField',
172 'integer unsigned': 'IntegerField',
173 'decimal': 'DecimalField',
174 'real': 'FloatField',
175 'text': 'TextField',
176 'char': 'CharField',
f9736ff @coleifer Adding introspection bits for date and time fields (thnx django)
authored
177 'date': 'DateField',
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
178 'datetime': 'DateTimeField',
f9736ff @coleifer Adding introspection bits for date and time fields (thnx django)
authored
179 'time': 'TimeField',
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
180 }
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
181
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
182 def get_conn_class(self):
183 return SqliteDatabase
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
184
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
185 def map_col(self, col):
186 col = col.lower()
187 if col in self.reverse_mapping:
188 return self.reverse_mapping[col]
189 elif re.search(r'^\s*(?:var)?char\s*\(\s*(\d+)\s*\)\s*$', col):
190 return 'CharField'
191 else:
192 return 'UnknownFieldType'
193
194 def get_columns(self, table):
195 curs = self.conn.execute('pragma table_info(%s)' % table)
196 col_dict = {}
197 for (_, name, col, not_null, _, is_pk) in curs.fetchall():
198 # cid, name, type, notnull, dflt_value, pk
199 if is_pk:
200 col_type = 'PrimaryKeyField'
201 else:
202 col_type = self.map_col(col)
203 col_dict[name] = col_type
204 return col_dict
205
206 def get_foreign_keys(self, table):
207 fks = []
208
209 curs = self.conn.execute("SELECT sql FROM sqlite_master WHERE tbl_name = ? AND type = ?", [table, "table"])
210 table_def = curs.fetchone()[0].strip()
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
211
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
212 try:
213 columns = re.search('\((.+)\)', table_def).groups()[0]
214 except AttributeError:
215 err('Unable to read table definition for "%s"' % table)
216 sys.exit(1)
217
218 for col_def in columns.split(','):
219 col_def = col_def.strip()
220 m = re.search('"?(.+?)"?\s+.+\s+references (.*) \(["|]?(.*)["|]?\)', col_def, re.I)
221 if not m:
222 continue
223
224 fk_column, rel_table, rel_pk = [s.strip('"') for s in m.groups()]
225 fks.append((fk_column, rel_table, rel_pk))
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
226
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
227 return fks
228
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
229
230 frame = '''from peewee import *
231
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
232 database = %s('%s', **%s)
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
233
234 class UnknownFieldType(object):
235 pass
236
237 class BaseModel(Model):
238 class Meta:
239 database = database
240 '''
241
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
242 engine_mapping = {
243 'postgresql': PgDB,
244 'sqlite': SqDB,
056e5a1 @coleifer mysql introspection hand checked and working
authored
245 'mysql': MySQLDB,
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
246 }
247
248 def get_db(engine):
249 if engine not in engine_mapping:
250 err('Unsupported engine: "%s"' % engine)
251 sys.exit(1)
252
253 db_class = engine_mapping[engine]
254 return db_class()
255
256 def introspect(engine, database, **connect):
257 db = get_db(engine)
258 db.connect(database, **connect)
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
259
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
260 tables = db.get_tables()
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
261
262 models = {}
263 table_to_model = {}
efb13e3 @coleifer Speed up foreign key lookups by caching them
authored
264 table_fks = {}
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
265
266 # first pass, just raw column names and peewee type
267 for table in tables:
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
268 models[table] = db.get_columns(table)
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
269 table_to_model[table] = tn(table)
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
270 table_fks[table] = db.get_foreign_keys(table)
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
271
272 # second pass, convert foreign keys, assign primary keys, and mark
273 # explicit column names where they don't match the "pythonic" ones
274 col_meta = {}
275 for table in tables:
276 col_meta[table] = {}
efb13e3 @coleifer Speed up foreign key lookups by caching them
authored
277 for column, rel_table, rel_pk in table_fks[table]:
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
278 models[table][column] = 'ForeignKeyField'
279 models[rel_table][rel_pk] = 'PrimaryKeyField'
280 col_meta[table][column] = {'to': table_to_model[rel_table]}
281
282 for column in models[table]:
283 col_meta[table].setdefault(column, {})
284 if column != cn(column):
285 col_meta[table][column]['db_column'] = "'%s'" % column
286
287 # write generated code to standard out
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
288 print frame % (db.get_conn_class().__name__, database, repr(connect))
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
289
290 # print the models
291 def print_model(model, seen):
efb13e3 @coleifer Speed up foreign key lookups by caching them
authored
292 for _, rel_table, _ in table_fks[model]:
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
293 if rel_table not in seen:
294 seen.add(rel_table)
295 print_model(rel_table, seen)
296
297 ttm = table_to_model[model]
298 print 'class %s(BaseModel):' % ttm
299 cols = models[model]
300 for column, field_class in ds(cols):
301 if column == 'id' and field_class in ('IntegerField', 'PrimaryKeyField'):
302 continue
303
304 field_params = ', '.join([
305 '%s=%s' % (k, v) for k, v in col_meta[model][column].items()
306 ])
307 print ' %s = %s(%s)' % (cn(column), field_class, field_params)
308 print
e469be1 @medwards Whitespace, hee-YAH!
medwards authored
309
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
310 print ' class Meta:'
311 print ' db_table = \'%s\'' % model
312 print
313 seen.add(model)
314
315 seen = set()
316 for model, cols in ds(models):
317 if model not in seen:
318 print_model(model, seen)
319
320 # misc
321 tn = lambda t: t.title().replace('_', '')
322 cn = lambda c: re.sub('_id$', '', c.lower())
323 ds = lambda d: sorted(d.items(), key=lambda t:t[0])
324
325 def err(msg):
326 print '\033[91m%s\033[0m' % msg
327
328
329 if __name__ == '__main__':
330 parser = OptionParser(usage='usage: %prog [options] database_name')
331 ao = parser.add_option
332 ao('-H', '--host', dest='host')
333 ao('-p', '--port', dest='port', type='int')
334 ao('-u', '--user', dest='user')
335 ao('-P', '--password', dest='password')
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
336 ao('-e', '--engine', dest='engine', default='postgresql')
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
337
338 options, args = parser.parse_args()
339 ops = ('host', 'port', 'user', 'password')
340 connect = dict((o, getattr(options, o)) for o in ops if getattr(options, o))
341
342 if len(args) < 1:
343 print 'error: missing required parameter "database"'
344 parser.print_help()
345 sys.exit(1)
346
347 database = args[-1]
6fb13f1 @coleifer Use the correct password argument in pwiz, fixes #74
authored
348
349 if options.engine == 'mysql' and 'password' in connect:
350 connect['passwd'] = connect.pop('password', None)
fd281cf @coleifer Adding a little helper to introspect postgresql databases and generate
authored
351
75c343d @coleifer Refactoring pwiz to work with both psql and sqlite...mysql coming soon
authored
352 introspect(options.engine, database, **connect)
Something went wrong with that request. Please try again.