fix: initial work on primary keys

Author: rustyconover

requirements-dev.lock

@@ -86,7 +86,7 @@ python-levenshtein==0.27.1
     # via query-farm-flight-server
 query-farm-duckdb-json-serialization==0.1.2
     # via query-farm-airport-test-server
-query-farm-flight-server==0.1.8
+query-farm-flight-server==0.1.10
     # via query-farm-airport-test-server
 rapidfuzz==3.13.0
     # via levenshtein

src/query_farm_airport_test_server/database_impl.py

@@ -161,6 +161,11 @@ class TableInfo:
     # To enable version history keep track of tables.
     table_versions: list[pa.Table] = field(default_factory=list)
 
+    primary_key_columns: list[str] = field(default_factory=list)
+    unique_columns: list[str] = field(default_factory=list)
+    multi_key_primary_keys: list[str] = field(default_factory=list)
+    extra_constraints: list[str] = field(default_factory=list)
+
     # the next row id to assign.
     row_id_counter: int = 0
 
@@ -219,6 +224,10 @@ def serialize(self) -> dict[str, Any]:
         return {
             "table_versions": [serialize_table_data(table) for table in self.table_versions],
             "row_id_counter": self.row_id_counter,
+            "primary_key_columns": self.primary_key_columns,
+            "unique_columns": self.unique_columns,
+            "multi_key_primary_keys": self.multi_key_primary_keys,
+            "extra_constraints": self.extra_constraints,
         }
 
     def deserialize(self, data: dict[str, Any]) -> "TableInfo":
@@ -228,6 +237,10 @@ def deserialize(self, data: dict[str, Any]) -> "TableInfo":
         self.table_versions = [deserialize_table_data(table) for table in data["table_versions"]]
         self.row_id_counter = data["row_id_counter"]
         self.endpoint_generator = None
+        self.primary_key_columns = data.get("primary_key_columns", [])
+        self.unique_columns = data.get("unique_columns", [])
+        self.multi_key_primary_keys = data.get("multi_key_primary_keys", [])
+        self.extra_constraints = data.get("extra_constraints", [])
         return self
 
 
@@ -374,7 +387,11 @@ def by_name(self, name: str) -> SchemaCollection:
 
     def serialize(self) -> dict[str, Any]:
         return {
-            "schemas": {name: schema.serialize() for name, schema in self.schemas_by_name.items()},
+            "schemas": {
+                name: schema.serialize()
+                for name, schema in self.schemas_by_name.items()
+                if name not in ("static_data", "remote_data", "utils")
+            },
             "version": self.version,
         }
 

src/query_farm_airport_test_server/server.py

@@ -395,7 +395,14 @@ def action_create_table(
 
             schema_with_row_id = parameters.arrow_schema.append(self.rowid_field)
 
-            table_info = TableInfo([schema_with_row_id.empty_table()], 0)
+            table_info = TableInfo(
+                table_versions=[schema_with_row_id.empty_table()],
+                row_id_counter=0,
+                primary_key_columns=parameters.primary_key_columns,
+                unique_columns=parameters.unique_columns,
+                multi_key_primary_keys=parameters.multi_key_primary_keys,
+                extra_constraints=parameters.extra_constraints,
+            )
 
             schema.tables_by_name[parameters.table_name] = table_info
 
@@ -653,6 +660,38 @@ def exchange_insert(
             for chunk in reader:
                 if chunk.data is not None:
                     new_rows = pa.Table.from_batches([chunk.data])
+
+                    for column_name in table_info.primary_key_columns:
+                        insert_values = new_rows.column(column_name)
+
+                        # Make sure that all elements are unique
+                        value_counts = pc.value_counts(insert_values)
+                        duplicates_struct = pc.filter(value_counts, pc.greater(value_counts.field("counts"), 1))
+                        duplicates = duplicates_struct.field("values")
+                        if len(duplicates) > 0:
+                            raise flight.FlightServerError(
+                                f"Cannot insert rows with duplicate values in primary key column {column_name}: {duplicates}",
+                                extra_info=msgpack.packb(
+                                    {
+                                        "exception_type": "ConstraintException",
+                                        "message": f"""Duplicate key "{column_name}: {duplicates[0]}" violates primary key constraint.""",
+                                    }
+                                ),
+                            )
+
+                        already_exists = pc.is_in(insert_values, value_set=existing_table.column(column_name))
+                        existing_rows = pc.filter(insert_values, already_exists)
+                        if len(existing_rows) > 0:
+                            raise flight.FlightServerError(
+                                f"Cannot insert rows with duplicate values in primary key column {column_name}: {duplicates}",
+                                extra_info=msgpack.packb(
+                                    {
+                                        "exception_type": "ConstraintException",
+                                        "message": f"""Duplicate key "{column_name}: {existing_rows[0]}" violates primary key constraint.""",
+                                    }
+                                ),
+                            )
+
                     assert new_rows.num_rows > 0
 
                     # append the row id column to the new rows.