Changeset 725 for trunk/pg.py


Ignore:
Timestamp:
Jan 12, 2016, 10:32:32 AM (4 years ago)
Author:
cito
Message:

Improve implementation and test for pkey()

File:
1 edited

Legend:

Unmodified
Added
Removed
  • trunk/pg.py

    r723 r725  
    3838from decimal import Decimal
    3939from collections import namedtuple
     40from itertools import groupby
    4041
    4142try:
     
    566567        If newpkey is set and is not a dictionary then set that
    567568        value as the primary key of the class.  If it is a dictionary
    568         then replace the _pkeys dictionary with a copy of it.
    569 
    570         """
     569        then replace the internal cache for primary keys with a copy of it.
     570
     571        """
     572        add_schema = self._add_schema
     573
    571574        # First see if the caller is supplying a dictionary
    572575        if isinstance(newpkey, dict):
    573576            # make sure that all classes have a namespace
    574             self._pkeys = dict([
    575                 (cl if '.' in cl else 'public.' + cl, pkey)
    576                 for cl, pkey in newpkey.items()])
     577            self._pkeys = dict((add_schema(cl), pkey)
     578                for cl, pkey in newpkey.items())
    577579            return self._pkeys
    578580
    579         qcl = self._add_schema(cl)  # build fully qualified class name
     581        qcl = add_schema(cl)  # build fully qualified class name
    580582        # Check if the caller is supplying a new primary key for the class
    581583        if newpkey:
     
    586588        if qcl not in self._pkeys:
    587589            # if not found, check again in case it was added after we started
    588             self._pkeys = {}
    589590            q = ("SELECT s.nspname, r.relname, a.attname"
    590591                " FROM pg_class r"
     
    595596                " AND NOT a.attisdropped"
    596597                " JOIN pg_index i ON i.indrelid = r.oid"
    597                 " AND i.indisprimary AND a.attnum = ANY (i.indkey)")
    598             for r in self.db.query(q).getresult():
    599                 cl, pkey = _join_parts(r[:2]), r[2]
    600                 self._pkeys.setdefault(cl, []).append(pkey)
    601             # (only) for composite primary keys, the values will be frozensets
    602             for cl, pkey in self._pkeys.items():
    603                 self._pkeys[cl] = frozenset(pkey) if len(pkey) > 1 else pkey[0]
    604             self._do_debug(self._pkeys)
     598                " AND i.indisprimary AND a.attnum = ANY (i.indkey)"
     599                " ORDER BY 1,2")
     600            rows = self.db.query(q).getresult()
     601            pkeys = {}
     602            for cl, group in groupby(rows, lambda row: row[:2]):
     603                cl = _join_parts(cl)
     604                pkey = [row[2] for row in group]
     605                pkeys[cl] = frozenset(pkey) if len(pkey) > 1 else pkey[0]
     606            self._do_debug(pkeys)
     607            self._pkeys = pkeys
    605608
    606609        # will raise an exception if primary key doesn't exist
Note: See TracChangeset for help on using the changeset viewer.