Merge pull request #17 from CartoDB/geocoder_check_quota

Geocoder check quota

.gitignore

@@ -0,0 +1 @@
+*.pyc

server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py

@@ -0,0 +1,27 @@
+import redis
+import user_service
+from datetime import date
+
+class QuotaService:
+    """ Class to manage all the quota operation for the Geocoder SQL API Extension """
+
+    def __init__(self, user_id, transaction_id, **kwargs):
+        self._user_service = user_service.UserService(user_id, **kwargs)
+        self.transaction_id = transaction_id
+
+    def check_user_quota(self):
+        """ Check if the current user quota surpasses the current quota """
+        # TODO We need to add the hard/soft limit flag for the geocoder
+        user_quota = self.user_service.user_quota()
+        today = date.today()
+        current_used = self.user_service.used_quota_month(today.year, today.month)
+        soft_geocoder_limit = self.user_service.soft_geocoder_limit()
+        return True if soft_geocoder_limit or (current_used + 1) < user_quota else False
+
+    def increment_geocoder_use(self, amount=1):
+        today = date.today()
+        self.user_service.increment_geocoder_use(today.year, today.month, self.transaction_id)
+
+    @property
+    def user_service(self):
+        return self._user_service

server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py

@@ -0,0 +1,79 @@
+import redis
+from datetime import date
+
+class UserService:
+  """ Class to manage all the user info """
+
+  GEOCODING_QUOTA_KEY = "geocoding_quota"
+  GEOCODING_SOFT_LIMIT_KEY = "soft_geocoder_limit"
+
+  REDIS_CONNECTION_KEY = "redis_connection"
+  REDIS_CONNECTION_HOST = "redis_host"
+  REDIS_CONNECTION_PORT = "redis_port"
+  REDIS_CONNECTION_DB = "redis_db"
+
+  REDIS_DEFAULT_USER_DB = 5
+  REDIS_DEFAULT_HOST = 'localhost'
+  REDIS_DEFAULT_PORT = 6379
+
+  def __init__(self, user_id, **kwargs):
+    self.user_id = user_id
+    if self.REDIS_CONNECTION_KEY in kwargs:
+      self._redis_connection = self.__get_redis_connection(redis_connection=kwargs[self.REDIS_CONNECTION_KEY])
+    else:
+      if self.REDIS_CONNECTION_HOST not in kwargs:
+        raise Exception("You have to provide redis configuration")
+      redis_config = self.__build_redis_config(kwargs)
+      self._redis_connection = self.__get_redis_connection(redis_config = redis_config)
+
+  def user_quota(self):
+      # Check for exceptions or redis timeout
+      user_quota = self._redis_connection.hget(self.__get_user_redis_key(), self.GEOCODING_QUOTA_KEY)
+      return int(user_quota) if user_quota and int(user_quota) >= 0 else 0
+
+  def soft_geocoder_limit(self):
+    """ Check what kind of limit the user has """
+    soft_limit = self._redis_connection.hget(self.__get_user_redis_key(), self.GEOCODING_SOFT_LIMIT_KEY)
+    return True if soft_limit == '1' else False
+
+  def used_quota_month(self, year, month):
+      """ Recover the used quota for the user in the current month """
+      # Check for exceptions or redis timeout
+      current_used = 0
+      for _, value in self._redis_connection.hscan_iter(self.__get_month_redis_key(year,month)):
+          current_used += int(value)
+      return current_used
+
+  def increment_geocoder_use(self, year, month, key, amount=1):
+      # TODO Manage exceptions or timeout
+      self._redis_connection.hincrby(self.__get_month_redis_key(year, month),key,amount)
+
+  @property
+  def redis_connection(self):
+      return self._redis_connection
+
+  def __get_redis_connection(self, redis_connection=None, redis_config=None):
+      if redis_connection:
+          conn = redis_connection
+      else:
+          conn = self.__create_redis_connection(redis_config)
+
+      return conn
+
+  def __create_redis_connection(self, redis_config):
+      pool = redis.ConnectionPool(host=redis_config['host'], port=redis_config['port'], db=redis_config['db'])
+      conn = redis.Redis(connection_pool=pool)
+      return conn
+
+  def __build_redis_config(self, config):
+      redis_host = config[self.REDIS_CONNECTION_HOST] if self.REDIS_CONNECTION_HOST in config else self.REDIS_DEFAULT_HOST
+      redis_port = config[self.REDIS_CONNECTION_PORT] if self.REDIS_CONNECTION_PORT in config else self.REDIS_DEFAULT_PORT
+      redis_db = config[self.REDIS_CONNECTION_DB] if self.REDIS_CONNECTION_DB in config else self.REDIS_DEFAULT_USER_DB
+      return {'host': redis_host, 'port': redis_port, 'db': redis_db}
+
+  def __get_month_redis_key(self, year, month):
+      today = date.today()
+      return "geocoder:{0}:{1}{2}".format(self.user_id, year, month)
+
+  def __get_user_redis_key(self):
+      return "geocoder:{0}".format(self.user_id)

