Refactor, extracted user_service

2015-11-10 16:32:53 +01:00
parent bb3b2cc2da
commit 7eab3e1f90
4 changed files with 86 additions and 77 deletions
--- a/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py
+++ b/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py
@@ -1,89 +1,25 @@
 import redis
+import user_service
 from datetime import date

 class QuotaService:
    """ Class to manage all the quota operation for the Geocoder SQL API Extension """

-    GEOCODING_QUOTA_KEY = "geocoding_quota"
-    REDIS_CONNECTION_KEY = "redis_connection"
-    REDIS_CONNECTION_HOST = "redis_host"
-    REDIS_CONNECTION_PORT = "redis_port"
-    REDIS_CONNECTION_DB = "redis_db"
-
    def __init__(self, logger, user_id, transaction_id, **kwargs):
        self.logger = logger
-        self.user_id = user_id
+        self.user_service = user_service.UserService(logger, user_id, **kwargs)
        self.transaction_id = transaction_id
-        self.cache = {}
-
-        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 "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 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.get_user_quota()
-        current_used = self.get_current_used_quota()
+        user_quota = self.user_service.get_user_quota()
+        current_used = self.user_service.get_current_used_quota()
        self.logger.debug("User quota: {0} --- Current used quota: {1}".format(user_quota, current_used))
        return True if (current_used + 1) < user_quota else False

-    def get_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 else 0
+    def increment_geocoder_use(self, amount=1):
+        self.user_service.increment_geocoder_use(self.transaction_id)

-    def get_current_used_quota(self):
-        """ 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()):
-            current_used += int(value)
-        return current_used
-
-    def increment_georeference_use(self, amount=1):
-        # TODO Manage exceptions or timeout
-        self.redis_connection.hincrby(self.__get_month_redis_key(), self.transaction_id,amount)
-
-    def get_redis_connection(self):
-        return self.redis_connection
-
-    def __get_redis_connection(self, redis_connection=None, redis_config=None):
-        if redis_connection:
-            self.__add_redis_connection_to_cache(redis_connection)
-            conn = redis_connection
-        else:
-            conn = self.__create_redis_connection(redis_config)
-
-        return conn
-
-    def __create_redis_connection(self, redis_config):
-        # Pool not needed
-        # Try to not create a connection every time, add it to a cache
-        self.logger.debug("Connecting to redis...")
-        pool = redis.ConnectionPool(host=redis_config['host'], port=redis_config['port'], db=redis_config['db'])
-        conn = redis.Redis(connection_pool=pool)
-        self.__add_redis_connection_to_cache(conn)
-        return conn
-
-    def __build_redis_config(self, config):
-        redis_host = config[self.REDIS_CONNECTION_HOST] if self.REDIS_CONNECTION_HOST in config else 'localhost'
-        redis_port = config[self.REDIS_CONNECTION_PORT] if self.REDIS_CONNECTION_PORT in config else 6379
-        redis_db = config[self.REDIS_CONNECTION_DB] if self.REDIS_CONNECTION_DB in config else 5
-        return {'host': redis_host, 'port': redis_port, 'db': redis_db}
-
-    def __add_redis_connection_to_cache(self, connection):
-        """ Cache the redis connection to avoid reach the limit of connections """
-        self.cache = {self.transaction_id: {self.REDIS_CONNECTION_KEY: connection}}
-
-    def __get_month_redis_key(self):
-        today = date.today()
-        return "geocoder:{0}:{1}".format(self.user_id, today.strftime("%Y%m"))
-
-    def __get_user_redis_key(self):
-        return "geocoder:{0}".format(self.user_id)
+    def get_user_service(self):
+        return self.user_service
--- a/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py
+++ b/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py
@@ -0,0 +1,73 @@
+import redis
+from datetime import date
+
+class UserService:
+  """ Class to manage all the user info """
+
+  GEOCODING_QUOTA_KEY = "geocoding_quota"
+  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, logger, user_id, **kwargs):
+    self.user_id = user_id
+    self.logger = logger
+    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 "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 get_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 else 0
+
+  def get_current_used_quota(self):
+      """ 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()):
+          current_used += int(value)
+      return current_used
+
+  def increment_geocoder_use(self, key, amount=1):
+      # TODO Manage exceptions or timeout
+      self.redis_connection.hincrby(self.__get_month_redis_key(),key,amount)
+
+  def get_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):
+      self.logger.debug("Connecting to redis...")
+      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_CONNECTION_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):
+      today = date.today()
+      return "geocoder:{0}:{1}".format(self.user_id, today.strftime("%Y%m"))
+
+  def __get_user_redis_key(self):
+      return "geocoder:{0}".format(self.user_id)
--- a/server/lib/python/cartodb_geocoder/example/client_func_example.sql
+++ b/server/lib/python/cartodb_geocoder/example/client_func_example.sql
@@ -1,7 +1,7 @@
 # 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)
+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']
@@ -9,7 +9,7 @@ $$
 $$ LANGUAGE plpythonu SECURITY DEFINER;

 CREATE OR REPLACE
-FUNCTION cartodb._Geocoder_Server_Conf()
+FUNCTION cartodb._geocoder_server_conf()
    RETURNS text AS
 $$
    conf = plpy.execute("SELECT cartodb.CDB_Conf_GetConf('geocoder_conf') conf")[0]['conf']
@@ -22,7 +22,7 @@ $$
      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)
+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);
--- a/server/lib/python/cartodb_geocoder/example/server_func_example.sql
+++ b/server/lib/python/cartodb_geocoder/example/server_func_example.sql
@@ -17,8 +17,8 @@ $$
  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_georeference_use()
-      SD[user_id] = {tx_id: {'redis_connection': qs.get_redis_connection()}}
+      qs.increment_geocoder_use()
+      SD[user_id] = {tx_id: {'redis_connection': qs.get_user_service().get_redis_connection()}}
      return result[0]["geom"]
    else:
      raise Exception('Something wrong with the georefence operation')