Moving the encrypt_password method from the connection class to the

psycopgmodule, and exported it from psycopg2.extensions as per review
comments.

lib/extensions.py

@@ -63,7 +63,7 @@ from psycopg2._psycopg import (                             # noqa
     string_types, binary_types, new_type, new_array_type, register_type,
     ISQLQuote, Notify, Diagnostics, Column,
     QueryCanceledError, TransactionRollbackError,
-    set_wait_callback, get_wait_callback, )
+    set_wait_callback, get_wait_callback, encrypt_password, )
 
 
 """Isolation level values."""

psycopg/connection_type.c

@@ -547,74 +547,6 @@ do { \
     EXC_IF_TPC_PREPARED(self, what); \
 } while(0)
 
-/* encrypt_password - Prepare the encrypted password form */
-#define psyco_encrypt_password_doc \
-"encrypt_password('password', 'user', ...) -- Prepares the encrypted form of a PostgreSQL password.\n\n" \
-"Accepted arguments are 'algorithm'."
-
-static PyObject *
-psyco_encrypt_password(connectionObject *self, PyObject *args, PyObject *kwargs)
-{
-    const char *password = NULL,
-               *user = NULL,
-               *algorithm = NULL;
-    char *encrypted = NULL;
-
-    PyObject *res = Py_None;
-
-    static char *kwlist[] = {"password", "user", "algorithm", NULL};
-
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "ss|s", kwlist, &password, &user, &algorithm)) {
-        return NULL;
-    }
-
-    if (self->server_version < 100000 ||
-        (algorithm && strcmp(algorithm, "md5") == 0)
-    ) {
-        encrypted = PQencryptPassword(password, user);
-
-        if (encrypted != NULL)
-        {
-            res = Text_FromUTF8(encrypted);
-            PQfreemem(encrypted);
-        }
-        return res;
-    }
-    else
-    {
-#if PG_VERSION_NUM >= 100000
-        encrypted = PQencryptPasswordConn(self->pgconn, password, user, algorithm);
-
-        if (!encrypted)
-        {
-            const char *msg;
-            msg = PQerrorMessage(self->pgconn);
-            if (msg && *msg) {
-                PyErr_Format(
-                    ProgrammingError,
-                    "Error encrypting the password!\n%s",
-                    msg
-                );
-                return NULL;
-            }
-        }
-        else
-        {
-            res = Text_FromUTF8(encrypted);
-            PQfreemem(encrypted);
-        }
-        return res;
-
-#else
-        PyErr_SetString(
-            NotSupportedError,
-            "Password encryption (other than 'md5' algorithm) is not supported for the server version >= 10 in libpq < 10"
-        );
-#endif
-    }
-    return NULL;
-}
-
 /* set_session - set default transaction characteristics */
 
 #define psyco_conn_set_session_doc \
@@ -1244,8 +1176,6 @@ static struct PyMethodDef connectionObject_methods[] = {
      METH_NOARGS, psyco_conn_isexecuting_doc},
     {"cancel", (PyCFunction)psyco_conn_cancel,
      METH_NOARGS, psyco_conn_cancel_doc},
-    {"encrypt_password", (PyCFunction)psyco_encrypt_password,
-     METH_VARARGS|METH_KEYWORDS, psyco_encrypt_password_doc},
     {NULL}
 };
 

psycopg/psycopgmodule.c

@@ -403,6 +403,109 @@ psyco_libpq_version(PyObject *self)
 #endif
 }
 
