Fix handling of composite pks in peewee.

1 files changed, 15 insertions, 0 deletions

diff --git a/archivist/peewee_fixes.py b/archivist/peewee_fixes.py
new file mode 100644
index 0000000..ff914da
--- /dev/null
+++ b/archivist/peewee_fixes.py
@@ -0,0 +1,15 @@
+from peewee import QueryCompiler, strip_parens, ForeignKeyField
+
+__all__ = []
+
+def _parse_select_query(self, node, alias_map, conv):
+    clone = node.clone()
+    if not node._explicit_selection:
+        if conv and isinstance(conv, ForeignKeyField):
+            clone._select = (conv.to_field,)
+        else:
+            clone._select = clone.model_class._meta.get_primary_key_fields()
+    sub, params = self.generate_select(clone, alias_map)
+    return '(%s)' % strip_parens(sub), params
+
+QueryCompiler._parse_select_query = _parse_select_query