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