Merge main

Author: TristanDamron

src/db_connector.py

@@ -35,6 +35,29 @@ def get_tables(self) -> list[str]:
         return tables
 
 
+    def get_groups(self) -> list[dict]:
+        groups = []
+        tables = self.get_tables()
+        for table in tables:
+            group = {}
+            split_table = table.split(".")
+            fields = self.get_fields_for_table(split_table[0], split_table[1])
+            for field in fields:
+                if "__fgid" in field["name"]:
+                    try:
+                        group[split_table[1]]
+                    except KeyError:
+                        group[split_table[1]] = {}
+                        group[split_table[1]]["layers"] = []
+                    group[split_table[1]]["layers"].append(field["name"].split("__fgid")[0])
+
+            if group:
+                groups.append(group)
+
+        return groups
+
+
+
     def get_fields_for_table(self, schema: str, table: str) -> list[dict]:
         fields = []
         self.cursor.execute(f"""SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = '{schema}' AND table_name = '{table}' AND column_name != 'geom';""")

src/feature_group.py

@@ -0,0 +1,42 @@
+from tabor_layer import TaborLayer
+
+
+class FeatureGroup(object):
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.schema = ""
+        self.owner = ""
+        self.layers = []
+
+
+    def add_layer(self, layer: TaborLayer) -> list[TaborLayer]:
+        if not layer.get_pk_field()[0]:
+            raise Exception(f"{layer.name} has no primary key but it is in group {self.name}")
+
+        if not self.schema:
+            self.schema = layer.schema
+
+        if not self.owner:
+            self.owner = layer.owner
+
+        if layer.schema != self.schema:
+            raise Exception(f"{layer.name} is in the wrong schema, expected {self.schema} but got {layer.schema}")
+
+        if layer.owner != self.owner:
+            raise Exception(f"{layer.name} has the wrong owner, expected {self.owner} but got {layer.owner}")
+
+        self.layers.append(layer)
+        return self.layers
+
+
+    def gather_foreign_keys(self) -> str:
+        foreign_keys = []
+
+        for layer in self.layers:
+            pk_name, pk_type = layer.get_pk_field()
+            foreign_keys.append(f"""{layer.name}__fgid {pk_type} references "{layer.schema}"."{layer.name}"({pk_name})""")
+        return f"({", ".join(foreign_keys)})"
+
+
+    def __str__(self) -> str:
+        return f"""CREATE TABLE IF NOT EXISTS "{self.schema}"."{self.name}" {self.gather_foreign_keys()}; """

src/tabor.py

