Procházet zdrojové kódy

Uvicorn app improvements (#3918)

* Upgrade Uvicorn

* Fix error in accessing queries attribute

* Improvements to app.py for uvicorn (prepared statements etc)

* Fix typo on variable name
Tom Christie před 7 roky
rodič
revize
e390271443
1 změnil soubory, kde provedl 103 přidání a 120 odebrání
  1. 103 120
      frameworks/Python/uvicorn/app.py

+ 103 - 120
frameworks/Python/uvicorn/app.py

@@ -2,11 +2,10 @@ import asyncio
 import asyncpg
 import jinja2
 import os
-import ujson as json
-from functools import partial
+import ujson
 from random import randint
 from operator import itemgetter
-from urllib import parse
+from urllib.parse import parse_qs
 
 
 async def setup():
@@ -20,9 +19,38 @@ async def setup():
     )
 
 
+READ_ROW_SQL = 'SELECT "randomnumber" FROM "world" WHERE id = $1'
+WRITE_ROW_SQL = 'UPDATE "world" SET "randomnumber"=$1 WHERE id=$2'
+ADDITIONAL_ROW = [0, 'Additional fortune added at request time.']
+
+JSON_RESPONSE = {
+    'type': 'http.response.start',
+    'status': 200,
+    'headers': [
+        [b'content-type', b'application/json'],
+    ]
+}
+
+HTML_RESPONSE = {
+    'type': 'http.response.start',
+    'status': 200,
+    'headers': [
+        [b'content-type', b'text/html; charset=utf-8'],
+    ]
+}
+
+PLAINTEXT_RESPONSE = {
+    'type': 'http.response.start',
+    'status': 200,
+    'headers': [
+        [b'content-type', b'text/plain; charset=utf-8'],
+    ]
+}
+
+
 pool = None
-additional = [0, 'Additional fortune added at request time.']
 key = itemgetter(1)
+json_dumps = ujson.dumps
 template = None
 path = os.path.join('templates', 'fortune.html')
 with open(path, 'r') as template_file:
@@ -33,23 +61,18 @@ loop = asyncio.get_event_loop()
 loop.run_until_complete(setup())
 
 
-def get_query_count(query_string):
-    # helper to deal with the querystring passed in
-    queries = parse.parse_qs(query_string).get(b'queries', [None])[0]
-    if queries:
-        try:
-            query_count = int(queries)
-            if query_count < 1:
-                return 1
-            if query_count > 500:
-                return 500
-            return query_count
-        except ValueError:
-            pass
-    return 1
+def get_num_queries(scope):
+    try:
+        query_string = scope['query_string']
+        query_count = int(parse_qs(query_string)[b'queries'][0])
+    except (KeyError, IndexError, ValueError):
+        return 1
 
-
-random_int = partial(randint, 1, 10000)
+    if query_count < 1:
+        return 1
+    if query_count > 500:
+        return 500
+    return query_count
 
 
 class JSONSerialization:
@@ -60,15 +83,8 @@ class JSONSerialization:
         pass
 
     async def __call__(self, receive, send):
