From 16f0fe5b350106d5cea1a178189cf310fd5eefbd Mon Sep 17 00:00:00 2001
From: Marshall Hallenbeck <marshall.hallenbeck@gmail.com>
Date: Wed, 8 Mar 2023 13:10:32 -0500
Subject: [PATCH] feat(mssqldb): update get_computers to be the same as smbdb

---
 cme/protocols/mssql/database.py | 39 ++++++++++++++++-----------------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/cme/protocols/mssql/database.py b/cme/protocols/mssql/database.py
index 77ff3b80..6911dc39 100755
--- a/cme/protocols/mssql/database.py
+++ b/cme/protocols/mssql/database.py
@@ -344,34 +344,33 @@ class database:
         """
         Return hosts from the database.
         """
+        q = select(self.ComputersTable)
+
         # if we're returning a single host by ID
         if self.is_computer_valid(filter_term):
-            results = self.conn.query(self.computers_table).filter(
-                self.computers_table.c.id == filter_term
-            ).first()
+            q.filter(
+                self.ComputersTable.c.id == filter_term
+            )
+            results = asyncio.run(self.conn.execute(q)).first()
+            # all() returns a list, so we keep the return format the same so consumers don't have to guess
+            return [results]
         # if we're filtering by domain controllers
         elif filter_term == 'dc':
+            q.filter(
+                self.ComputersTable.c.dc == 1
+            )
             if domain:
-                results = self.conn.query(self.computers_table).filter(
-                    self.computers_table.c.dc == 1,
-                    func.lower(self.computers_table.c.domain) == func.lower(domain)
-                ).all()
-            else:
-                results = self.conn.query(self.computers_table).filter(
-                    self.computers_table.c.dc == 1
-                ).all()
+                q.filter(
+                    func.lower(self.ComputersTable.c.domain) == func.lower(domain)
+                )
         # if we're filtering by ip/hostname
         elif filter_term and filter_term != "":
-            results = self.conn.query(self.computers_table).filter((
-                    func.lower(self.computers_table.c.ip).like(func.lower(f"%{filter_term}%")) |
-                    func.lower(self.computers_table.c.hostname).like(func.lower(f"%{filter_term}"))
-            )).all()
-        # otherwise return all computers
-        else:
-            results = self.conn.query(self.computers_table).all()
+            q.filter(
+                func.lower(self.ComputersTable.c.ip).like(func.lower(f"%{filter_term}%")) |
+                func.lower(self.ComputersTable.c.hostname).like(func.lower(f"%{filter_term}"))
+            )
 
-        self.conn.commit()
-        self.conn.close()
+        results = asyncio.run(self.conn.execute(q)).all()
         return results
 
     def clear_database(self):