Extracted redis connection and quota check to be reusable

server/extension/sql/0.0.1/11_redis_helper.sql

@@ -0,0 +1,11 @@
+-- Get the connection to redis from cache or create a new one
+CREATE OR REPLACE FUNCTION cdb_geocoder_server._connect_to_redis(user_id name)
+RETURNS boolean AS $$
+  if user_id in GD and 'redis_connection' in GD[user_id]:
+    return False
+  else:
+    from cartodb_geocoder import redis_helper
+    redis_connection = redis_helper.RedisHelper('localhost', 6379, 5).redis_connection()
+    GD[user_id] = {'redis_connection': redis_connection}
+    return True
+$$ LANGUAGE plpythonu;

server/extension/sql/0.0.1/30_admin0.sql

@@ -11,14 +11,23 @@ RETURNS Geometry AS $$
         plpy.error('The api_key must be provided')
 
     #--TODO: rate limiting check
-    #--TODO: quota check
+    #--This will create and cache a redis connection, if needed, in the GD object for the current user
+    redis_conn_plan = plpy.prepare("SELECT cdb_geocoder_server._connect_to_redis($1)", ["name"])
+    redis_conn_result = plpy.execute(redis_conn_plan, [user_id], 1)
+    qs = quota_service.QuotaService(user_id, tx_id, GD[user_id]['redis_connection'])
+
+    if not qs.check_user_quota():
+      plpy.error("Not enough quota for this user")
 
     #-- Copied from the doc, see http://www.postgresql.org/docs/9.4/static/plpython-database.html
     plan = plpy.prepare("SELECT cdb_geocoder_server._geocode_admin0_polygon($1) AS mypolygon", ["text"])
-    rv = plpy.execute(plan, [country_name], 1)
-
-    plpy.debug('Returning from Returning from geocode_admin0_polygons')
-    return rv[0]["mypolygon"]
+    result = plpy.execute(plan, [country_name], 1)
+    if result.status() == 5 and result.nrows() == 1:
+      qs.increment_geocoder_use()
+      plpy.debug('Returning from geocode_admin0_polygons')
+      return result[0]["mypolygon"]
+    else:
+      plpy.error('Something wrong with the georefence operation')
 $$ LANGUAGE plpythonu;
 
 

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

@@ -1,12 +1,11 @@
-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)
+    def __init__(self, user_id, transaction_id, redis_connection):
+        self._user_service = user_service.UserService(user_id, redis_connection)
         self.transaction_id = transaction_id
 
     def check_user_quota(self):
@@ -16,6 +15,7 @@ class QuotaService:
         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):

server/lib/python/cartodb_geocoder/cartodb_geocoder/redis_helper.py

@@ -0,0 +1,19 @@
+import redis
+
+class RedisHelper:
+
+  REDIS_DEFAULT_USER_DB = 5
+  REDIS_DEFAULT_HOST = 'localhost'
+  REDIS_DEFAULT_PORT = 6379
+
+  def __init__(self, host, port=REDIS_DEFAULT_PORT, db=REDIS_DEFAULT_USER_DB):
+    self.host = host
+    self.port = port
+    self.db = db
+
+  def redis_connection(self):
+    return self.__create_redis_connection()
+
+  def __create_redis_connection(self):
+    pool = redis.ConnectionPool(host=self.host, port=self.port, db=self.db)
+    return redis.Redis(connection_pool=pool)

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

@@ -1,4 +1,4 @@
-import redis
+import redis_helper
 from datetime import date
 
 class UserService:
@@ -12,19 +12,9 @@ class UserService:
   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):
+  def __init__(self, user_id, redis_connection):
     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)
+    self._redis_connection = redis_connection
 
   def user_quota(self):
       # Check for exceptions or redis timeout
@@ -52,25 +42,6 @@ class UserService:
   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)