Add precondition check for database validation

data/database.py

@@ -30,6 +30,7 @@ from data.fields import (ResumableSHA256Field, ResumableSHA1Field, JSONField, Ba
 from data.text import match_mysql, match_like
 from data.read_slave import ReadSlaveModel
 from util.names import urn_generator
+from util.validation import validate_postgres_precondition
 
 
 logger = logging.getLogger(__name__)
@@ -70,6 +71,12 @@ SCHEME_RANDOM_FUNCTION = {
 }
 
 
+PRECONDITION_VALIDATION = {
+  'postgresql': validate_postgres_precondition,
+  'postgresql+psycopg2': validate_postgres_precondition,
+}
+
+
 _EXTRA_ARGS = {
   'mysql': dict(charset='utf8mb4'),
   'mysql+pymysql': dict(charset='utf8mb4'),
@@ -284,6 +291,25 @@ def validate_database_url(url, db_kwargs, connect_timeout=5):
       pass
 
 
+def validate_database_precondition(url, db_kwargs, connect_timeout=5):
+  """ Validates that we can connect to the given database URL and the database meets our
+      precondition. Raises an exception if the validation fails. """
+  db_kwargs = db_kwargs.copy()
+  try:
+    driver = _db_from_url(url, db_kwargs, connect_timeout=connect_timeout, allow_retry=False,
+                          allow_pooling=False)
+    driver.connect()
+    pre_condition_check = PRECONDITION_VALIDATION.get(make_url(url).drivername)
+    if pre_condition_check:
+      pre_condition_check(driver)
+
+  finally:
+    try:
+      driver.close()
+    except:
+      pass
+
+
 def _wrap_for_retry(driver):
   return type('Retrying' + driver.__class__.__name__, (RetryOperationalError, driver), {})