瀏覽代碼

Add more functionality to Postgresql driver

oz123 10 年之前
父節點
當前提交
1fa1d2dbe4
共有 2 個文件被更改,包括 22 次插入19 次删除
  1. 14 18
      pwman/data/drivers/postgresql.py
  2. 8 1
      pwman/tests/test_postgresql.py

+ 14 - 18
pwman/data/drivers/postgresql.py

@@ -194,24 +194,14 @@ class PostgresqlDatabase(Database):
         self._checktags()
         self._commit()
 
-    def addnodes(self, nodes):
-        cursor = self._get_cur()
-        for n in nodes:
-            sql = "INSERT INTO %sNODES(DATA) VALUES(%%(data)s)" % (
-                self._prefix)
-            if not isinstance(n, Node):
-                raise DatabaseException("Tried to insert foreign object into "
-                                        "database [%s]", n)
-            values = {"data": cPickle.dumps(n)}
-            try:
-                cursor.execute(sql, values)
-            except pgdb.DatabaseError as e:
-                raise DatabaseException("Postgresql: %s" % (e))
-            id = self._lastrowid("NODES")
-            n.set_id(id)
-
-            self._setnodetags(n)
-        self._commit()
+    def add_node(self, node):
+        sql = ("INSERT INTO NODE(USERNAME, PASSWORD, URL, NOTES)"
+               "VALUES(%s, %s, %s, %s)")
+        node_tags = list(node)
+        node, tags = node_tags[:4], node_tags[-1]
+        self._cur.execute(sql, (node))
+        #self._setnodetags(self._cur.lastrowid, tags)
+        self._con.commit()
 
     def removenodes(self, nodes):
         cursor = self._get_cur()
@@ -397,6 +387,12 @@ class PostgresqlDatabase(Database):
         except pg.ProgrammingError:
             self._con.rollback()
 
+    def save_crypto_info(self, seed, digest):
+        """save the random seed and the digested key"""
+        self._cur.execute("DELETE  FROM CRYPTO")
+        self._cur.execute("INSERT INTO CRYPTO VALUES(%s, %s)", (seed, digest))
+        self._con.commit()
+
     def savekey(self, key):
         salt, digest = key.split('$6$')
         sql = "INSERT INTO CRYPTO(SEED, DIGEST) VALUES(%s,%s)"

+ 8 - 1
pwman/tests/test_postgresql.py

@@ -46,7 +46,6 @@ class TestPostGresql(unittest.TestCase):
         self.db._con.commit()
 
     def test_1_con(self):
-
         self.assertIsInstance(self.db._cur, pg._psycopg.cursor)
 
     def test_2_create_tables(self):
@@ -59,6 +58,14 @@ class TestPostGresql(unittest.TestCase):
         secretkey = self.db.loadkey()
         self.assertEqual(secretkey, 'SECRET$6$KEY')
 
+    def test_4_save_crypto(self):
+        self.db.save_crypto_info("TOP", "SECRET")
+        secretkey = self.db.loadkey()
+        self.assertEqual(secretkey, 'TOP$6$SECRET')
+
+    def test_5_add_node(self):
+        self.db.add_node(("TBONE", "S3K43T", "example.org", "some note"))
+
 
 if __name__ == '__main__':