server/lib/python/cartodb_geocoder/example/client_func_example.sql

@@ -0,0 +1,29 @@
+# cdb_conf geocoder config example
+INSERT INTO cdb_conf VALUES ('geocoder_conf', '{"geocoder_db": {"host": "localhost", "port": "5432", db": "cartodb_dev_user_274bf952-8568-4598-9efd-be92ed3d2ead_db", "user": "development_cartodb_user_274bf952-8568-4598-9efd-be92ed3d2ead"}, "redis": {"host": "localhost", "port": 6379, "db": 5 } }')
+
+CREATE OR REPLACE FUNCTION cartodb._geocoder_admin0_polygons(search text)
+    RETURNS Geometry AS
+$$
+    db_connection_str = plpy.execute("SELECT * FROM cartodb._Geocoder_Server_Conf() conf;")[0]['conf']
+    return plpy.execute("SELECT cartodb._Geocoder_Admin0_Polygons('{0}', session_user, txid_current(), '{1}') as geom".format(search, db_connection_str))[0]['geom']
+$$ LANGUAGE plpythonu SECURITY DEFINER;
+
+CREATE OR REPLACE
+FUNCTION cartodb._geocoder_server_conf()
+    RETURNS text AS
+$$
+    conf = plpy.execute("SELECT cartodb.CDB_Conf_GetConf('geocoder_conf') conf")[0]['conf']
+    if conf is None:
+      raise "There is no geocoder server configuration "
+    else:
+      import json
+      params = json.loads(conf)
+      db_params = params['geocoder_db']
+      return "host={0} port={1} dbname={2} user={3}".format(db_params['host'],db_params['port'],db_params['db'],db_params['user'])
+$$ LANGUAGE 'plpythonu';
+
+CREATE OR REPLACE FUNCTION cartodb._geocoder_admin0_polygons(search text, user_id name, tx_id bigint, db_connection_str text)
+RETURNS Geometry AS $$
+    CONNECT db_connection_str;
+    SELECT geocode_admin0(search, tx_id, user_id);
+$$ LANGUAGE plproxy;

server/lib/python/cartodb_geocoder/example/server_func_example.sql

@@ -0,0 +1,28 @@
+CREATE OR REPLACE
+FUNCTION geocode_admin0(search text, tx_id bigint, user_id name)
+  RETURNS Geometry AS
+$$
+  import logging
+  from cartodb_geocoder import quota_service
+
+  LOG_FILENAME = '/tmp/plpython.log'
+  logging.basicConfig(filename=LOG_FILENAME,level=logging.DEBUG)
+
+  if user_id in SD and tx_id in SD[user_id] and 'redis_connection' in SD[user_id][tx_id]:
+    logging.debug("Using redis cached connection...")
+    qs = quota_service.QuotaService(user_id, tx_id, redis_connection=SD[user_id][tx_id]['redis_connection'])
+  else:
+    qs = quota_service.QuotaService(user_id, tx_id, redis_host='localhost', redis_port=6379, redis_db=5)
+
+  if qs.check_user_quota():
+    result = plpy.execute("SELECT geom FROM geocode_admin0_polygons(Array[\'{0}\']::text[])".format(search))
+    if result.status() == 5 and result.nrows() == 1:
+      qs.increment_geocoder_use()
+      SD[user_id] = {tx_id: {'redis_connection': qs.user_service.redis_connection}}
+      return result[0]["geom"]
+    else:
+      raise Exception('Something wrong with the georefence operation')
+  else:
+    raise Exception('Not enough quota for this user')
+
+$$ LANGUAGE plpythonu;

