Merge pull request #15 from geoCML/write-constraints

Write constraints from Postgres database to Tabor file

Author: TristanDamron

src/constraint.py

@@ -1,5 +1,6 @@
 from tabor_constraint_type import TaborConstraintType
 
+
 class Constraint(object):
     def __init__(self, constraint: dict, layer: str) -> None:
         self.constraint = constraint
@@ -17,6 +18,10 @@ def on(self, other_layer: str) -> str:
         return f"""CREATE OR REPLACE FUNCTION {self.layer}_on_{other_layer}() RETURNS trigger AS $$ DECLARE overlap boolean; BEGIN SELECT Count(*) INTO overlap FROM {other_layer} WHERE ST_Contains({other_layer}.geom, NEW.geom); IF NOT overlap THEN RAISE EXCEPTION '{self.layer} is not on {other_layer}'; END IF; RETURN NEW; END; $$ LANGUAGE 'plpgsql'; CREATE CONSTRAINT TRIGGER {self.layer}_on_{other_layer} AFTER INSERT OR UPDATE ON {self.layer} FOR EACH ROW EXECUTE FUNCTION {self.layer}_on_{other_layer}();"""
 
 
+    def as_dict(self) -> dict:
+        return self.constraint
+
+
     def __str__(self) -> str:
         if self.constraint_type.type == "on":
             try:
@@ -25,3 +30,23 @@ def __str__(self) -> str:
                 raise Exception("Constraint 'on' needs a relative layer value")
             return self.on(layer)
         return ""
+
+
+class Trigger(object):
+    """
+    Triggers are strings that can be used to derive Constraints
+
+    e.g. trees_on_grass -> { "name": "on", "layer": "grass" }
+    """
+    def __init__(self, name: str) -> None:
+        self.constraint = self.derive_constraint_from_name(name)
+
+
+    def derive_constraint_from_name(self, name) -> Constraint:
+        if "_on_" in name:
+            return Constraint({
+                "name": "on",
+                "layer": name.split("_on_")[1]
+            }, name.split("_on_")[0])
+
+        raise Exception(f"Cannot derive a constraint from trigger '{name}'")

src/db_connector.py

@@ -62,3 +62,8 @@ def get_geometry_type_for_table(self, schema: str, table: str) -> str:
             return ""
 
         return result[0]
+
+
+    def get_triggers_for_table(self, table: str) -> set[str]:
+        self.cursor.execute(f"""SELECT trigger_name FROM information_schema.triggers WHERE event_object_table = '{table}'""")
+        return {trigger[0] for trigger in self.cursor.fetchall()}

src/tabor.py

@@ -31,7 +31,7 @@ def write(file_path: str, db: str, username: str, password: str, host: str, port
 
             data[table] = {}
             data[table]["fields"] = db_connector.get_fields_for_table(schema, table_name)
-
+            data[table]["constraints"] = db_connector.get_triggers_for_table(table_name)
 
             geom_type = db_connector.get_geometry_type_for_table(schema, table_name)
             if geom_type:

src/tabor_file.py

@@ -1,6 +1,7 @@
 import re
 from yaml import safe_dump, safe_load
 
+from constraint import Constraint, Trigger
 from tabor_layer import TaborLayer
 from consts import VERSION
 
@@ -24,9 +25,9 @@ def __init__(self, path: str, psql_data: dict | None = None) -> None:
                         except KeyError:
                             geometry = None
 
-                        constraints: dict = {}
+                        constraints: list[dict] = []
                         try:
-                            constraints = layer["constraints"]
+                            constraints: list[dict] = layer["constraints"]
                         except KeyError:
                             pass
 
@@ -43,17 +44,17 @@ def __init__(self, path: str, psql_data: dict | None = None) -> None:
                 except KeyError:
                     geom = None
 
-
-                constraints: dict = {}
+                derived_constraints: list[dict] = []
                 try:
-                    constraints = values["constraints"]
+                    for trigger in values["constraints"]:
+                        derived_constraints.append(Trigger(trigger).constraint.constraint)
                 except KeyError:
                     pass
 
-                self.add_layer(table.split(".")[1], table.split(".")[0], geom, values["owner"], values["fields"], constraints)
+                self.add_layer(table.split(".")[1], table.split(".")[0], geom, values["owner"], values["fields"], derived_constraints)
 
 
-    def add_layer(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: dict) -> TaborLayer:
+    def add_layer(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: list[dict]) -> TaborLayer:
         self.layers.append(TaborLayer(name, schema, geometry, owner, fields, constraints))
         return self.layers[len(self.layers) - 1]
 

src/tabor_layer.py

@@ -4,7 +4,7 @@
 
 
 class TaborLayer(object):
-    def __init__(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: dict) -> None:
+    def __init__(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: list[dict]) -> None:
         self.name = name
         self.schema = schema
 
@@ -43,14 +43,18 @@ def get_pk_field(self) -> str:
                 return field.name
         return ""
 
+
     def as_dict(self) -> dict:
         var_dict = {
             "name": self.name,
             "schema": self.schema,
             "owner": self.owner,
-            "fields": [field.as_dict() for field in self.fields]
+            "fields": [field.as_dict() for field in self.fields],
         }
 
+        if self.constraints:
+            var_dict["constraints"] = [constraint.as_dict() for constraint in self.constraints]
+
         if self.geometry:
             var_dict["geometry"] = str(self.geometry)