def balance(campaign, currency, get_or_create=False, for_update=False): """Return a campaigns balance for ``currency``, if ``get_or_create=True``, if an existing balance isn't found, one will be created and returned. :param campaign: Campaign for which to retrieve balance. :type campaign: :class:`pooldlib.postgresql.models.Campaign` :param currency: Currenty type of the balance to retrieve :type currency::class: `pooldlib.postgresql.models.Currency` :param get_or_create: If `True` and an existing balance is not found one will be created for ``currency``. :type get_or_create: boolean :param for_update: If `True`, the ``SELECT FOR UPDATE`` directive will be used when retrieving the target balance. :type for_update: boolean :returns: :class:`pooldlib.postgresql.models.Balance` """ b = _balance.get(for_update=for_update, currency_id=currency.id, type='campaign', campaign_id=campaign.id) if not b and get_or_create: b = _balance.create_for_campaign(campaign, currency) if isinstance(b, (tuple, list)): b = b[0] return b or None
def balance_for_currency(self, currency, get_or_create=True, for_update=False): """Return a user's :class:`pooldlib.postgresql.models.Balance` for a given :class:`pooldlib.postgresql.models.Currency` or currency code (:class:`str`). :param currency: Currency of balance to return. :type currency: :class:`pooldlib.postgresql.models.Balance` or string. :param get_or_create: Should a balance be created for the user in the target currency if none is found? :type get_or_create: boolean, default `True` :param for_update: If `True` the `FOR UPDATE` directive will be used, locking the row for an `UPDATE` query. :type for_update: boolean, default `False` """ from pooldlib.api import balance from pooldlib.postgresql import (db, Currency as CurrencyModel, Balance as BalanceModel) if isinstance(currency, basestring): currency = CurrencyModel.get(currency) if self.__class__.__name__ == 'User': balance = balance.get(for_update=for_update, user_id=self.id, currency_id=currency.id) else: balance = balance.get(for_update=for_update, campaign_id=self.id, currency_id=currency.id) # If we don't find a balance for the user, create if requested to if not balance and get_or_create: balance = BalanceModel() balance.enabled = True balance.amount = Decimal('0.0000') balance.currency = currency balance.type = self.__class__.__name__.lower() if self.__class__.__name__ == 'User': balance.user = self else: balance.campaign = self db.session.add(balance) self.balances.append(balance) db.session.flush() # Balance.get returns a list if for_update=False. A user should only # have a single balance for a given currency, so return the # CurrencyModel object instead of a list. if isinstance(balance, list): balance = balance[0] return balance