From 55dc75615a2093cbffaa1ad5f01f7611dafda867 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Tue, 10 Oct 2023 14:41:45 +0000
Subject: [PATCH] Python: Add Peewee model

---
 python/ql/lib/semmle/python/frameworks/Peewee.qll | 9 +++++++--
 python/ql/src/meta/ClassHierarchy/Find.ql         | 7 +++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/python/ql/lib/semmle/python/frameworks/Peewee.qll b/python/ql/lib/semmle/python/frameworks/Peewee.qll
index f923b463b3e..eedc328d98d 100644
--- a/python/ql/lib/semmle/python/frameworks/Peewee.qll
+++ b/python/ql/lib/semmle/python/frameworks/Peewee.qll
@@ -11,14 +11,17 @@ private import semmle.python.dataflow.new.TaintTracking
 private import semmle.python.Concepts
 private import semmle.python.ApiGraphs
 private import semmle.python.frameworks.PEP249
+private import semmle.python.frameworks.data.ModelsAsData
 
 /**
+ * INTERNAL: Do not use.
+ *
  * Provides models for the `peewee` PyPI package.
  * See
  * - https://pypi.org/project/peewee/
  * - https://docs.peewee-orm.com/en/latest/index.html
  */
-private module Peewee {
+module Peewee {
   /** Provides models for the `peewee.Database` class and subclasses. */
   module Database {
     /** Gets a reference to the `peewee.Database` class or any subclass. */
@@ -31,7 +34,7 @@ private module Peewee {
             .getMember(["SqliteDatabase", "MySQLDatabase", "PostgresqlDatabase"])
             .getASubclass*()
       or
-      // Ohter known subclasses, semi auto generated by using
+      // Other known subclasses, semi auto generated by using
       // ```ql
       // class DBClass extends Class, SelfRefMixin {
       //   DBClass() {
@@ -153,6 +156,8 @@ private module Peewee {
             .getMember("sqliteq")
             .getMember("SqliteQueueDatabase")
             .getASubclass*()
+      or
+      result = ModelOutput::getATypeNode("peewee.Database~Subclass").getASubclass*()
     }
 
     /** Gets a reference to an instance of `peewee.Database` or any subclass. */
diff --git a/python/ql/src/meta/ClassHierarchy/Find.ql b/python/ql/src/meta/ClassHierarchy/Find.ql
index 9cebc36f6d0..fed4dcee3de 100644
--- a/python/ql/src/meta/ClassHierarchy/Find.ql
+++ b/python/ql/src/meta/ClassHierarchy/Find.ql
@@ -28,6 +28,7 @@ private import semmle.python.frameworks.SqlAlchemy
 private import semmle.python.frameworks.Tornado
 private import semmle.python.frameworks.Urllib3
 private import semmle.python.frameworks.Pydantic
+private import semmle.python.frameworks.Peewee
 import semmle.python.frameworks.data.internal.ApiGraphModelsExtensions as Extensions
 
 class FlaskViewClasses extends FindSubclassesSpec {
@@ -417,6 +418,12 @@ class PydanticBaseModel extends FindSubclassesSpec {
   override API::Node getAlreadyModeledClass() { result = Pydantic::BaseModel::subclassRef() }
 }
 
+class PeeweeDatabase extends FindSubclassesSpec {
+  PeeweeDatabase() { this = "peewee.Database~Subclass" }
+
+  override API::Node getAlreadyModeledClass() { result = Peewee::Database::subclassRef() }
+}
+
 bindingset[fullyQualified]
 predicate fullyQualifiedToYamlFormat(string fullyQualified, string type2, string path) {
   exists(int firstDot | firstDot = fullyQualified.indexOf(".", 0, 0) |