diff --git a/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py b/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py index 3f4fca5..7c170cb 100644 --- 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) \ No newline at end of file + def get_user_service(self): + return self.user_service \ No newline at end of file diff --git a/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py b/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py new file mode 100644 index 0000000..5e7c578 --- /dev/null +++ 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) \ No newline at end of file diff --git a/server/lib/python/cartodb_geocoder/example/client_func_example.sql b/server/lib/python/cartodb_geocoder/example/client_func_example.sql index 7c0a580..4ae26d0 100644 --- 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); diff --git a/server/lib/python/cartodb_geocoder/example/server_func_example.sql b/server/lib/python/cartodb_geocoder/example/server_func_example.sql index 4f7bbc5..2190cda 100644 --- 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')