Merge pull request #6776 from yoff/python/model-asyncpg

Python: Model `asyncpg`

python/ql/test/library-tests/frameworks/asyncpg/ConceptsTest.ql

@@ -0,0 +1,2 @@
+import python
+import experimental.meta.ConceptsTest

python/ql/test/library-tests/frameworks/asyncpg/test.py

@@ -0,0 +1,105 @@
+import asyncio
+import asyncpg
+
+async def test_connection():
+    conn = await asyncpg.connect()
+
+    try:
+        # The file-like object is passed in as a keyword-only argument.
+        # See https://magicstack.github.io/asyncpg/current/api/index.html#asyncpg.connection.Connection.copy_from_query
+        await conn.copy_from_query("sql", output="filepath")  # $ getSql="sql" getAPathArgument="filepath"
+        await conn.copy_from_query("sql", "arg1", "arg2", output="filepath")  # $ getSql="sql" getAPathArgument="filepath"
+
+        await conn.copy_from_table("table", output="filepath")  # $ getAPathArgument="filepath"
+        await conn.copy_to_table("table", source="filepath")  # $ getAPathArgument="filepath"
+
+        await conn.execute("sql")  # $ getSql="sql"
+        await conn.executemany("sql")  # $ getSql="sql"
+        await conn.fetch("sql")  # $ getSql="sql"
+        await conn.fetchrow("sql")  # $ getSql="sql"
+        await conn.fetchval("sql")  # $ getSql="sql"
+
+    finally:
+        await conn.close()
+
+
+async def test_prepared_statement():
+    conn = await asyncpg.connect()
+
+    try:
+        pstmt = await conn.prepare("psql")  # $ constructedSql="psql"
+        pstmt.executemany()  # $ getSql="psql"
+        pstmt.fetch()  # $ getSql="psql"
+        pstmt.fetchrow()  # $ getSql="psql"
+        pstmt.fetchval()  # $ getSql="psql"
+
+    finally:
+        await conn.close()
+
+# The sql statement is executed when the `CursorFactory` (obtained by e.g. `conn.cursor()`) is awaited.
+# See https://magicstack.github.io/asyncpg/current/api/index.html#asyncpg.cursor.CursorFactory
+async def test_cursor():
+    conn = await asyncpg.connect()
+
+    try:
+        async with conn.transaction():
+            cursor = await conn.cursor("sql")  # $ getSql="sql" constructedSql="sql"
+            await cursor.fetch()
+
+            pstmt = await conn.prepare("psql")  # $ constructedSql="psql"
+            pcursor = await pstmt.cursor()  # $ getSql="psql"
+            await pcursor.fetch()
+
+            async for record in conn.cursor("sql"):  # $ getSql="sql" constructedSql="sql"
+                pass
+
+            async for record in pstmt.cursor():  # $ getSql="psql"
+                pass
+
+            cursor_factory = conn.cursor("sql")  # $ constructedSql="sql"
+            cursor = await cursor_factory  # $ getSql="sql"
+
+            pcursor_factory = pstmt.cursor()
+            pcursor = await pcursor_factory  # $ getSql="psql"
+
+    finally:
+        await conn.close()
+
+async def test_connection_pool():
+    pool = await asyncpg.create_pool()
+
+    try:
+        await pool.copy_from_query("sql", output="filepath")  # $ getSql="sql" getAPathArgument="filepath"
+        await pool.copy_from_query("sql", "arg1", "arg2", output="filepath")  # $ getSql="sql" getAPathArgument="filepath"
+        await pool.copy_from_table("table", output="filepath")  # $ getAPathArgument="filepath"
+        await pool.copy_to_table("table", source="filepath")  # $ getAPathArgument="filepath"
+
+        await pool.execute("sql")  # $ getSql="sql"
+        await pool.executemany("sql")  # $ getSql="sql"
+        await pool.fetch("sql")  # $ getSql="sql"
+        await pool.fetchrow("sql")  # $ getSql="sql"
+        await pool.fetchval("sql")  # $ getSql="sql"
+
+        async with pool.acquire() as conn:
+            await conn.execute("sql")  # $ getSql="sql"
+
+        conn = await pool.acquire()
+        try:
+            await conn.fetch("sql")  # $ getSql="sql"
+        finally:
+            await pool.release(conn)
+
+    finally:
+        await pool.close()
+
+    async with asyncpg.create_pool() as pool:
+        await pool.execute("sql")  # $ getSql="sql"
+
+        async with pool.acquire() as conn:
+            await conn.execute("sql")  # $ getSql="sql"
+
+        conn = await pool.acquire()
+        try:
+            await conn.fetch("sql")  # $ getSql="sql"
+        finally:
+            await pool.release(conn)

