created method for retrieving column type from postgres

Author: dantownsend

tests/base.py

@@ -50,11 +50,28 @@ def run_sync(self, query):
         _Table = type("_Table", (Table,), {})
         return _Table.raw(query).run_sync()
 
-    def table_exists(self, tablename: str):
+    def table_exists(self, tablename: str) -> bool:
         _Table: t.Type[Table] = type(tablename.upper(), (Table,), {})
         _Table._meta.tablename = tablename
         return _Table.table_exists().run_sync()
 
+    def get_postgres_column_type_str(
+        self, tablename: str, column_name: str
+    ) -> str:
+        """
+        Fetches the column type as a string, from the database.
+        """
+        query = """
+            SELECT data_type FROM information_schema.columns
+            WHERE table_name = '{tablename}'
+            AND table_catalog = 'piccolo'
+            AND column_name = '{column_name}'
+        """.format(
+            tablename=tablename, column_name=column_name
+        )
+        response = self.run_sync(query)
+        return response[0]["data_type"].upper()
+
     def create_tables(self):
         if ENGINE.engine_type == "postgres":
             self.run_sync(

tests/table/test_alter.py

@@ -151,15 +151,12 @@ def test_integer_to_bigint(self):
         )
         alter_query.run_sync()
 
-        query = """
-            SELECT data_type FROM information_schema.columns
-            WHERE table_name = 'band'
-            AND table_catalog = 'piccolo'
-            AND column_name = 'popularity'
-            """
-
-        response = Band.raw(query).run_sync()
-        self.assertEqual(response[0]["data_type"].upper(), "BIGINT")
+        self.assertEqual(
+            self.get_postgres_column_type_str(
+                tablename="band", column_name="popularity"
+            ),
+            "BIGINT",
+        )
 
         popularity = (
             Band.select(Band.popularity).first().run_sync()["popularity"]
@@ -177,15 +174,12 @@ def test_integer_to_varchar(self):
         )
         alter_query.run_sync()
 
-        query = """
-            SELECT data_type FROM information_schema.columns
-            WHERE table_name = 'band'
-            AND table_catalog = 'piccolo'
-            AND column_name = 'popularity'
-            """
-
-        response = Band.raw(query).run_sync()
-        self.assertEqual(response[0]["data_type"].upper(), "CHARACTER VARYING")
+        self.assertEqual(
+            self.get_postgres_column_type_str(
+                tablename="band", column_name="popularity"
+            ),
+            "CHARACTER VARYING",
+        )
 
         popularity = (
             Band.select(Band.popularity).first().run_sync()["popularity"]