Port to psycopg3 (#1324)

* WIP.

* Add some comments about porting from psycopg 2 to 3 (#1318)

* WIP

* Disable _set_wait_callback()

* TransactionStatus.

* First working query.

* More pg3 changes.

* test_pgexecute still fails.

* Fix bytea support.

* Fix json and enum unicode.

* Get unit tests to pass.

* Behave tests still break, WIP.

* Prompt seems to be displayed fine, why don't the tests see the whitespace?

* Python version.

* Fix test.

* Black.

* Added black to dev reqs.

* nbu link for donations.

* Use psycopg.sql to format statement.

* Special case for show help in pgbouncer.

* Fix test.

* Added integration test.

* Install pgbouncer in ci.

* Fix integration test.

* Remove tmate session.

* Revert commenting out python versions.

* Pin pgspecial to >=2.

* Changelog.

Co-authored-by: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Co-authored-by: Amjith Ramanujam <amjith.r@gmail.com>

12 files changed, 101 insertions, 84 deletions

diff --git a/tests/features/basic_commands.feature b/tests/features/basic_commands.feature
index 99f893e2..cd15306b 100644
--- a/tests/features/basic_commands.feature
+++ b/tests/features/basic_commands.feature
@@ -49,7 +49,6 @@ Feature: run the cli,
       when we send "\?" command
       then we see help output
 
-  @wip
   Scenario: run the cli with dsn and password
     When we launch dbcli using dsn_password
       then we send password
diff --git a/tests/features/db_utils.py b/tests/features/db_utils.py
index 6898394e..595c6c2c 100644
--- a/tests/features/db_utils.py
+++ b/tests/features/db_utils.py
@@ -1,5 +1,4 @@
-from psycopg2 import connect
-from psycopg2.extensions import AsIs
+from psycopg import connect
 
 
 def create_db(
@@ -17,13 +16,10 @@ def create_db(
     """
     cn = create_cn(hostname, password, username, "postgres", port)
 
-    # ISOLATION_LEVEL_AUTOCOMMIT = 0
-    # Needed for DB creation.
-    cn.set_isolation_level(0)
-
+    cn.autocommit = True
     with cn.cursor() as cr:
-        cr.execute("drop database if exists %s", (AsIs(dbname),))
-        cr.execute("create database %s", (AsIs(dbname),))
+        cr.execute(f"drop database if exists {dbname}")
+        cr.execute(f"create database {dbname}")
 
     cn.close()
 
@@ -41,13 +37,26 @@ def create_cn(hostname, password, username, dbname, port):
     :return: psycopg2.connection
     """
     cn = connect(
-        host=hostname, user=username, database=dbname, password=password, port=port
+        host=hostname, user=username, dbname=dbname, password=password, port=port
     )
 
-    print(f"Created connection: {cn.dsn}.")
+    print(f"Created connection: {cn.info.get_parameters()}.")
     return cn
 
 
+def pgbouncer_available(hostname="localhost", password=None, username="postgres"):
+    cn = None
+    try:
+        cn = create_cn(hostname, password, username, "pgbouncer", 6432)
+        return True
+    except:
+        print("Pgbouncer is not available.")
+    finally:
+        if cn:
+            cn.close()
+    return False
+
+
 def drop_db(hostname="localhost", username=None, password=None, dbname=None, port=None):
     """
     Drop database.
@@ -58,12 +67,11 @@ def drop_db(hostname="localhost", username=None, password=None, dbname=None, por
     """
     cn = create_cn(hostname, password, username, "postgres", port)
 
-    # ISOLATION_LEVEL_AUTOCOMMIT = 0
     # Needed for DB drop.
-    cn.set_isolation_level(0)
+    cn.autocommit = True
 
     with cn.cursor() as cr:
-        cr.execute("drop database if exists %s", (AsIs(dbname),))
+        cr.execute(f"drop database if exists {dbname}")
 
     close_cn(cn)
 
@@ -74,5 +82,6 @@ def close_cn(cn=None):
     :param connection: psycopg2.connection
     """
     if cn:
+        cn_params = cn.info.get_parameters()
         cn.close()
-        print(f"Closed connection: {cn.dsn}.")
+        print(f"Closed connection: {cn_params}.")
diff --git a/tests/features/environment.py b/tests/features/environment.py
index 215c85cd..b93c12ec 100644
--- a/tests/features/environment.py
+++ b/tests/features/environment.py
@@ -111,7 +111,11 @@ def before_all(context):
         context.conf["dbname"],
         context.conf["port"],
     )
-
+    context.pgbouncer_available = dbutils.pgbouncer_available(
+        hostname=context.conf["host"],
+        password=context.conf["pass"],
+        username=context.conf["user"],
+    )
     context.fixture_data = fixutils.read_fixture_files()
 
     # use temporary directory as config home
@@ -164,7 +168,19 @@ def before_scenario(context, scenario):
     if scenario.name == "list databases":
         # not using the cli for that
         return
-    wrappers.run_cli(context)
+    currentdb = None
+    if "pgbouncer" in scenario.feature.tags:
+        if context.pgbouncer_available:
+            os.environ["PGDATABASE"] = "pgbouncer"
+            os.environ["PGPORT"] = "6432"
+            currentdb = "pgbouncer"
+        else:
+            scenario.skip()
+    else:
+        # set env vars back to normal test database
+        os.environ["PGDATABASE"] = context.conf["dbname"]
+        os.environ["PGPORT"] = context.conf["port"]
+    wrappers.run_cli(context, currentdb=currentdb)
     wrappers.wait_prompt(context)
 
 
@@ -172,13 +188,17 @@ def after_scenario(context, scenario):
     """Cleans up after each scenario completes."""
     if hasattr(context, "cli") and context.cli and not context.exit_sent:
         # Quit nicely.
-        if not context.atprompt:
+        if not getattr(context, "atprompt", False):
             dbname = context.currentdb
-            context.cli.expect_exact(f"{dbname}> ", timeout=15)
-        context.cli.sendcontrol("c")
-        context.cli.sendcontrol("d")
+            context.cli.expect_exact(f"{dbname}>", timeout=5)
+        try:
+            context.cli.sendcontrol("c")
+            context.cli.sendcontrol("d")
+        except Exception as x:
+            print("Failed cleanup after scenario:")
+            print(x)
         try:
-            context.cli.expect_exact(pexpect.EOF, timeout=15)
+            context.cli.expect_exact(pexpect.EOF, timeout=5)
         except pexpect.TIMEOUT:
             print(f"--- after_scenario {scenario.name}: kill cli")
             context.cli.kill(signal.SIGKILL)
diff --git a/tests/features/pgbouncer.feature b/tests/features/pgbouncer.feature
new file mode 100644
index 00000000..14cc5ad8
--- /dev/null
+++ b/tests/features/pgbouncer.feature
@@ -0,0 +1,12 @@
+@pgbouncer
+Feature: run pgbouncer,
+  call the help command,
+  exit the cli
+
+  Scenario: run "show help" command
+     When we send "show help" command
+      then we see the pgbouncer help output
+
+  Scenario: run the cli and exit
+     When we send "ctrl + d"
+      then dbcli exits
diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py
index a7c99eea..7c878143 100644
--- a/tests/features/steps/basic_commands.py
+++ b/tests/features/steps/basic_commands.py
@@ -69,7 +69,7 @@ def step_ctrl_d(context):
     context.cli.sendline(r"\pset pager off")
     wrappers.wait_prompt(context)
     context.cli.sendcontrol("d")
-    context.cli.expect(pexpect.EOF, timeout=15)
+    context.cli.expect(pexpect.EOF, timeout=5)
     context.exit_sent = True
 
 
diff --git a/tests/features/steps/crud_database.py b/tests/features/steps/crud_database.py
index 3f5d0e71..87cdc85b 100644
--- a/tests/features/steps/crud_database.py
+++ b/tests/features/steps/crud_database.py
@@ -59,7 +59,7 @@ def step_see_prompt(context):
     Wait to see the prompt.
     """
     db_name = getattr(context, "currentdb", context.conf["dbname"])
-    wrappers.expect_exact(context, f"{db_name}> ", timeout=5)
+    wrappers.expect_exact(context, f"{db_name}>", timeout=5)
     context.atprompt = True
 
 
diff --git a/tests/features/steps/pgbouncer.py b/tests/features/steps/pgbouncer.py
new file mode 100644
index 00000000..f1569823
--- /dev/null
+++ b/tests/features/steps/pgbouncer.py
@@ -0,0 +1,22 @@
+"""
+Steps for behavioral style tests are defined in this module.
+Each step is defined by the string decorating it.
+This string is used to call the step in "*.feature" file.
+"""
+
+from behave import when, then
+import wrappers
+
+
+@when('we send "show help" command')
+def step_send_help_command(context):
+    context.cli.sendline("show help")
+
+
+@then("we see the pgbouncer help output")
+def see_pgbouncer_help(context):
+    wrappers.expect_exact(
+        context,
+        "SHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION",
+        timeout=3,
+    )
diff --git a/tests/features/steps/wrappers.py b/tests/features/steps/wrappers.py
index 0ca83669..61805170 100644
--- a/tests/features/steps/wrappers.py
+++ b/tests/features/steps/wrappers.py
@@ -70,4 +70,5 @@ def run_cli(context, run_args=None, prompt_check=True, currentdb=None):
 
 def wait_prompt(context):
     """Make sure prompt is displayed."""
-    expect_exact(context, "{0}> ".format(context.conf["dbname"]), timeout=5)
+    prompt_str = "{0}>".format(context.currentdb)
+    expect_exact(context, [prompt_str + " ", prompt_str, pexpect.EOF], timeout=3)
diff --git a/tests/test_exceptionals.py b/tests/test_exceptionals.py
deleted file mode 100644
index e69de29b..00000000
--- a/tests/test_exceptionals.py
+++ /dev/null
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 109674cb..7fef4658 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -1,6 +1,6 @@
 from textwrap import dedent
 
-import psycopg2
+import psycopg
 import pytest
 from unittest.mock import patch, MagicMock
 from pgspecial.main import PGSpecial, NO_QUERY
@@ -428,7 +428,7 @@ def test_describe_special(executor, command, verbose, pattern, pgspecial):
 @dbtest
 @pytest.mark.parametrize("sql", ["invalid sql", "SELECT 1; select error;"])
 def test_raises_with_no_formatter(executor, sql):
-    with pytest.raises(psycopg2.ProgrammingError):
+    with pytest.raises(psycopg.ProgrammingError):
         list(executor.run(sql))
 
 
@@ -513,13 +513,6 @@ def test_short_host(executor):
         assert executor.short_host == "localhost1"
 
 
-class BrokenConnection:
-    """Mock a connection that failed."""
-
-    def cursor(self):
-        raise psycopg2.InterfaceError("I'm broken!")
-
-
 class VirtualCursor:
     """Mock a cursor to virtual database like pgbouncer."""
 
@@ -549,13 +542,15 @@ def test_exit_without_active_connection(executor):
         aliases=(":q",),
     )
 
-    with patch.object(executor, "conn", BrokenConnection()):
+    with patch.object(
+        executor.conn, "cursor", side_effect=psycopg.InterfaceError("I'm broken!")
+    ):
         # we should be able to quit the app, even without active connection
         run(executor, "\\q", pgspecial=pgspecial)
         quit_handler.assert_called_once()
 
         # an exception should be raised when running a query without active connection
-        with pytest.raises(psycopg2.InterfaceError):
+        with pytest.raises(psycopg.InterfaceError):
             run(executor, "select 1", pgspecial=pgspecial)
 
 
diff --git a/tests/test_plan.wiki b/tests/test_plan.wiki
deleted file mode 100644
index 6812f18a..00000000
--- a/tests/test_plan.wiki
+++ /dev/null
@@ -1,38 +0,0 @@
-= Gross Checks =
-    * [ ] Check connecting to a local database.
-    * [ ] Check connecting to a remote database.
-    * [ ] Check connecting to a database with a user/password.
-    * [ ] Check connecting to a non-existent database.
-    * [ ] Test changing the database.
-
-    == PGExecute ==
-    * [ ] Test successful execution given a cursor.
-    * [ ] Test unsuccessful execution with a syntax error.
-    * [ ] Test a series of executions with the same cursor without failure.
-    * [ ] Test a series of executions with the same cursor with failure.
-    * [ ] Test passing in a special command.
-
-    == Naive Autocompletion ==
-    * [ ] Input empty string, ask for completions - Everything.
-    * [ ] Input partial prefix, ask for completions - Stars with prefix.
-    * [ ] Input fully autocompleted string, ask for completions - Only full match
-    * [ ] Input non-existent prefix, ask for completions - nothing
-    * [ ] Input lowercase prefix - case insensitive completions
-
-    == Smart Autocompletion ==
-    * [ ] Input empty string and check if only keywords are returned.
-    * [ ] Input SELECT prefix and check if only columns are returned.
-    * [ ] Input SELECT blah - only keywords are returned.
-    * [ ] Input SELECT * FROM - Table names only
-
-    == PGSpecial ==
-    * [ ] Test \d
-    * [ ] Test \d tablename
-    * [ ] Test \d tablena*
-    * [ ] Test \d non-existent-tablename
-    * [ ] Test \d index
-    * [ ] Test \d sequence
-    * [ ] Test \d view
-
-    == Exceptionals ==
-    * [ ] Test the 'use' command to change db.
diff --git a/tests/utils.py b/tests/utils.py
index 460ea469..67d769fd 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,8 +1,6 @@
 import pytest
-import psycopg2
-import psycopg2.extras
+import psycopg
 from pgcli.main import format_output, OutputSettings
-from pgcli.pgexecute import register_json_typecasters
 from os import getenv
 
 POSTGRES_USER = getenv("PGUSER", "postgres")
@@ -12,12 +10,12 @@ POSTGRES_PASSWORD = getenv("PGPASSWORD", "postgres")
 
 
 def db_connection(dbname=None):
-    conn = psycopg2.connect(
+    conn = psycopg.connect(
         user=POSTGRES_USER,
         host=POSTGRES_HOST,
         password=POSTGRES_PASSWORD,
         port=POSTGRES_PORT,
-        database=dbname,
+        dbname=dbname,
     )
     conn.autocommit = True
     return conn
@@ -26,11 +24,10 @@ def db_connection(dbname=None):
 try:
     conn = db_connection()
     CAN_CONNECT_TO_DB = True
-    SERVER_VERSION = conn.server_version
-    json_types = register_json_typecasters(conn, lambda x: x)
-    JSON_AVAILABLE = "json" in json_types
-    JSONB_AVAILABLE = "jsonb" in json_types
-except:
+    SERVER_VERSION = conn.info.parameter_status("server_version")
+    JSON_AVAILABLE = True
+    JSONB_AVAILABLE = True
+except Exception as x:
     CAN_CONNECT_TO_DB = JSON_AVAILABLE = JSONB_AVAILABLE = False
     SERVER_VERSION = 0