From da12d6628dae7880e4f1e41f1ae7ae198b9e33dc Mon Sep 17 00:00:00 2001 From: Mario de Frutos Date: Fri, 13 Nov 2015 16:44:26 +0100 Subject: [PATCH] Extracted redis connection and quota check to be reusable --- .../extension/sql/0.0.1/11_redis_helper.sql | 11 ++++++ server/extension/sql/0.0.1/30_admin0.sql | 19 +++++++--- .../cartodb_geocoder/quota_service.py | 6 ++-- .../cartodb_geocoder/redis_helper.py | 19 ++++++++++ .../cartodb_geocoder/user_service.py | 35 ++----------------- 5 files changed, 50 insertions(+), 40 deletions(-) create mode 100644 server/extension/sql/0.0.1/11_redis_helper.sql create mode 100644 server/lib/python/cartodb_geocoder/cartodb_geocoder/redis_helper.py diff --git a/server/extension/sql/0.0.1/11_redis_helper.sql b/server/extension/sql/0.0.1/11_redis_helper.sql new file mode 100644 index 0000000..593d9fb --- /dev/null +++ b/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; \ No newline at end of file diff --git a/server/extension/sql/0.0.1/30_admin0.sql b/server/extension/sql/0.0.1/30_admin0.sql index bb1fc2f..e5a5eeb 100644 --- a/server/extension/sql/0.0.1/30_admin0.sql +++ b/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; 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 5d61fad..eeeb97c 100644 --- a/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py +++ b/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): diff --git a/server/lib/python/cartodb_geocoder/cartodb_geocoder/redis_helper.py b/server/lib/python/cartodb_geocoder/cartodb_geocoder/redis_helper.py new file mode 100644 index 0000000..ddaa5af --- /dev/null +++ b/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) \ 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 index 3e22d48..3a0956b 100644 --- a/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py +++ b/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)