-        content = json.dumps({'message': 'Hello, world!'}).encode('utf-8')
-        await send({
-            'type': 'http.response.start',
-            'status': 200,
-            'headers': [
-                [b'content-type', b'application/json'],
-                [b'content-length', str(len(content)).encode()]
-            ]
-        })
+        content = json_dumps({'message': 'Hello, world!'}).encode('utf-8')
+        await send(JSON_RESPONSE)
         await send({
             'type': 'http.response.body',
             'body': content,
@@ -84,62 +100,53 @@ class SingleDatabaseQuery:
         pass
 
     async def __call__(self, receive, send):
+        row_id = randint(1, 10000)
         connection = await pool.acquire()
         try:
-            row = await connection.fetchrow('SELECT id, "randomnumber" FROM "world" WHERE id = ' + str(random_int()))
-            world = {'id': row[0], 'randomNumber': row[1]}
-
-            content = json.dumps(world).encode('utf-8')
-            await send({
-                'type': 'http.response.start',
-                'status': 200,
-                'headers': [
-                    [b'content-type', b'application/json'],
-                    [b'content-length', str(len(content)).encode()]
-                ]
-            })
-            await send({
-                'type': 'http.response.body',
-                'body': content,
-                'more_body': False
-            })
+            number = await connection.fetchval(READ_ROW_SQL, row_id)
+            world = {'id': row_id, 'randomNumber': number}
         finally:
             await pool.release(connection)
 
+        content = json_dumps(world).encode('utf-8')
+        await send(JSON_RESPONSE)
+        await send({
+            'type': 'http.response.body',
+            'body': content,
+            'more_body': False
+        })
+
+
 
 class MultipleDatabaseQueries:
     """
     Test type 3: Multiple database queries
     """
     def __init__(self, scope):
-        self.queries = get_query_count(scope.get('query_string', {}))
+        self.scope = scope
 
     async def __call__(self, receive, send):
+        num_queries = get_num_queries(self.scope)
+        row_ids = [randint(1, 10000) for _ in range(num_queries)]
+        worlds = []
+
         connection = await pool.acquire()
         try:
-            worlds = []
-            for i in range(self.queries):
-                sql = 'SELECT id, "randomnumber" FROM "world" WHERE id = ' + str(random_int())
-                row = await connection.fetchrow(sql)
-                worlds.append({'id': row[0], 'randomNumber': row[1]})
-
-            content = json.dumps(worlds).encode('utf-8')
-            await send({
-                'type': 'http.response.start',
-                'status': 200,
-                'headers': [
-                    [b'content-type', b'application/json'],
-                    [b'content-length', str(len(content)).encode()]
-                ]
-            })
-            await send({
-                'type': 'http.response.body',
-                'body': content,
-                'more_body': False
-            })
+            statement = await connection.prepare(READ_ROW_SQL)
+            for row_id in row_ids:
+                number = await statement.fetchval(row_id)
+                worlds.append({'id': row_id, 'randomNumber': number})
         finally:
             await pool.release(connection)
 
+        content = json_dumps(worlds).encode('utf-8')
+        await send(JSON_RESPONSE)
+        await send({
+            'type': 'http.response.body',
+            'body': content,
+            'more_body': False
+        })
+
 
 class Fortunes:
     """
@@ -152,59 +159,49 @@ class Fortunes:
         connection = await pool.acquire()
         try:
             fortunes = await connection.fetch('SELECT * FROM Fortune')
-            fortunes.append(additional)
-            fortunes.sort(key=key)
-            content = template.render(fortunes=fortunes).encode('utf-8')
-            await send({
-                'type': 'http.response.start',
-                'status': 200,
-                'headers': [
-                    [b'content-type', b'text/html; charset=utf-8'],
-                    [b'content-length', str(len(content)).encode()]
-                ]
-            })
-            await send({
-                'type': 'http.response.body',
-                'body': content,
-                'more_body': False
-            })
         finally:
             await pool.release(connection)
 
+        fortunes.append(ADDITIONAL_ROW)
+        fortunes.sort(key=key)
+        content = template.render(fortunes=fortunes).encode('utf-8')
+        await send(HTML_RESPONSE)
+        await send({
+            'type': 'http.response.body',
+            'body': content,
+            'more_body': False
+        })
+
 
 class DatabaseUpdates:
     """
     Test type 5: Database updates
     """
     def __init__(self, scope):
-        self.queries = get_query_count(scope.get('query_string', {}))
+        self.scope = scope
 
     async def __call__(self, receive, send):
+        num_queries = get_num_queries(self.scope)
+        updates = [(randint(1, 10000), randint(1, 10000)) for _ in range(num_queries)]
+        worlds = [{'id': row_id, 'randomNumber': number} for row_id, number in updates]
+
         connection = await pool.acquire()
         try:
-            worlds = []
-            for i in range(self.queries):
-                row = await connection.fetchrow('SELECT id FROM "world" WHERE id=' + str(random_int()))
-                worlds.append({'id': row[0], 'randomNumber': random_int()})
-                await connection.execute('UPDATE "world" SET "randomnumber"=%s WHERE id=%s' % (random_int(), row[0]))
-
-            content = json.dumps(worlds).encode('utf-8')
-            await send({
-                'type': 'http.response.start',
-                'status': 200,
-                'headers': [
-                    [b'content-type', b'application/json'],
-                    [b'content-length', str(len(content)).encode()]
-                ]
-            })
-            await send({
-                'type': 'http.response.body',
-                'body': content,
-                'more_body': False
-            })
+            statement = await connection.prepare(READ_ROW_SQL)
+            for row_id, _ in updates:
+                await statement.fetchval(row_id)
+            await connection.executemany(WRITE_ROW_SQL, updates)
         finally:
             await pool.release(connection)
 
+        content = json_dumps(worlds).encode('utf-8')
+        await send(JSON_RESPONSE)
+        await send({
+            'type': 'http.response.body',
+            'body': content,
+            'more_body': False
+        })
+
 
 class Plaintext:
     """
@@ -215,14 +212,7 @@ class Plaintext:
 
     async def __call__(self, receive, send):
         content = b'Hello, world!'
-        await send({
-            'type': 'http.response.start',
-            'status': 200,
-            'headers': [
-                [b'content-type', b'text/plain'],
-                [b'content-length', str(len(content)).encode()]
-            ]
-        })
+        await send(PLAINTEXT_RESPONSE)
         await send({
             'type': 'http.response.body',
             'body': content,
@@ -236,14 +226,7 @@ class Handle404:
 
     async def __call__(self, receive, send):
         content = b'Not found'
-        await send({
-            'type': 'http.response.start',
-            'status': 200,
-            'headers': [
-                [b'content-type', b'text/plain'],
-                [b'content-length', str(len(content)).encode()]
-            ]
-        })
+        await send(PLAINTEXT_RESPONSE)
         await send({
             'type': 'http.response.body',
             'body': content,