Merge pull request #18 from eberonsh/master

Add customized primary key field to Command
luojilab · Aug 9, 2021 · e3b48de · e3b48de
2 parents 582195a + ee803ff
commit e3b48de
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 19 deletions.
diff --git a/mirage/management/commands/mirage.py b/mirage/management/commands/mirage.py
@@ -15,29 +15,41 @@ def add_arguments(self, parser):
         parser.add_argument("--total", type=int, required=False)
         parser.add_argument("--limit", type=int, required=False)
         parser.add_argument("--tofield", type=str, required=False)
+        parser.add_argument("--idfield", type=str, required=False)
 
     def handle(self, *args, **options):
-        app = options['app']
-        model = options['model']
-        field = options['field']
-        method = options['method']
-        offset = options['offset'] or 0
-        total = options['total']
-        limit = options['limit'] or 1000
-        tofield = options['tofield']
+        app = options["app"]
+        model = options["model"]
+        field = options["field"]
+        method = options["method"]
+        offset = options["offset"] or 0
+        total = options["total"]
+        limit = options["limit"] or 1000
+        tofield = options["tofield"]
+        idfield = options["idfield"] or "id"
 
-        if method == 'encrypt':
-            Migrator(app, model, field).encrypt(offset=offset, total=total, limit=limit)
-        elif method == 'decrypt':
-            Migrator(app, model, field).decrypt(offset=offset, total=total, limit=limit)
-        elif method == 'copy_to':
+        if method == "encrypt":
+            Migrator(app, model, field, idfield=idfield).encrypt(
+                offset=offset, total=total, limit=limit
+            )
+        elif method == "decrypt":
+            Migrator(app, model, field, idfield=idfield).decrypt(
+                offset=offset, total=total, limit=limit
+            )
+        elif method == "copy_to":
             assert tofield is not None
-            Migrator(app, model, field, tofield=tofield).copy_to(offset=offset, total=total, limit=limit)
-        elif method == 'decrypt_to':
+            Migrator(app, model, field, tofield=tofield, idfield=idfield).copy_to(
+                offset=offset, total=total, limit=limit
+            )
+        elif method == "decrypt_to":
             assert tofield is not None
-            Migrator(app, model, field, tofield=tofield).decrypt_to(offset=offset, total=total, limit=limit)
-        elif method == 'encrypt_to':
+            Migrator(app, model, field, tofield=tofield, idfield=idfield).decrypt_to(
+                offset=offset, total=total, limit=limit
+            )
+        elif method == "encrypt_to":
             assert tofield is not None
-            Migrator(app, model, field, tofield=tofield).encrypt_to(offset=offset, total=total, limit=limit)
+            Migrator(app, model, field, tofield=tofield, idfield=idfield).encrypt_to(
+                offset=offset, total=total, limit=limit
+            )
         else:
             return
diff --git a/mirage/tools.py b/mirage/tools.py
@@ -62,7 +62,7 @@ def executor(self, apps=None, schema_editor=None, offset=0, total=None, limit=10
                 else:
                     cursor.execute(
                         f"select {self.idfield}, {self.field} from {db_table} where "
-                        "{self.idfield}>{offset} order by {self.idfield} limit {limit};"
+                        f"{self.idfield}>{offset} order by {self.idfield} limit {limit};"
                     )
                 for query in cursor.fetchall():
                     if method in ['encrypt', 'encrypt_to']: