Complete rewrite of NGXImage

2 files changed, 105 insertions, 99 deletions

diff --git a/app/imutil.py b/app/imutil.py
index 015a2c1..05d5436 100644
--- a/app/imutil.py
+++ b/app/imutil.py
@@ -1,6 +1,7 @@
 from PIL import Image
 from app import blure
-from sanic.response import raw
+from sanic.response import raw, BaseHTTPResponse
+from sanic.exceptions import NotFound
 from pathlib import Path
 from io import BytesIO
 
@@ -14,32 +15,15 @@ class InvalidImageFormat(ValueError):
     pass
 
 
-class NGXImage:
-    def __init__(self, id: int):
-        self.filename = blure.url.to_url(id)
-
-    @staticmethod
-    def _load_pic(filename):
-        p = Path(_IMAGE_PATH.format(filename))
-        if not p.is_file():
-            return NGXImage.not_found()
-
-        return raw(b'',
-                   content_type='image',
-                   headers={'X-Accel-Redirect': _IMAGE_URL.format(filename)},
-                   status=200)
-
-    def orig(self):
-        return self._load_pic(self.filename)
+async def is_image_exists(id: int):
+    async with blure.pool.acquire() as conn:
+        rec = await conn.fetchval('SELECT TRUE FROM pics WHERE id=$1', id)
+        return rec is not None
 
-    def thumb(self):
-        return self._load_pic(self.filename + '_thumb')
 
-    @staticmethod
-    def not_found():
-        return raw(_NOT_FOUND_IMAGE,
-                   content_type=_NOT_FOUND_IMAGE_CONTENT_TYPE,
-                   status=404)
+class NGXImage:
+    def __init__(self, from_id: int):
+        self.id = from_id
 
     @staticmethod
     def pillow_format(content_type: str):
@@ -55,25 +39,81 @@ class NGXImage:
         else:
             raise InvalidImageFormat(f'{content_type} is not supported')
 
-    def save(self, body: BytesIO, content_type: str):
-        image_path = Path(_IMAGE_PATH.format(self.filename))
-        thumb_path = Path(_IMAGE_PATH.format(self.filename + '_thumb'))
+    @classmethod
+    async def create_from_bytes(cls, bytes_io: BytesIO, content_type: str):
+        async with blure.pool.acquire() as conn:
+            # TODO: make it a transaction
+            new_id = await conn.fetchval(
+                '''
+                    INSERT INTO pics(src_url, content_type)
+                    VALUES ($1, $2)
+                    RETURNING id
+                ''',
+                '',
+                content_type
+            )
+
+            ngx_image = NGXImage(from_id=new_id)
+
+            with ngx_image.orig_path.open('wb') as f:
+                f.write(bytes_io.getvalue())
+
+            image = Image.open(bytes_io)
+            image.thumbnail(blure.config.CUT_SIZES[2])
+            image.save(
+                ngx_image.thumb_path,
+                format=cls.pillow_format(content_type)
+            )
+
+            return ngx_image
+
+    async def __aenter__(self):
+        if not await is_image_exists(self.id):
+            raise NotFound(
+                f'image {blure.url.to_url(self.id)} does not exist'
+            )
+        else:
+            return self
 
-        with image_path.open('wb') as f:
-            f.write(body.getvalue())
+    async def __aexit__(self, *exc):
+        pass
 
-        with thumb_path.open('wb') as f:
-            im = Image.open(body)
-            thumb_stream = BytesIO()
-            im.thumbnail(blure.config.CUT_SIZES[2])
-            im.save(thumb_stream, format=self.pillow_format(content_type))
-            f.write(thumb_stream.getvalue())
+    @property
+    def orig_path(self):
+        return Path(_IMAGE_PATH.format(blure.url.to_url(self.id)))
 
-    def delete_from_disk(self):
-        image_path = Path(_IMAGE_PATH.format(self.filename))
-        thumb_path = Path(_IMAGE_PATH.format(self.filename + '_thumb'))
+    @property
+    def thumb_path(self):
+        return Path(_IMAGE_PATH.format(blure.url.to_url(self.id)) + '_thumb')
 
-        if image_path.exists():
-            image_path.unlink()
-        if thumb_path.exists():
-            thumb_path.unlink()
+    @classmethod
+    def _send_image(cls, filepath: Path) -> BaseHTTPResponse:
+        if not filepath.is_file():
+            return cls.not_found()
+
+        return raw(b'',
+                   content_type='image',
+                   headers={'X-Accel-Redirect': str(filepath)[4:]},  # FIXME: this should NOT be '[4:]' # noqa
+                   status=200)
+
+    def orig(self):
+        return self._send_image(self.orig_path)
+
+    def thumb(self):
+        return self._send_image(self.thumb_path)
+
+    @staticmethod
+    def not_found():
+        return raw(_NOT_FOUND_IMAGE,
+                   content_type=_NOT_FOUND_IMAGE_CONTENT_TYPE,
+                   status=404)
+
+    async def delete_from_db(self):
+        async with blure.pool.acquire() as conn:
+            await conn.execute('DELETE FROM pics WHERE id=$1', self.id)
+
+    def delete_from_disk(self):
+        if self.orig_path.exists():
+            self.orig_path.unlink()
+        if self.thumb_path.exists():
+            self.thumb_path.unlink()
diff --git a/app/views.py b/app/views.py
index f81cfd4..fcd3818 100644
--- a/app/views.py
+++ b/app/views.py
@@ -2,16 +2,11 @@ from io import BytesIO
 from .request_routine import db_route
 from sanic.response import text, redirect, json
 from jinja2_sanic import render_template