server/lib/python/cartodb_geocoder/requirements.txt

@@ -0,0 +1 @@
+redis-py==2.10.5

server/lib/python/cartodb_geocoder/setup.py

@@ -0,0 +1,40 @@
+"""
+CartoDB Geocoder Python Library
+
+See:
+https://github.com/CartoDB/geocoder-api
+"""
+
+from setuptools import setup, find_packages
+
+setup(
+    name='cartodb_geocoder',
+
+    version='0.0.1',
+
+    description='CartoDB Geocoder Python Library',
+
+    url='https://github.com/CartoDB/geocoder-api',
+
+    author='Data Services Team - CartoDB',
+    author_email='dataservices@cartodb.com',
+
+    license='MIT',
+
+    classifiers=[
+        'Development Status :: 4 - Beta',
+        'Intended Audience :: Mapping comunity',
+        'Topic :: Maps :: Mapping Tools',
+        'License :: OSI Approved :: MIT License',
+        'Programming Language :: Python :: 2.7',
+    ],
+
+    keywords='maps api mapping tools geocoder',
+
+    packages=find_packages(exclude=['contrib', 'docs', 'tests']),
+
+    extras_require={
+        'dev': ['unittest'],
+        'test': ['unittest', 'nose', 'mockredispy'],
+    }
+)

server/lib/python/cartodb_geocoder/test/test_quota_service.py

@@ -0,0 +1,34 @@
+from mockredis import MockRedis
+from cartodb_geocoder import quota_service
+from unittest import TestCase
+from nose.tools import assert_raises
+
+
+class TestQuotaService(TestCase):
+
+  def setUp(self):
+    self.fake_redis_connection = MockRedis()
+    self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 100)
+    self.fake_redis_connection.hset('geocoder:user_id','soft_geocoder_limit', 0)
+    self.qs = quota_service.QuotaService('user_id', 'tx_id', redis_connection = self.fake_redis_connection)
+
+  def test_should_return_true_if_quota_with_no_use(self):
+    assert self.qs.check_user_quota() == True
+
+  def test_should_return_true_if_quota_is_not_completely_used(self):
+    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10)
+    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10)
+    assert self.qs.check_user_quota() == True
+
+  def test_should_return_false_if_quota_is_surpassed(self):
+    self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 1)
+    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10)
+    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10)
+    assert self.qs.check_user_quota() == False
+
+  def test_should_return_true_if_quota_is_surpassed_but_soft_limit_is_enabled(self):
+    self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 1)
+    self.fake_redis_connection.hset('geocoder:user_id','soft_geocoder_limit', 1)
+    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10)
+    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10)
+    assert self.qs.check_user_quota() == True

server/lib/python/cartodb_geocoder/test/test_user_service.py

@@ -0,0 +1,37 @@
+from mockredis import MockRedis
+from cartodb_geocoder import user_service
+from unittest import TestCase
+from nose.tools import assert_raises
+
+
+class TestUserService(TestCase):
+
+    def setUp(self):
+      self.fake_redis_connection = MockRedis()
+      self.us = user_service.UserService('user_id', redis_connection = self.fake_redis_connection)
+
+    def test_user_quota_should_be_10(self):
+      self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 10)
+      assert self.us.user_quota() == 10
+
+    def test_should_return_0_if_negative_quota(self):
+      self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', -10)
+      assert self.us.user_quota() == 0
+
+    def test_should_return_0_if_not_user(self):
+      assert self.us.user_quota() == 0
+
+    def test_user_used_quota_for_a_month(self):
+      self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10)
+      self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10)
+      assert self.us.used_quota_month(2015, 11) == 20
+
+    def test_user_not_amount_in_used_quota_for_month_should_be_0(self):
+      assert self.us.used_quota_month(2015, 11) == 0
+
+    def test_increment_used_quota(self):
+      self.us.increment_geocoder_use(2015, 11, 'tx_id', 1)
+      assert self.us.used_quota_month(2015, 11) == 1
+
+    def test_exception_if_not_redis_config(self):
+      assert_raises(Exception, user_service.UserService, 'user_id')