From 70491ce3a09356a8b29c8a61c3e24bada4e9e073 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Tue, 10 Oct 2023 14:44:36 +0000
Subject: [PATCH] Python: Add aioch model

---
 python/ql/lib/semmle/python/frameworks/Aioch.qll | 3 +++
 python/ql/src/meta/ClassHierarchy/Find.ql        | 7 +++++++
 2 files changed, 10 insertions(+)

diff --git a/python/ql/lib/semmle/python/frameworks/Aioch.qll b/python/ql/lib/semmle/python/frameworks/Aioch.qll
index ac22bef6dde..3ac9994f113 100644
--- a/python/ql/lib/semmle/python/frameworks/Aioch.qll
+++ b/python/ql/lib/semmle/python/frameworks/Aioch.qll
@@ -9,6 +9,7 @@ private import python
 private import semmle.python.Concepts
 private import semmle.python.ApiGraphs
 private import semmle.python.frameworks.ClickhouseDriver
+private import semmle.python.frameworks.data.ModelsAsData
 
 /**
  * INTERNAL: Do not use.
@@ -24,6 +25,8 @@ module Aioch {
     /** Gets a reference to the `aioch.Client` class or any subclass. */
     API::Node subclassRef() {
       result = API::moduleImport("aioch").getMember("Client").getASubclass*()
+      or
+      result = ModelOutput::getATypeNode("aioch.Client~Subclass").getASubclass*()
     }
 
     /** Gets a reference to an instance of `clickhouse_driver.Client` or any subclass. */
diff --git a/python/ql/src/meta/ClassHierarchy/Find.ql b/python/ql/src/meta/ClassHierarchy/Find.ql
index fed4dcee3de..4b9ce1fdbb8 100644
--- a/python/ql/src/meta/ClassHierarchy/Find.ql
+++ b/python/ql/src/meta/ClassHierarchy/Find.ql
@@ -29,6 +29,7 @@ private import semmle.python.frameworks.Tornado
 private import semmle.python.frameworks.Urllib3
 private import semmle.python.frameworks.Pydantic
 private import semmle.python.frameworks.Peewee
+private import semmle.python.frameworks.Aioch
 import semmle.python.frameworks.data.internal.ApiGraphModelsExtensions as Extensions
 
 class FlaskViewClasses extends FindSubclassesSpec {
@@ -424,6 +425,12 @@ class PeeweeDatabase extends FindSubclassesSpec {
   override API::Node getAlreadyModeledClass() { result = Peewee::Database::subclassRef() }
 }
 
+class AiochClient extends FindSubclassesSpec {
+  AiochClient() { this = "aioch.Client~Subclass" }
+
+  override API::Node getAlreadyModeledClass() { result = Aioch::Client::subclassRef() }
+}
+
 bindingset[fullyQualified]
 predicate fullyQualifiedToYamlFormat(string fullyQualified, string type2, string path) {
   exists(int firstDot | firstDot = fullyQualified.indexOf(".", 0, 0) |