support unix_socket

Author: paulomach

src/charm.py

@@ -99,15 +99,20 @@ def _database_config(self):
         if None in [username, password, endpoints]:
             return {}
 
-        [host, port] = endpoints.split(":")
-
-        return {
+        config = {
             "user": username,
             "password": password,
-            "host": host,
-            "port": port,
             "database": DATABASE_NAME,
         }
+        if endpoints.startswith("file://"):
+            config["unix_socket"] = endpoints[7:]
+            config["port"] = "socket"
+        else:
+            host, port = endpoints.split(":")
+            config["host"] = host
+            config["port"] = port
+
+        return config
 
     # ==============
     # Helpers
@@ -127,7 +132,7 @@ def _start_continuous_writes(self, starting_number: int) -> None:
                 "src/continuous_writes.py",
                 self._database_config["user"],
                 self._database_config["password"],
-                self._database_config["host"],
+                self._database_config.get("host", self._database_config["unix_socket"]),
                 self._database_config["port"],
                 self._database_config["database"],
                 CONTINUOUS_WRITE_TABLE_NAME,

src/continuous_writes.py

@@ -55,12 +55,15 @@ def main():
     database_config = {
         "user": username,
         "password": password,
-        "host": host,
-        "port": port,
         "database": database,
         "use_pure": True,
         "connection_timeout": 5,
     }
+    if port == "socket":
+        database_config["unix_socket"] = host
+    else:
+        database_config["port"] = port
+        database_config["host"] = host
 
     continuous_writes(database_config, table_name, int(starting_number))