python/ql/test/library-tests/frameworks/flask_sqlalchemy/SqlExecution.py

@@ -12,7 +12,7 @@ db = SQLAlchemy(app)
 # - https://github.com/pallets/flask-sqlalchemy/blob/931ec00d1e27f51508e05706eef41cc4419a0b32/src/flask_sqlalchemy/__init__.py#L765
 # - https://github.com/pallets/flask-sqlalchemy/blob/931ec00d1e27f51508e05706eef41cc4419a0b32/src/flask_sqlalchemy/__init__.py#L99-L109
 
-assert str(type(db.text("Foo"))) == "<class 'sqlalchemy.sql.elements.TextClause'>"
+assert str(type(db.text("Foo"))) == "<class 'sqlalchemy.sql.elements.TextClause'>"  # $ constructedSql="Foo"
 
 # also has engine/session instantiated
 
@@ -44,8 +44,8 @@ assert result.fetchall() == [("Foo",)]
 
 
 # text
-t = db.text("foo")
+t = db.text("foo")  # $ constructedSql="foo"
 assert isinstance(t, sqlalchemy.sql.expression.TextClause)
 
-t = db.text(text="foo")
+t = db.text(text="foo")  # $ constructedSql="foo"
 assert isinstance(t, sqlalchemy.sql.expression.TextClause)

python/ql/test/library-tests/frameworks/sqlalchemy/SqlExecution.py

@@ -46,7 +46,7 @@ with engine.begin() as connection:
     connection.execute("some sql")  # $ getSql="some sql"
 
 # Injection requiring the text() taint-step
-t = text("some sql")
+t = text("some sql")  # $ constructedSql="some sql"
 session.query(User).filter(t)
 session.query(User).group_by(User.id).having(t)
 session.query(User).group_by(t).first()

python/ql/test/library-tests/frameworks/sqlalchemy/new_tests.py

@@ -6,7 +6,7 @@ import sqlalchemy.orm
 # either v1.4 or v2.0, such that we cover both.
 
 raw_sql = "select 'FOO'"
-text_sql = sqlalchemy.text(raw_sql)
+text_sql = sqlalchemy.text(raw_sql)  # $ constructedSql=raw_sql
 
 Base = sqlalchemy.orm.declarative_base()
 
@@ -169,7 +169,7 @@ assert session.query(For14).all()[0].id == 14
 
 # and now we can do the actual querying
 
-text_foo = sqlalchemy.text("'FOO'")
+text_foo = sqlalchemy.text("'FOO'")  # $ constructedSql="'FOO'"
 
 # filter_by is only vulnerable to injection if sqlalchemy.text is used, which is evident
 # from the logs produced if this file is run
@@ -298,7 +298,7 @@ with engine.connect() as conn:
     assert scalar_result == "FOO"
 
     # This is a contrived example
-    select = sqlalchemy.select(sqlalchemy.text("'BAR'"))
+    select = sqlalchemy.select(sqlalchemy.text("'BAR'"))  # $ constructedSql="'BAR'"
     result = conn.execute(select) # $ getSql=select
     assert result.fetchall() == [("BAR",)]
 

python/ql/test/library-tests/frameworks/sqlalchemy/taint_test.py

@@ -8,14 +8,14 @@ def test_taint():
 
     ensure_tainted(ts) # $ tainted
 
-    t1 = sqlalchemy.text(ts)
-    t2 = sqlalchemy.text(text=ts)
-    t3 = sqlalchemy.sql.text(ts)
-    t4 = sqlalchemy.sql.text(text=ts)
-    t5 = sqlalchemy.sql.expression.text(ts)
-    t6 = sqlalchemy.sql.expression.text(text=ts)
-    t7 = sqlalchemy.sql.expression.TextClause(ts)
-    t8 = sqlalchemy.sql.expression.TextClause(text=ts)
+    t1 = sqlalchemy.text(ts)  # $ constructedSql=ts
+    t2 = sqlalchemy.text(text=ts)  # $ constructedSql=ts
+    t3 = sqlalchemy.sql.text(ts)  # $ constructedSql=ts
+    t4 = sqlalchemy.sql.text(text=ts)  # $ constructedSql=ts
+    t5 = sqlalchemy.sql.expression.text(ts)  # $ constructedSql=ts
+    t6 = sqlalchemy.sql.expression.text(text=ts)  # $ constructedSql=ts
+    t7 = sqlalchemy.sql.expression.TextClause(ts)  # $ constructedSql=ts
+    t8 = sqlalchemy.sql.expression.TextClause(text=ts)  # $ constructedSql=ts
 
     # Since we flag user-input to a TextClause with its' own query, we don't want to
     # have a taint-step for it as that would lead to us also giving an alert for normal