diff --git a/docs/index.rst b/docs/index.rst index 4a7f2dc..cf92acd 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -14,6 +14,7 @@ Basic metrics tracking for your `Flask`_ application. The core of library is ver The following is optional: * `freegeoip.net `_ integration for storing geography of the visitor. +* Alternatively you can define a custom geolocation function, for use with the geoip2 package for example. * Unique visitor tracking if you are wanting to use Flask's cookie storage. * Summation hooks for live count of common web analysis statistics such as hit counts. @@ -55,6 +56,8 @@ Usage app.config['TRACK_USAGE_USE_FREEGEOIP'] = False # You can use a different instance of freegeoip like so # app.config['TRACK_USAGE_FREEGEOIP_ENDPOINT'] = 'https://example.org/api/' + # (You can also define a custom geolocation function when initializing the extension) + app.config['TRACK_USAGE_INCLUDE_OR_EXCLUDE_VIEWS'] = 'include' # We will just print out the data for the example @@ -68,7 +71,9 @@ Usage t = TrackUsage(app, [ PrintWriter(), OutputWriter(transform=lambda s: "OUTPUT: " + str(s)) - ]) + ], + custom_geolocation_function # (Optional) + ) # Include the view in the metrics @t.include @@ -156,7 +161,8 @@ TRACK_USAGE_USE_FREEGEOIP **Default**: False -Turn FreeGeoIP integration on or off. If set to true, then geography information is also stored in the usage logs. +Turn FreeGeoIP integration on or off. If set to true, then the geography information is also stored in the usage logs. +Alternatively you can define a custom geolocation lookup function when initializing the extension. .. versionchanged:: 1.1. The default server for using geoip integration changed to extreme-ip-lookup.com diff --git a/src/flask_track_usage/__init__.py b/src/flask_track_usage/__init__.py index fc8d14f..fb31ba2 100644 --- a/src/flask_track_usage/__init__.py +++ b/src/flask_track_usage/__init__.py @@ -56,7 +56,7 @@ class TrackUsage(object): Tracks basic usage of Flask applications. """ - def __init__(self, app=None, storage=None, _fake_time=None): + def __init__(self, app=None, storage=None, ip_lookup_func=None, _fake_time=None): """ Create the instance. @@ -78,9 +78,9 @@ def __init__(self, app=None, storage=None, _fake_time=None): self._fake_time = _fake_time if app is not None and storage is not None: - self.init_app(app, storage) + self.init_app(app, storage, ip_lookup_func) - def init_app(self, app, storage): + def init_app(self, app, storage, ip_lookup_func): """ Initialize the instance with the app. @@ -90,6 +90,7 @@ def init_app(self, app, storage): """ self.app = app self._storages = storage + self.ip_lookup_func = ip_lookup_func self._use_freegeoip = app.config.get( 'TRACK_USAGE_USE_FREEGEOIP', False) self._freegeoip_endpoint = app.config.get( @@ -190,7 +191,10 @@ def after_request(self, response): data['username'] = str(ctx.request.authorization.username) elif getattr(self.app, 'login_manager', None) and current_user and not current_user.is_anonymous: data['username'] = str(current_user) - if self._use_freegeoip: + if self.ip_lookup_func: + clean_ip = quote_plus(str(ctx.request.remote_addr)) + data['ip_info'] = self.ip_lookup_func(clean_ip) + elif self._use_freegeoip: clean_ip = quote_plus(str(ctx.request.remote_addr)) if '{ip}' in self._freegeoip_endpoint: url = self._freegeoip_endpoint.format(ip=clean_ip) diff --git a/src/flask_track_usage/storage/sql.py b/src/flask_track_usage/storage/sql.py index 35422a1..8ff4f8e 100644 --- a/src/flask_track_usage/storage/sql.py +++ b/src/flask_track_usage/storage/sql.py @@ -79,6 +79,7 @@ def set_up(self, engine=None, metadata=None, table_name="flask_usage", """ import sqlalchemy as sql + from sqlalchemy.dialects.postgresql import JSONB if db: self._eng = db.engine self._metadata = db.metadata @@ -106,7 +107,7 @@ def set_up(self, engine=None, metadata=None, table_name="flask_usage", sql.Column('remote_addr', sql.String(24)), sql.Column('xforwardedfor', sql.String(24)), sql.Column('authorization', sql.Boolean), - sql.Column('ip_info', sql.String(1024)), + sql.Column('ip_info', JSONB()), sql.Column('path', sql.String(128)), sql.Column('speed', sql.Float), sql.Column('datetime', sql.DateTime), @@ -118,7 +119,10 @@ def set_up(self, engine=None, metadata=None, table_name="flask_usage", else: self._metadata.reflect(bind=self._eng) self.track_table = self._metadata.tables[table_name] - + # Patch the ip_info column type after reflection + if 'ip_info' in self.track_table.c: + self.track_table.c.ip_info.type = JSONB() + def store(self, data): """ Executed on "function call". @@ -131,16 +135,7 @@ def store(self, data): """ user_agent = data["user_agent"] utcdatetime = datetime.datetime.fromtimestamp(data['date']) - if data["ip_info"]: - t = {} - for key in data["ip_info"]: - t[key] = data["ip_info"][key] - if not len(json.dumps(t)) < 1024: - del t[key] - break - ip_info_str = json.dumps(t) - else: - ip_info_str = None + with self._eng.begin() as con: stmt = self.track_table.insert().values( url=data['url'], @@ -156,7 +151,7 @@ def store(self, data): remote_addr=data["remote_addr"], xforwardedfor=data["xforwardedfor"], authorization=data["authorization"], - ip_info=ip_info_str, + ip_info=data["ip_info"], path=data["path"], speed=data["speed"], datetime=utcdatetime,