@@ -17,15 +17,21 @@ def load(file_path: str, db: str, username: str, password: str, host: str, port:
 
         tabor_src = TaborFile(file_path)
         psql = tabor_src.to_psql()
+        layers = psql["layers"]
+        groups = psql["groups"]
 
-        for _,layer in psql.items():
+        for _,layer in layers.items():
             for _,value in layer.items():
                 if type(value) is list:
                     for query in value:
                         db_connector.execute_commit_query(query)
                 else:
                     db_connector.execute_commit_query(value)
 
+        for _,group in groups.items():
+            for _,value in group.items():
+                db_connector.execute_commit_query(value)
+
         print(f"Loaded {file_path} to {db}")
     except:
         raise Exception("Failed to load .tabor file to PostGIS database")
@@ -45,12 +51,14 @@ def write(file_path: str, db: str, username: str, password: str, host: str, port
         data = {}
         db_connector = DBConnector(db, username, password, host, port)
         tables = db_connector.get_tables()
+        groups = db_connector.get_groups()
+        group_names = [list(group)[0] for group in groups]
 
         for table in tables:
             schema = table.split(".")[0]
             table_name = table.split(".")[1]
 
-            if table_name in ignore_tables:
+            if table_name in ignore_tables or table_name in group_names:
                 continue
 
             data[table] = {}
@@ -63,6 +71,11 @@ def write(file_path: str, db: str, username: str, password: str, host: str, port
 
             data[table]["owner"] = username
 
+            for group in groups:
+                for group_name,layers in group.items():
+                    if table_name in layers["layers"]:
+                        data[table]["group"] = group_name
+
         tabor_src = TaborFile(file_path, psql_data=data)
         tabor_src.write()
         print(f"Wrote .tabor file to {file_path}")

src/tabor_file.py

@@ -2,11 +2,14 @@
 from yaml import safe_dump, safe_load
 
 from constraint import Constraint, Trigger
+from feature_group import FeatureGroup
 from tabor_layer import TaborLayer
 from consts import VERSION
 
 
 class TaborFile(object):
+    # TODO: Use hashmaps!
+    groups: list[FeatureGroup] = []
     layers: list[TaborLayer] = []
 
     def __init__(self, path: str, psql_data: dict | None = None) -> None:
@@ -31,7 +34,14 @@ def __init__(self, path: str, psql_data: dict | None = None) -> None:
                         except KeyError:
                             pass
 
-                        self.add_layer(layer["name"], layer["schema"], geometry, layer["owner"], layer["fields"], constraints)
+
+                        group = ""
+                        try:
+                            group = layer["group"]
+                        except KeyError:
+                            pass
+
+                        self.add_layer(layer["name"], layer["schema"], geometry, layer["owner"], layer["fields"], constraints, group)
 
             except FileNotFoundError:
                 raise Exception(f"Failed to read .tabor file at {self.path}, does that path exist?")
@@ -51,24 +61,43 @@ def __init__(self, path: str, psql_data: dict | None = None) -> None:
                 except KeyError:
                     pass
 
-                self.add_layer(table.split(".")[1], table.split(".")[0], geom, values["owner"], values["fields"], derived_constraints)
+                derived_group = None
+                try:
+                    derived_group = values["group"]
+                except KeyError:
+                    pass
+                self.add_layer(table.split(".")[1], table.split(".")[0], geom, values["owner"], values["fields"], derived_constraints, derived_group)
 
 
-    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))
+    def add_layer(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: list[dict], group: str | None) -> TaborLayer:
+        layer = TaborLayer(name, schema, geometry, owner, fields, constraints, group)
+        if group:
+            feature_group = self.add_or_get_group(group)
+            feature_group.add_layer(layer)
+        self.layers.append(layer)
         return self.layers[len(self.layers) - 1]
 
 
+    def add_or_get_group(self, name: str) -> FeatureGroup:
+        for group in self.groups:
+            if group.name == name:
+                return group
+
+        self.groups.append(FeatureGroup(name))
+        return self.groups[len(self.groups) - 1]
+
+
     def to_psql(self) -> dict:
         result = {}
+        result["layers"] = {}
         for layer in self.layers:
-            result[layer.name] = {}
+            result["layers"][layer.name] = {}
             fields = []
             for field in layer.fields:
                 fields.append(field.as_psql())
 
-            if layer.get_pk_field():
-                pk_query = f""", PRIMARY KEY ({layer.get_pk_field()})"""
+            if layer.get_pk_field()[0]:
+                pk_query = f""", PRIMARY KEY ({layer.get_pk_field()[0]})"""
             else:
                 pk_query = ""
 
@@ -77,15 +106,21 @@ def to_psql(self) -> dict:
             else:
                 geom_query = ""
 
-            result[layer.name]["schema"] = f"""CREATE TABLE IF NOT EXISTS "{layer.schema}"."{layer.name}" ({", ".join(fields)}{geom_query}{pk_query});"""
-            result[layer.name]["owner"] = f"""ALTER TABLE "{layer.schema}"."{layer.name}" OWNER TO {layer.owner};"""
+            result["layers"][layer.name]["schema"] = f"""CREATE TABLE IF NOT EXISTS "{layer.schema}"."{layer.name}" ({", ".join(fields)}{geom_query}{pk_query});"""
+            result["layers"][layer.name]["owner"] = f"""ALTER TABLE "{layer.schema}"."{layer.name}" OWNER TO {layer.owner};"""
 
-            result[layer.name]["constraints"] = []
+            result["layers"][layer.name]["constraints"] = []
             for constraint in layer.constraints:
-                result[layer.name]["constraints"].append(str(constraint))
+                result["layers"][layer.name]["constraints"].append(str(constraint))
 
             if layer.geometry:
-                result[layer.name]["geometry"] = f"""ALTER TABLE "{layer.schema}"."{layer.name}" ALTER COLUMN geom TYPE Geometry({layer.derive_geometry_type()});"""
+                result["layers"][layer.name]["geometry"] = f"""ALTER TABLE "{layer.schema}"."{layer.name}" ALTER COLUMN geom TYPE Geometry({layer.derive_geometry_type()});"""
+
+        result["groups"] = {}
+        for group in self.groups:
+            result["groups"][group.name] = {}
+            result["groups"][group.name]["schema"] = str(group)
+            result["groups"][group.name]["owner"] = f"""ALTER TABLE "{group.schema}"."{group.name}" OWNER TO {group.owner};"""
 
         return result
 

src/tabor_layer.py

@@ -1,10 +1,11 @@
 from geometry_type import GeometryType
 from constraint import Constraint
+from tabor_field_type import TaborFieldType
 from tabor_field import TaborField
 
 
 class TaborLayer(object):
-    def __init__(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: list[dict]) -> None:
+    def __init__(self, name: str, schema: str, geometry: str | None, owner: str, fields: dict, constraints: list[dict], group: str | None) -> None:
         self.name = name
         self.schema = schema
 
@@ -13,6 +14,11 @@ def __init__(self, name: str, schema: str, geometry: str | None, owner: str, fie
         else:
             self.geometry = None
 
+        if group:
+            self.group = group
+        else:
+            self.group = None
+
         self.constraints: list[Constraint] = []
         if constraints:
             for constraint in constraints:
@@ -37,11 +43,11 @@ def add_field(self, name: str, type: str, pk: bool) -> None:
         self.fields.append(TaborField(name, type, pk))
 
 
-    def get_pk_field(self) -> str:
+    def get_pk_field(self) -> tuple[str, TaborFieldType]:
         for field in self.fields:
             if field.pk:
-                return field.name
-        return ""
+                return (field.name, field.type)
+        return ("", TaborFieldType(""))
 
 
     def as_dict(self) -> dict:
@@ -58,6 +64,9 @@ def as_dict(self) -> dict:
         if self.geometry:
             var_dict["geometry"] = str(self.geometry)
 
+        if self.group:
+            var_dict["group"] = self.group
+
         return var_dict
 
 
@@ -72,3 +81,7 @@ def derive_geometry_type(self) -> str:
             "multi polygon": "Multipolygon"
         }
         return postgis_geometry_types[self.geometry.type]
+
+
+    def __str__(self) -> str:
+        return f"{self.schema}.{self.name}"