improve version string parsing for Postgres

Author: dantownsend

piccolo/engine/postgres.py

@@ -227,6 +227,19 @@ def __init__(self, config: t.Dict[str, t.Any]) -> None:
         )
         super().__init__()
 
+    def _parse_raw_version_string(self, version_string: str) -> float:
+        """
+        The format of the version string isn't always consistent. Sometimes
+        it's just the version number e.g. '9.6.18', and sometimes
+        it contains specific build information e.g.
+        '12.4 (Ubuntu 12.4-0ubuntu0.20.04.1)'. Just extract the major and
+        minor version numbers.
+        """
+        version_segment = version_string.split(" ")[0]
+        major, minor = version_segment.split(".")[:2]
+        version = float(f"{major}.{minor}")
+        return version
+
     def get_version(self) -> float:
         """
         Returns the version of Postgres being run.
@@ -248,10 +261,10 @@ def get_version(self) -> float:
             print(exception)
             return 0.0
         else:
-            server_version = response[0]["server_version"]
-            major, minor = server_version.split(".")[:2]
-            version = float(f"{major}.{minor}")
-            return version
+            version_string = response[0]["server_version"]
+            return self._parse_raw_version_string(
+                version_string=version_string
+            )
 
     def prep_database(self):
         loop = asyncio.new_event_loop()

tests/engine/test_version_parsing.py

@@ -0,0 +1,29 @@
+from unittest import TestCase
+
+from piccolo.engine.postgres import PostgresEngine
+from ..base import postgres_only
+
+
+@postgres_only
+class TestVersionParsing(TestCase):
+    def test_version_parsing(self):
+        """
+        Make sure the version number can correctly be parsed from a range
+        of known formats.
+        """
+        engine = PostgresEngine(config={})
+
+        self.assertEqual(
+            engine._parse_raw_version_string(version_string="9.4"), 9.4
+        )
+
+        self.assertEqual(
+            engine._parse_raw_version_string(version_string="9.4.1"), 9.4
+        )
+
+        self.assertEqual(
+            engine._parse_raw_version_string(
+                version_string="12.4 (Ubuntu 12.4-0ubuntu0.20.04.1)"
+            ),
+            12.4,
+        )