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(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