+/* encrypt_password - Prepare the encrypted password form */
+#define psyco_encrypt_password_doc \
+"encrypt_password(password, user, [conn_or_curs], [algorithm]) -- Prepares the encrypted form of a PostgreSQL password.\n\n"
+
+static PyObject *
+psyco_encrypt_password(PyObject *self, PyObject *args)
+{
+    char *encrypted = NULL;
+
+    PyObject *obj = NULL,
+             *res = Py_None,
+             *password = NULL,
+             *user = NULL,
+             *algorithm = NULL;
+
+    connectionObject *conn = NULL;
+
+    static char *kwlist[] = {"password", "user", "scope", "algorithm", NULL};
+
+    if (!PyArg_ParseTuple(args, "OO|OO",
+            &password, &user, &obj, &algorithm)) {
+        return NULL;
+    }
+
+    if (obj != NULL && obj != Py_None) {
+        if (PyObject_TypeCheck(obj, &cursorType)) {
+            conn = ((cursorObject*)obj)->conn;
+        }
+        else if (PyObject_TypeCheck(obj, &connectionType)) {
+            conn = (connectionObject*)obj;
+        }
+        else {
+            PyErr_SetString(PyExc_TypeError,
+                    "argument 3 must be a connection or a cursor");
+            return NULL;
+        }
+    }
+
+    /* for ensure_bytes */
+    Py_INCREF(user);
+    Py_INCREF(password);
+    if (algorithm) {
+        Py_INCREF(algorithm);
+    }
+
+    if (!(user = psycopg_ensure_bytes(user))) { goto exit; }
+    if (!(password = psycopg_ensure_bytes(password))) { goto exit; }
+    if (algorithm && !(algorithm = psycopg_ensure_bytes(algorithm))) {
+        goto exit;
+    }
+
+    /* Use the libpq API 'PQencryptPassword', when no connection object is
+       available, or the algorithm is set to as 'md5', or the database server
+       version < 10 */
+    if (conn == NULL || conn->server_version < 100000 ||
+            (algorithm != NULL && algorithm != Py_None &&
+             strcmp(Bytes_AS_STRING(algorithm), "md5") == 0)) {
+        encrypted = PQencryptPassword(Bytes_AS_STRING(password),
+            Bytes_AS_STRING(user));
+
+        if (encrypted != NULL) {
+            res = Text_FromUTF8(encrypted);
+            PQfreemem(encrypted);
+        }
+        goto exit;
+    }
+
+#if PG_VERSION_NUM >= 100000
+    encrypted = PQencryptPasswordConn(conn->pgconn, Bytes_AS_STRING(password),
+            Bytes_AS_STRING(user),
+            algorithm ?  Bytes_AS_STRING(algorithm) : NULL);
+
+    if (!encrypted) {
+        const char *msg = PQerrorMessage(conn->pgconn);
+        if (msg && *msg) {
+            PyErr_Format(ProgrammingError, msg);
+            res = NULL;
+            goto exit;
+        }
+    }
+    else {
+        res = Text_FromUTF8(encrypted);
+        PQfreemem(encrypted);
+    }
+#else
+    PyErr_SetString(
+            NotSupportedError,
+            "Password encryption (other than 'md5' algorithm) is not supported for the server version >= 10 in libpq < 10"
+            );
+    res = NULL;
+#endif
+
+exit:
+    Py_XDECREF(user);
+    Py_XDECREF(password);
+    if (algorithm) {
+        Py_XDECREF(algorithm);
+    }
+
+    return res;
+}
+
+
 /* psyco_encodings_fill
 
    Fill the module's postgresql<->python encoding table */
@@ -852,6 +955,8 @@ static PyMethodDef psycopgMethods[] = {
      METH_O, psyco_set_wait_callback_doc},
     {"get_wait_callback",  (PyCFunction)psyco_get_wait_callback,
      METH_NOARGS, psyco_get_wait_callback_doc},
+    {"encrypt_password", (PyCFunction)psyco_encrypt_password,
+     METH_VARARGS, psyco_encrypt_password_doc},
 
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };

tests/test_connection.py

@@ -1392,17 +1392,20 @@ class TransactionControlTests(ConnectingTestCase):
 
         # MD5 algorithm
         self.assertEqual(
-            self.conn.encrypt_password('psycopg2', 'ashesh', 'md5'),
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn, 'md5'),
             'md594839d658c28a357126f105b9cb14cfc'
         )
 
         if libpq_version() < 100000:
             self.assertRaises(
                 psycopg2.NotSupportedError,
-                self.conn.encrypt_password, 'psycopg2', 'ashesh'
+                ext.encrypt_password, 'psycopg2', 'ashesh', self.conn,
+                'scram-sha-256'
             )
         else:
-            enc_password = self.conn.encrypt_password('psycopg2', 'ashesh')
+            enc_password = ext.encrypt_password(
+                'psycopg2', 'ashesh', self.conn
+            )
             if server_encryption_algorithm == 'md5':
                 self.assertEqual(
                     enc_password, 'md594839d658c28a357126f105b9cb14cfc'
@@ -1411,22 +1414,22 @@ class TransactionControlTests(ConnectingTestCase):
                 self.assertEqual(enc_password[:14], 'SCRAM-SHA-256$')
 
             self.assertEqual(
-                self.conn.encrypt_password(
+                ext.encrypt_password(
-                    'psycopg2', 'ashesh', 'scram-sha-256'
+                    'psycopg2', 'ashesh', self.conn, 'scram-sha-256'
                 )[:14], 'SCRAM-SHA-256$'
             )
 
     @skip_after_postgres(10)
     def test_encrypt_password_pre_10(self):
         self.assertEqual(
-            self.conn.encrypt_password('psycopg2', 'ashesh'),
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn),
             'md594839d658c28a357126f105b9cb14cfc'
         )
 
         # Encryption algorithm will be ignored for postgres version < 10, it
         # will always use MD5.
         self.assertEqual(
-            self.conn.encrypt_password('psycopg2', 'ashesh', 'abc'),
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn, 'abc'),
             'md594839d658c28a357126f105b9cb14cfc'
         )