From 6fb7b527426b564302d6fd8ca11dfa7b78cc2e7a Mon Sep 17 00:00:00 2001
From: Gord Thompson <gord@gordthompson.com>
Date: Sun, 12 Jan 2020 20:08:22 -0500
Subject: Refactor test provisioning to dialect-level files

Fixes: #5085

<!-- Provide a general summary of your proposed changes in the Title field above -->

Move dialect-specific provisioning code to dialect-level copies of provision.py.

<!-- go over following points. check them with an `x` if they do apply, (they turn into clickable checkboxes once the PR is submitted, so no need to do everything at once)

-->

This pull request is:

- [ ] A documentation / typographical error fix
	- Good to go, no issue or tests are needed
- [x] A short code fix
	- please include the issue number, and create an issue if none exists, which
	  must include a complete example of the issue.  one line code fixes without an
	  issue and demonstration will not be accepted.
	- Please include: `Fixes: #<issue number>` in the commit message
	- please include tests.   one line code fixes without tests will not be accepted.
- [ ] A new feature implementation
	- please include the issue number, and create an issue if none exists, which must
	  include a complete example of how the feature would look.
	- Please include: `Fixes: #<issue number>` in the commit message
	- please include tests.

**Have a nice day!**

Closes: #5092
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/5092
Pull-request-sha: 25b9b7a9800549fb823576af8674e8d33ff4b2c1

Change-Id: Ie0b4a69aa472a60bdbd825e04c8595382bcc98e1
---
 lib/sqlalchemy/dialects/postgresql/provision.py | 67 +++++++++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 lib/sqlalchemy/dialects/postgresql/provision.py

(limited to 'lib/sqlalchemy/dialects/postgresql/provision.py')

diff --git a/lib/sqlalchemy/dialects/postgresql/provision.py b/lib/sqlalchemy/dialects/postgresql/provision.py
new file mode 100644
index 000000000..404da93a8
--- /dev/null
+++ b/lib/sqlalchemy/dialects/postgresql/provision.py
@@ -0,0 +1,67 @@
+import logging
+import time
+
+from ... import exc
+from ... import text
+from ...testing.provision import create_db
+from ...testing.provision import drop_db
+from ...testing.provision import temp_table_keyword_args
+
+
+log = logging.getLogger(__name__)
+
+
+@create_db.for_db("postgresql")
+def _pg_create_db(cfg, eng, ident):
+    template_db = cfg.options.postgresql_templatedb
+
+    with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
+        try:
+            _pg_drop_db(cfg, conn, ident)
+        except Exception:
+            pass
+        if not template_db:
+            template_db = conn.scalar("select current_database()")
+
+        attempt = 0
+        while True:
+            try:
+                conn.execute(
+                    "CREATE DATABASE %s TEMPLATE %s" % (ident, template_db)
+                )
+            except exc.OperationalError as err:
+                attempt += 1
+                if attempt >= 3:
+                    raise
+                if "accessed by other users" in str(err):
+                    log.info(
+                        "Waiting to create %s, URI %r, "
+                        "template DB %s is in use sleeping for .5",
+                        ident,
+                        eng.url,
+                        template_db,
+                    )
+                    time.sleep(0.5)
+            except:
+                raise
+            else:
+                break
+
+
+@drop_db.for_db("postgresql")
+def _pg_drop_db(cfg, eng, ident):
+    with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
+        conn.execute(
+            text(
+                "select pg_terminate_backend(pid) from pg_stat_activity "
+                "where usename=current_user and pid != pg_backend_pid() "
+                "and datname=:dname"
+            ),
+            dname=ident,
+        )
+        conn.execute("DROP DATABASE %s" % ident)
+
+
+@temp_table_keyword_args.for_db("postgresql")
+def _postgresql_temp_table_keyword_args(cfg, eng):
+    return {"prefixes": ["TEMPORARY"]}
-- 
cgit v1.2.1