-from app import log, blure
+from app import blure
 from requests import get as fetch_url
 from .imutil import NGXImage
 from .util import URLDecodeError
-from sanic.exceptions import NotFound, InvalidUsage as BadRequest
-
-
-async def is_image_exists(pg, id: int):
-    rec = await pg.fetchval('SELECT 1 FROM pics WHERE id=$1', id)
-    return rec is not None
+from sanic.exceptions import NotFound, InvalidUsage as BadRequest, ServerError
 
 
 @blure.exception(NotFound)
@@ -19,6 +14,11 @@ async def not_found(req, exc):
     return render_template('404.html.j2', req, dict())
 
 
+@blure.exception(URLDecodeError)
+async def handle_urldecode(req, exc):
+    raise ServerError(str(exc), status_code=400)
+
+
 @db_route('/')
 async def index(ctx):
     records = await ctx.pg.fetch('SELECT id FROM pics')
@@ -28,46 +28,24 @@ async def index(ctx):
 
 @db_route('/i/<url>')
 async def raw_image(ctx, url):
-    try:
-        id = ctx.app.url.to_id(url)
-    except URLDecodeError as err:
-        log.warn(f'URL not decoded: {err}')
-        return NGXImage.not_found()
-
-    if not await is_image_exists(ctx.pg, id):
-        log.error('Image not in db')
-        return NGXImage.not_found()
-
-    return NGXImage(id).orig()
+    async with NGXImage(blure.url.to_id(url)) as image:
+        return image.orig()
 
 
 @db_route('/t/<url>')
 async def thumb_image(ctx, url):
-    try:
-        id = ctx.app.url.to_id(url)
-    except URLDecodeError:
-        return NGXImage.not_found()
-
-    if not await is_image_exists(ctx.pg, id):
-        log.error('Image not in db')
-        return NGXImage.not_found()
-
-    return NGXImage(id).thumb()
+    async with NGXImage(blure.url.to_id(url)) as image:
+        return image.orig()
 
 
 @db_route('/p/<url>')
 async def pic_profile(ctx, url):
-    log.warn(url)
-    try:
-        id = ctx.app.url.to_id(url)
-        if await is_image_exists(ctx.pg, id):
-            return render_template('profile.html.j2',
-                                   ctx.r,
-                                   dict(url=url, tags=[]))
-        else:
-            raise NotFound('Url not found')
-    except URLDecodeError:
-        raise NotFound('Invalid url')
+    # Check image exists
+    # This will be used to get meta
+    async with NGXImage(blure.url.to_id(url)):
+        return render_template('profile.html.j2',
+                               ctx.r,
+                               dict(url=url, tags=[]))
 
 
 @db_route('/c/push', methods=['POST'])
@@ -75,27 +53,15 @@ async def pic_push(ctx):
     try:
         file = ctx.r.files['im'][0]
         image_stream = BytesIO(file.body)
-        ext = file.name.split('.')[-1]
 
         if len(ctx.r.form['content-type']) != 1:
             raise BadRequest('Need only one content-type')
 
         content_type = ctx.r.form['content-type'][0]
 
-        id = await ctx.pg.fetchval(
-            '''
-            INSERT INTO pics(src_url, content_type)
-            VALUES ($1, $2)
-            RETURNING id
-            ''',
-            '',
-            ext
-        )
-
-        im = NGXImage(id)
-        im.save(image_stream, content_type)
+        image = await NGXImage.create_from_bytes(image_stream, content_type)
 
-        return text('new url is ' + ctx.app.url.to_url(id))
+        return text(ctx.app.url.to_url(image.id))
     except KeyError:
         return text('you did not post anything')
 
@@ -112,7 +78,7 @@ async def push_url(ctx):
 @db_route('/c/delete/<url>', methods=['POST'])
 async def delete_pic(ctx, url):
     id = ctx.app.url.to_id(url)
-    await ctx.pg.execute('DELETE FROM pics WHERE id=$1', id)
-    im = NGXImage(id)
-    im.delete_from_disk()
+    async with NGXImage(id) as image:
+        await image.delete_from_db()
+        image.delete_from_disk()
     return redirect(ctx.app.url_for('index'))