Changeset 765 for trunk/pg.py


Ignore:
Timestamp:
Jan 18, 2016, 6:21:44 PM (4 years ago)
Author:
cito
Message:

Improve support for access by primary key

Composite primary keys are now returned as tuples instead of frozensets,
where the ordering of the tuple reflects the primary key index.

Primary keys now takes precedence if both OID and primary key are available
(this was solved the other way around in 4.x). Use of OIDs is thus slightly
more discouraged, though it still works as before for tables with OIDs where
no primary key is available.

This changeset also clarifies some docstrings, makes the code a bit clearer,
handles and tests some more edge cases (pg module still has 100% coverage).

File:
1 edited

Legend:

Unmodified
Added
Removed
  • trunk/pg.py

    r763 r765  
    355355            raise ValueError
    356356        return d
     357
     358    _num_types = frozenset('int float num money'
     359        ' int2 int4 int8 float4 float8 numeric money'.split())
    357360
    358361    def _prepare_num(self, d):
     
    608611        """Execute a SQL command string.
    609612
    610         This method simply sends a SQL query to the database. If the query is
     613        This method simply sends a SQL query to the database.  If the query is
    611614        an insert statement that inserted exactly one row into a table that
    612615        has OIDs, the return value is the OID of the newly inserted row.
    613616        If the query is an update or delete statement, or an insert statement
    614617        that did not insert exactly one row in a table with OIDs, then the
    615         number of rows affected is returned as a string. If it is a statement
     618        number of rows affected is returned as a string.  If it is a statement
    616619        that returns rows as a result (usually a select statement, but maybe
    617620        also an "insert/update ... returning" statement), this method returns
    618621        a Query object that can be accessed via getresult() or dictresult()
    619         or simply printed. Otherwise, it returns `None`.
     622        or simply printed.  Otherwise, it returns `None`.
    620623
    621624        The query can contain numbered parameters of the form $1 in place
    622         of any data constant. Arguments given after the query string will
    623         be substituted for the corresponding numbered parameter. Parameter
     625        of any data constant.  Arguments given after the query string will
     626        be substituted for the corresponding numbered parameter.  Parameter
    624627        values can also be given as a single list or tuple argument.
    625628        """
     
    630633        return self.db.query(qstr, args)
    631634
    632     def pkey(self, table, flush=False):
     635    def pkey(self, table, composite=False, flush=False):
    633636        """Get or set the primary key of a table.
    634637
    635         Composite primary keys are represented as frozensets. Note that
    636         this raises a KeyError if the table does not have a primary key.
     638        Single primary keys are returned as strings unless you
     639        set the composite flag.  Composite primary keys are always
     640        represented as tuples.  Note that this raises a KeyError
     641        if the table does not have a primary key.
    637642
    638643        If flush is set then the internal cache for primary keys will
    639         be flushed. This may be necessary after the database schema or
     644        be flushed.  This may be necessary after the database schema or
    640645        the search path has been changed.
    641646        """
     
    647652            pkey = pkeys[table]
    648653        except KeyError:  # cache miss, check the database
    649             q = ("SELECT a.attname FROM pg_index i"
     654            q = ("SELECT a.attname, a.attnum, i.indkey FROM pg_index i"
    650655                " JOIN pg_attribute a ON a.attrelid = i.indrelid"
    651656                " AND a.attnum = ANY(i.indkey)"
    652657                " AND NOT a.attisdropped"
    653658                " WHERE i.indrelid=%s::regclass"
    654                 " AND i.indisprimary") % (
     659                " AND i.indisprimary ORDER BY a.attnum") % (
    655660                    self._prepare_qualified_param(table, 1),)
    656661            pkey = self.db.query(q, (table,)).getresult()
    657662            if not pkey:
    658663                raise KeyError('Table %s has no primary key' % table)
     664            # we want to use the order defined in the primary key index here,
     665            # not the order as defined by the columns in the table
    659666            if len(pkey) > 1:
    660                 pkey = frozenset(k[0] for k in pkey)
     667                indkey = [int(k) for k in pkey[0][2].split()]
     668                pkey = sorted(pkey, key=lambda row: indkey.index(row[1]))
     669                pkey = tuple(row[0] for row in pkey)
    661670            else:
    662671                pkey = pkey[0][0]
    663672            pkeys[table] = pkey  # cache it
     673        if composite and not isinstance(pkey, tuple):
     674            pkey = (pkey,)
    664675        return pkey
    665676
     
    755766        """Get a row from a database table or view.
    756767
    757         This method is the basic mechanism to get a single row.  The keyname
    758         that the key specifies a unique row.  If keyname is not specified
    759         then the primary key for the table is used.  If row is a dictionary
    760         then the value for the key is taken from it and it is modified to
    761         include the new values, replacing existing values where necessary.
    762         For a composite key, keyname can also be a sequence of key names.
     768        This method is the basic mechanism to get a single row.  It assumes
     769        that the keyname specifies a unique row.  It must be the name of a
     770        single column or a tuple of column names.  If the keyname is not
     771        specified, then the primary key for the table is used.
     772
     773        If row is a dictionary, then the value for the key is taken from it.
     774        Otherwise, the row must be a single value or a tuple of values
     775        corresponding to the passed keyname or primary key.  The fetched row
     776        from the table will be returned as a new dictionary or used to replace
     777        the existing values when row was passed as aa dictionary.
     778
    763779        The OID is also put into the dictionary if the table has one, but
    764780        in order to allow the caller to work with multiple tables, it is
    765         munged as "oid(table)".
    766         """
    767         if table.endswith('*'):  # scan descendant tables?
    768             table = table[:-1].rstrip()  # need parent table name
     781        munged as "oid(table)" using the actual name of the table.
     782        """
     783        if table.endswith('*'):  # hint for descendant tables can be ignored
     784            table = table[:-1].rstrip()
     785        attnames = self.get_attnames(table)
     786        qoid = _oid_key(table) if 'oid' in attnames else None
     787        if keyname and isinstance(keyname, basestring):
     788            keyname = (keyname,)
     789        if qoid and isinstance(row, dict) and qoid in row and 'oid' not in row:
     790            row['oid'] = row[qoid]
    769791        if not keyname:
    770             # use the primary key by default
    771             try:
    772                 keyname = self.pkey(table)
    773             except KeyError:
    774                 raise _prg_error('Table %s has no primary key' % table)
    775         attnames = self.get_attnames(table)
     792            try:  # if keyname is not specified, try using the primary key
     793                keyname = self.pkey(table, True)
     794            except KeyError:  # the table has no primary key
     795                # try using the oid instead
     796                if qoid and isinstance(row, dict) and 'oid' in row:
     797                    keyname = ('oid',)
     798                else:
     799                    raise _prg_error('Table %s has no primary key' % table)
     800            else:  # the table has a primary key
     801                # check whether all key columns have values
     802                if isinstance(row, dict) and not set(keyname).issubset(row):
     803                    # try using the oid instead
     804                    if qoid and 'oid' in row:
     805                        keyname = ('oid',)
     806                    else:
     807                        raise KeyError(
     808                            'Missing value in row for specified keyname')
     809        if not isinstance(row, dict):
     810            if not isinstance(row, (tuple, list)):
     811                row = [row]
     812            if len(keyname) != len(row):
     813                raise KeyError(
     814                    'Differing number of items in keyname and row')
     815            row = dict(zip(keyname, row))
    776816        params = []
    777817        param = partial(self._prepare_param, params=params)
    778818        col = self.escape_identifier
    779         # We want the oid for later updates if that isn't the key.
    780         # To allow users to work with multiple tables, we munge
    781         # the name of the "oid" key by adding the name of the table.
    782         qoid = _oid_key(table)
    783         if keyname == 'oid':
    784             if isinstance(row, dict):
    785                 if qoid not in row:
    786                     raise _prg_error('%s not in row' % qoid)
    787             else:
    788                 row = {qoid: row}
    789             what = '*'
    790             where = 'oid = %s' % param(row[qoid], 'int')
    791         else:
    792             keyname = [keyname] if isinstance(
    793                 keyname, basestring) else sorted(keyname)
    794             if not isinstance(row, dict):
    795                 if len(keyname) > 1:
    796                     raise _prg_error('Composite key needs dict as row')
    797                 row = dict((k, row) for k in keyname)
    798             what = ', '.join(col(k) for k in attnames)
    799             where = ' AND '.join('%s = %s' % (
    800                 col(k), param(row[k], attnames[k])) for k in keyname)
     819        what = 'oid, *' if qoid else '*'
     820        where = ' AND '.join('%s = %s' % (
     821            col(k), param(row[k], attnames[k])) for k in keyname)
     822        if 'oid' in row:
     823            if qoid:
     824                row[qoid] = row['oid']
     825            del row['oid']
    801826        q = 'SELECT %s FROM %s WHERE %s LIMIT 1' % (
    802827            what, self._escape_qualified_name(table), where)
     
    808833                table, where, self._list_params(params)))
    809834        for n, value in res[0].items():
    810             if n == 'oid':
     835            if qoid and n == 'oid':
    811836                n = qoid
    812             elif attnames.get(n) == 'bytea':
     837            elif value is not None and attnames.get(n) == 'bytea':
    813838                value = self.unescape_bytea(value)
    814839            row[n] = value
     
    831856        although PostgreSQL does.
    832857        """
    833         if 'oid' in kw:
    834             del kw['oid']
     858        if table.endswith('*'):  # hint for descendant tables can be ignored
     859            table = table[:-1].rstrip()
    835860        if row is None:
    836861            row = {}
    837862        row.update(kw)
     863        if 'oid' in row:
     864            del row['oid']  # do not insert oid
    838865        attnames = self.get_attnames(table)
     866        qoid = _oid_key(table) if 'oid' in attnames else None
    839867        params = []
    840868        param = partial(self._prepare_param, params=params)
     
    846874                values.append(param(row[n], attnames[n]))
    847875        names, values = ', '.join(names), ', '.join(values)
    848         ret = 'oid, *' if 'oid' in attnames else '*'
     876        ret = 'oid, *' if qoid else '*'
    849877        q = 'INSERT INTO %s (%s) VALUES (%s) RETURNING %s' % (
    850878            self._escape_qualified_name(table), names, values, ret)
    851879        self._do_debug(q, params)
    852880        q = self.db.query(q, params)
    853         res = q.dictresult()  # this will always return a row
    854         for n, value in res[0].items():
    855             if n == 'oid':
    856                 n = _oid_key(table)
    857             elif attnames.get(n) == 'bytea' and value is not None:
    858                 value = self.unescape_bytea(value)
    859             row[n] = value
     881        res = q.dictresult()
     882        if res:  # this should always be true
     883            for n, value in res[0].items():
     884                if qoid and n == 'oid':
     885                    n = qoid
     886                elif value is not None and attnames.get(n) == 'bytea':
     887                    value = self.unescape_bytea(value)
     888                row[n] = value
    860889        return row
    861890
     
    864893
    865894        Similar to insert but updates an existing row.  The update is based
    866         on the OID value as munged by get or passed as keyword, or on the
    867         primary key of the table.  The dictionary is modified to reflect
    868         any changes caused by the update due to triggers, rules, default
    869         values, etc.
    870         """
    871         # Update always works on the oid which get() returns if available,
    872         # otherwise use the primary key.  Fail if neither.
    873         # Note that we only accept oid key from named args for safety.
    874         qoid = _oid_key(table)
    875         if 'oid' in kw:
    876             kw[qoid] = kw.pop('oid')
     895        on the primary key of the table or the OID value as munged by get
     896        or passed as keyword.
     897
     898        The dictionary is then modified to reflect any changes caused by the
     899        update due to triggers, rules, default values, etc.
     900        """
     901        if table.endswith('*'):
     902            table = table[:-1].rstrip()  # need parent table name
     903        attnames = self.get_attnames(table)
     904        qoid = _oid_key(table) if 'oid' in attnames else None
    877905        if row is None:
    878906            row = {}
     907        elif 'oid' in row:
     908            del row['oid']  # only accept oid key from named args for safety
    879909        row.update(kw)
    880         attnames = self.get_attnames(table)
     910        if qoid and qoid in row and 'oid' not in row:
     911            row['oid'] = row[qoid]
     912        try:  # try using the primary key
     913            keyname = self.pkey(table, True)
     914        except KeyError:  # the table has no primary key
     915            # try using the oid instead
     916            if qoid and 'oid' in row:
     917                keyname = ('oid',)
     918            else:
     919                raise _prg_error('Table %s has no primary key' % table)
     920        else:  # the table has a primary key
     921            # check whether all key columns have values
     922            if not set(keyname).issubset(row):
     923                # try using the oid instead
     924                if qoid and 'oid' in row:
     925                    keyname = ('oid',)
     926                else:
     927                    raise KeyError('Missing primary key in row')
    881928        params = []
    882929        param = partial(self._prepare_param, params=params)
    883930        col = self.escape_identifier
    884         if qoid in row:
    885             where = 'oid = %s' % param(row[qoid], 'int')
    886             keyname = []
    887         else:
    888             try:
    889                 keyname = self.pkey(table)
    890             except KeyError:
    891                 raise _prg_error('Table %s has no primary key' % table)
    892             keyname = [keyname] if isinstance(
    893                 keyname, basestring) else sorted(keyname)
    894             try:
    895                 where = ' AND '.join('%s = %s' % (
    896                     col(k), param(row[k], attnames[k])) for k in keyname)
    897             except KeyError:
    898                 raise _prg_error('Update operation needs primary key or oid')
     931        where = ' AND '.join('%s = %s' % (
     932            col(k), param(row[k], attnames[k])) for k in keyname)
     933        if 'oid' in row:
     934            if qoid:
     935                row[qoid] = row['oid']
     936            del row['oid']
     937        values = []
    899938        keyname = set(keyname)
    900         keyname.add('oid')
    901         values = []
    902939        for n in attnames:
    903940            if n in row and n not in keyname:
     
    906943            return row
    907944        values = ', '.join(values)
    908         ret = 'oid, *' if 'oid' in attnames else '*'
     945        ret = 'oid, *' if qoid else '*'
    909946        q = 'UPDATE %s SET %s WHERE %s RETURNING %s' % (
    910947            self._escape_qualified_name(table), values, where, ret)
     
    914951        if res:  # may be empty when row does not exist
    915952            for n, value in res[0].items():
    916                 if n == 'oid':
     953                if qoid and n == 'oid':
    917954                    n = qoid
    918                 elif attnames.get(n) == 'bytea' and value is not None:
     955                elif value is not None and attnames.get(n) == 'bytea':
    919956                    value = self.unescape_bytea(value)
    920957                row[n] = value
     
    9641001        only available since PostgreSQL 9.5.
    9651002        """
    966         if 'oid' in kw:
    967             del kw['oid']
     1003        if table.endswith('*'):  # hint for descendant tables can be ignored
     1004            table = table[:-1].rstrip()
    9681005        if row is None:
    9691006            row = {}
     1007        if 'oid' in row:
     1008            del row['oid']  # do not insert oid
     1009        if 'oid' in kw:
     1010            del kw['oid']  # do not update oid
    9701011        attnames = self.get_attnames(table)
     1012        qoid = _oid_key(table) if 'oid' in attnames else None
    9711013        params = []
    9721014        param = partial(self._prepare_param,params=params)
     
    9791021        names, values = ', '.join(names), ', '.join(values)
    9801022        try:
    981             keyname = self.pkey(table)
     1023            keyname = self.pkey(table, True)
    9821024        except KeyError:
    9831025            raise _prg_error('Table %s has no primary key' % table)
    984         keyname = [keyname] if isinstance(
    985             keyname, basestring) else sorted(keyname)
    9861026        target = ', '.join(col(k) for k in keyname)
    9871027        update = []
     
    9981038            return row
    9991039        do = 'update set %s' % ', '.join(update) if update else 'nothing'
    1000         ret = 'oid, *' if 'oid' in attnames else '*'
     1040        ret = 'oid, *' if qoid else '*'
    10011041        q = ('INSERT INTO %s AS included (%s) VALUES (%s)'
    10021042            ' ON CONFLICT (%s) DO %s RETURNING %s') % (
     
    10121052            raise  # re-raise original error
    10131053        res = q.dictresult()
    1014         if update:  # may be empty with "do nothing"
     1054        if res:  # may be empty with "do nothing"
    10151055            for n, value in res[0].items():
    1016                 if n == 'oid':
    1017                     n = _oid_key(table)
    1018                 elif attnames.get(n) == 'bytea' and value is not None:
     1056                if qoid and n == 'oid':
     1057                    n = qoid
     1058                elif value is not None and attnames.get(n) == 'bytea':
    10191059                    value = self.unescape_bytea(value)
    10201060                row[n] = value
     
    10381078            if n == 'oid':
    10391079                continue
    1040             if t in ('int', 'integer', 'smallint', 'bigint',
    1041                     'float', 'real', 'double precision',
    1042                     'num', 'numeric', 'money'):
     1080            if t in self._num_types:
    10431081                row[n] = 0
    1044             elif t in ('bool', 'boolean'):
     1082            elif t == 'bool':
    10451083                row[n] = self._make_bool(False)
    10461084            else:
     
    10521090
    10531091        This method deletes the row from a table.  It deletes based on the
    1054         OID value as munged by get or passed as keyword, or on the primary
    1055         key of the table.  The return value is the number of deleted rows
    1056         (i.e. 0 if the row did not exist and 1 if the row was deleted).
    1057         """
    1058         # Like update, delete works on the oid.
    1059         # One day we will be testing that the record to be deleted
    1060         # isn't referenced somewhere (or else PostgreSQL will).
    1061         # Note that we only accept oid key from named args for safety.
    1062         qoid = _oid_key(table)
    1063         if 'oid' in kw:
    1064             kw[qoid] = kw.pop('oid')
     1092        primary key of the table or the OID value as munged by get() or
     1093        passed as keyword.
     1094
     1095        The return value is the number of deleted rows (i.e. 0 if the row
     1096        did not exist and 1 if the row was deleted).
     1097
     1098        Note that if the row cannot be deleted because e.g. it is still
     1099        referenced by another table, this method raises a ProgrammingError.
     1100        """
     1101        if table.endswith('*'):  # hint for descendant tables can be ignored
     1102            table = table[:-1].rstrip()
     1103        attnames = self.get_attnames(table)
     1104        qoid = _oid_key(table) if 'oid' in attnames else None
    10651105        if row is None:
    10661106            row = {}
     1107        elif 'oid' in row:
     1108            del row['oid']  # only accept oid key from named args for safety
    10671109        row.update(kw)
     1110        if qoid and qoid in row and 'oid' not in row:
     1111            row['oid'] = row[qoid]
     1112        try:  # try using the primary key
     1113            keyname = self.pkey(table, True)
     1114        except KeyError:  # the table has no primary key
     1115            # try using the oid instead
     1116            if qoid and 'oid' in row:
     1117                keyname = ('oid',)
     1118            else:
     1119                raise _prg_error('Table %s has no primary key' % table)
     1120        else:  # the table has a primary key
     1121            # check whether all key columns have values
     1122            if not set(keyname).issubset(row):
     1123                # try using the oid instead
     1124                if qoid and 'oid' in row:
     1125                    keyname = ('oid',)
     1126                else:
     1127                    raise KeyError('Missing primary key in row')
    10681128        params = []
    10691129        param = partial(self._prepare_param, params=params)
    1070         if qoid in row:
    1071             where = 'oid = %s' % param(row[qoid], 'int')
    1072         else:
    1073             try:
    1074                 keyname = self.pkey(table)
    1075             except KeyError:
    1076                 raise _prg_error('Table %s has no primary key' % table)
    1077             keyname = [keyname] if isinstance(
    1078                 keyname, basestring) else sorted(keyname)
    1079             attnames = self.get_attnames(table)
    1080             col = self.escape_identifier
    1081             try:
    1082                 where = ' AND '.join('%s = %s' % (
    1083                     col(k), param(row[k], attnames[k])) for k in keyname)
    1084             except KeyError:
    1085                 raise _prg_error('Delete operation needs primary key or oid')
     1130        col = self.escape_identifier
     1131        where = ' AND '.join('%s = %s' % (
     1132            col(k), param(row[k], attnames[k])) for k in keyname)
     1133        if 'oid' in row:
     1134            if qoid:
     1135                row[qoid] = row['oid']
     1136            del row['oid']
    10861137        q = 'DELETE FROM %s WHERE %s' % (
    10871138            self._escape_qualified_name(table), where)
Note: See TracChangeset for help on using the changeset viewer.