/sqlautocode/loader.py

 1import sqlalchemy
 2from sqlalchemy.databases import postgres
 3
 4
 5class AutoLoader(object):
 6    pass
 7
 8class PGIndexLoader(AutoLoader):
 9    """ SA does not load indexes for us """
10
11    sql4indexes = "SELECT indexname, tablename, indexdef FROM pg_indexes"
12
13    def __init__(self, db):
14        ix = {}
15        for name, tbl_name, sqltext in db.execute(self.sql4indexes):
16            ix.setdefault(tbl_name, []).append( (name, sqltext) )
17        self._indexes = ix
18
19    def indexes(self, table):
20        return [self._index_from_def(name, sqltext, table)
21                for name, sqltext in self._indexes.get(table.name, ())]
22
23    def _index_from_def(self, name, sqltext, table):
24        # CREATE UNIQUE INDEX name ON "tablename" USING btree (columnslist)
25        unique = ' UNIQUE ' in sqltext
26        cols = sqltext.split(' (')[1].split(')')[0].split(',')
27        cols = [table.columns[cname.strip().replace('"', '')]
28                for cname in cols]
29        name = name.encode('utf-8')
30        return sqlalchemy.Index(name, unique=unique, *cols)
31
32postgres.PGDialect.indexloader = PGIndexLoader