Skip to content

Commit 9975fac

Browse files
RasmusWLRasmusWL
committed
Python: Make asyncio version of PEP249 modeling library
so it's also easy to modeling asyncio libraries Also ports aiomysql/aiopg to use this new modeling
1 parent 2b4ebf7 commit 9975fac

File tree

3 files changed

+111
-75
lines changed

3 files changed

+111
-75
lines changed

‎python/ql/lib/semmle/python/frameworks/Aiomysql.qll

Lines changed: 16 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,10 @@ private import python
99
private import semmle.python.dataflow.new.DataFlow
1010
private import semmle.python.Concepts
1111
private import semmle.python.ApiGraphs
12+
private import semmle.python.frameworks.PEP249
1213

1314
/** Provides models for the `aiomysql` PyPI package. */
1415
private module Aiomysql {
15-
private import semmle.python.internal.Awaited
16-
1716
/**
1817
* Gets a `ConnectionPool` that is created when the result of `aiomysql.create_pool()` is awaited.
1918
* See https://aiomysql.readthedocs.io/en/stable/pool.html
@@ -23,49 +22,29 @@ private module Aiomysql {
2322
}
2423

2524
/**
26-
* Gets a `Connection` that is created when
25+
* A Connection that is created when
2726
* - the result of `aiomysql.connect()` is awaited.
2827
* - the result of calling `acquire` on a `ConnectionPool` is awaited.
29-
* See https://aiomysql.readthedocs.io/en/stable/connection.html#connection
28+
* See
29+
* - https://aiomysql.readthedocs.io/en/stable/connection.html#connection
30+
* - https://aiomysql.readthedocs.io/en/stable/pool.html#Pool.acquire
3031
*/
31-
API::Node connection() {
32-
result = API::moduleImport("aiomysql").getMember("connect").getReturn().getAwaited()
33-
or
34-
result = connectionPool().getMember("acquire").getReturn().getAwaited()
32+
class AiomysqlConnection extends PEP249::AsyncDatabaseConnection {
33+
AiomysqlConnection() {
34+
this = API::moduleImport("aiomysql").getMember("connect").getReturn().getAwaited()
35+
or
36+
this = connectionPool().getMember("acquire").getReturn().getAwaited()
37+
}
3538
}
3639

3740
/**
38-
* Gets a `Cursor` that is created when
41+
* An additional cursor, that is created when
3942
* - the result of calling `cursor` on a `ConnectionPool` is awaited.
40-
* - the result of calling `cursor` on a `Connection` is awaited.
41-
* See https://aiomysql.readthedocs.io/en/stable/cursors.html
43+
* See
44+
* - https://aiomysql.readthedocs.io/en/stable/pool.html##Pool.cursor
4245
*/
43-
API::Node cursor() {
44-
result = connectionPool().getMember("cursor").getReturn().getAwaited()
45-
or
46-
result = connection().getMember("cursor").getReturn().getAwaited()
47-
}
48-
49-
/**
50-
* A query. Calling `execute` on a `Cursor` constructs a query.
51-
* See https://aiomysql.readthedocs.io/en/stable/cursors.html#Cursor.execute
52-
*/
53-
class CursorExecuteCall extends SqlConstruction::Range, API::CallNode {
54-
CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
55-
56-
override DataFlow::Node getSql() { result = this.getParameter(0, "operation").asSink() }
57-
}
58-
59-
/**
60-
* An awaited query. Awaiting the result of calling `execute` executes the query.
61-
* See https://aiomysql.readthedocs.io/en/stable/cursors.html#Cursor.execute
62-
*/
63-
class AwaitedCursorExecuteCall extends SqlExecution::Range {
64-
CursorExecuteCall executeCall;
65-
66-
AwaitedCursorExecuteCall() { this = executeCall.getReturn().getAwaited().asSource() }
67-
68-
override DataFlow::Node getSql() { result = executeCall.getSql() }
46+
class AiomysqlCursor extends PEP249::AsyncDatabaseCursor {
47+
AiomysqlCursor() { this = connectionPool().getMember("cursor").getReturn().getAwaited() }
6948
}
7049

7150
/**

‎python/ql/lib/semmle/python/frameworks/Aiopg.qll

Lines changed: 16 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,10 @@ private import python
99
private import semmle.python.dataflow.new.DataFlow
1010
private import semmle.python.Concepts
1111
private import semmle.python.ApiGraphs
12+
private import semmle.python.frameworks.PEP249
1213

1314
/** Provides models for the `aiopg` PyPI package. */
1415
private module Aiopg {
15-
private import semmle.python.internal.Awaited
16-
1716
/**
1817
* Gets a `ConnectionPool` that is created when the result of `aiopg.create_pool()` is awaited.
1918
* See https://aiopg.readthedocs.io/en/stable/core.html#pool
@@ -23,49 +22,29 @@ private module Aiopg {
2322
}
2423

2524
/**
26-
* Gets a `Connection` that is created when
25+
* A Connection that is created when
2726
* - the result of `aiopg.connect()` is awaited.
2827
* - the result of calling `acquire` on a `ConnectionPool` is awaited.
29-
* See https://aiopg.readthedocs.io/en/stable/core.html#connection
28+
* See
29+
* - https://aiopg.readthedocs.io/en/stable/core.html#connection
30+
* - https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Pool.acquire
3031
*/
31-
API::Node connection() {
32-
result = API::moduleImport("aiopg").getMember("connect").getReturn().getAwaited()
33-
or
34-
result = connectionPool().getMember("acquire").getReturn().getAwaited()
32+
class AiopgConnection extends PEP249::AsyncDatabaseConnection {
33+
AiopgConnection() {
34+
this = API::moduleImport("aiopg").getMember("connect").getReturn().getAwaited()
35+
or
36+
this = connectionPool().getMember("acquire").getReturn().getAwaited()
37+
}
3538
}
3639

3740
/**
38-
* Gets a `Cursor` that is created when
41+
* An additional cursor, that is created when
3942
* - the result of calling `cursor` on a `ConnectionPool` is awaited.
40-
* - the result of calling `cursor` on a `Connection` is awaited.
41-
* See https://aiopg.readthedocs.io/en/stable/core.html#cursor
42-
*/
43-
API::Node cursor() {
44-
result = connectionPool().getMember("cursor").getReturn().getAwaited()
45-
or
46-
result = connection().getMember("cursor").getReturn().getAwaited()
47-
}
48-
49-
/**
50-
* A query. Calling `execute` on a `Cursor` constructs a query.
51-
* See https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Cursor.execute
52-
*/
53-
class CursorExecuteCall extends SqlConstruction::Range, API::CallNode {
54-
CursorExecuteCall() { this = cursor().getMember("execute").getACall() }
55-
56-
override DataFlow::Node getSql() { result = this.getParameter(0, "operation").asSink() }
57-
}
58-
59-
/**
60-
* An awaited query. Awaiting the result of calling `execute` executes the query.
61-
* See https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Cursor.execute
43+
* See
44+
* - https://aiopg.readthedocs.io/en/stable/core.html#aiopg.Pool.cursor
6245
*/
63-
class AwaitedCursorExecuteCall extends SqlExecution::Range {
64-
CursorExecuteCall execute;
65-
66-
AwaitedCursorExecuteCall() { this = execute.getReturn().getAwaited().asSource() }
67-
68-
override DataFlow::Node getSql() { result = execute.getSql() }
46+
class AiopgCursor extends PEP249::AsyncDatabaseCursor {
47+
AiopgCursor() { this = connectionPool().getMember("cursor").getReturn().getAwaited() }
6948
}
7049

7150
/**

‎python/ql/lib/semmle/python/frameworks/PEP249.qll

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ module PEP249 {
5252
result in ["sql", "statement", "operation", "query", "query_string", "sql_script"]
5353
}
5454

55+
private string getExecuteMethodName() { result in ["execute", "executemany", "executescript"] }
56+
5557
/**
5658
* A call to an execute method on a database cursor or a connection, such as `execute`
5759
* or `executemany`.
@@ -68,7 +70,7 @@ module PEP249 {
6870
exists(API::Node start |
6971
start instanceof DatabaseCursor or start instanceof DatabaseConnection
7072
|
71-
this = start.getMember(["execute", "executemany", "executescript"]).getACall()
73+
this = start.getMember(getExecuteMethodName()).getACall()
7274
)
7375
}
7476

@@ -77,6 +79,82 @@ module PEP249 {
7779
}
7880
}
7981

82+
// ---------------------------------------------------------------------------
83+
// asyncio implementations
84+
// ---------------------------------------------------------------------------
85+
//
86+
// we differentiate between normal and asyncio implementations, since we model the
87+
// `execute` call differently -- as a SqlExecution vs SqlConstruction, since the SQL
88+
// is only executed in asyncio after being awaited (which might happen in something
89+
// like `asyncio.gather`)
90+
/**
91+
* An API graph node representing a module that implements PEP 249 using asyncio.
92+
*/
93+
abstract class AsyncPEP249ModuleApiNode extends API::Node {
94+
/** Gets a string representation of this element. */
95+
override string toString() { result = this.(API::Node).toString() }
96+
}
97+
98+
/**
99+
* An API graph node representing a asyncio database connection (after being awaited).
100+
*/
101+
abstract class AsyncDatabaseConnection extends API::Node {
102+
/** Gets a string representation of this element. */
103+
override string toString() { result = this.(API::Node).toString() }
104+
}
105+
106+
private class DefaultAsyncDatabaseConnection extends AsyncDatabaseConnection {
107+
DefaultAsyncDatabaseConnection() {
108+
this = any(AsyncPEP249ModuleApiNode mod).getMember("connect").getReturn().getAwaited()
109+
}
110+
}
111+
112+
/**
113+
* An API graph node representing a asyncio database cursor (after being awaited).
114+
*/
115+
abstract class AsyncDatabaseCursor extends API::Node {
116+
/** Gets a string representation of this element. */
117+
override string toString() { result = this.(API::Node).toString() }
118+
}
119+
120+
private class DefaultAsyncDatabaseCursor extends AsyncDatabaseCursor {
121+
DefaultAsyncDatabaseCursor() {
122+
this = any(AsyncDatabaseConnection conn).getMember("cursor").getReturn().getAwaited()
123+
}
124+
}
125+
126+
/**
127+
* A call to an execute method on an asyncio database cursor or an asyncio connection,
128+
* such as `execute` or `executemany`.
129+
*
130+
* (This is not an SqlExecution, since that only happens when the coroutine is
131+
* awaited)
132+
*
133+
* See ExecuteMethodCall for more details.
134+
*/
135+
private class AsyncExecuteMethodCall extends SqlConstruction::Range, API::CallNode {
136+
AsyncExecuteMethodCall() {
137+
exists(API::Node start |
138+
start instanceof AsyncDatabaseCursor or start instanceof AsyncDatabaseConnection
139+
|
140+
this = start.getMember(getExecuteMethodName()).getACall()
141+
)
142+
}
143+
144+
override DataFlow::Node getSql() {
145+
result in [this.getArg(0), this.getArgByName(getSqlKwargName()),]
146+
}
147+
}
148+
149+
/** Actual execution of the AsyncExecuteMethodCall coroutine. */
150+
private class AwaitedAsyncExecuteMethodCall extends SqlExecution::Range {
151+
AsyncExecuteMethodCall execute;
152+
153+
AwaitedAsyncExecuteMethodCall() { this = execute.getReturn().getAwaited().asSource() }
154+
155+
override DataFlow::Node getSql() { result = execute.getSql() }
156+
}
157+
80158
// ---------------------------------------------------------------------------
81159
// old impl
82160
// ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)