summaryrefslogtreecommitdiff
path: root/lib/sqlalchemy/dialects/postgresql/base.py
diff options
context:
space:
mode:
authorMike Bayer <mike_mp@zzzcomputing.com>2020-08-18 22:53:09 -0400
committerMike Bayer <mike_mp@zzzcomputing.com>2020-08-19 10:43:38 -0400
commitd1005e130558b33fd455be6d994e8e863799a318 (patch)
tree4e6b3acdf0476ddcdb61afbce782c064afd9a12a /lib/sqlalchemy/dialects/postgresql/base.py
parent0901190bb440580f0664fe3f6310173762b908e0 (diff)
downloadsqlalchemy-d1005e130558b33fd455be6d994e8e863799a318.tar.gz
Implement DDL visitor for PG ENUM with schema translate support
Fixed issue where the :class:`_postgresql.ENUM` type would not consult the schema translate map when emitting a CREATE TYPE or DROP TYPE during the test to see if the type exists or not. Additionally, repaired an issue where if the same enum were encountered multiple times in a single DDL sequence, the "check" query would run repeatedly rather than relying upon a cached value. Fixes: #5520 Change-Id: I79f46e29ac0168e873ff178c242f8d78f6679aeb
Diffstat (limited to 'lib/sqlalchemy/dialects/postgresql/base.py')
-rw-r--r--lib/sqlalchemy/dialects/postgresql/base.py57
1 files changed, 47 insertions, 10 deletions
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 7717a2526..543746687 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1011,6 +1011,7 @@ from ...sql import expression
from ...sql import roles
from ...sql import sqltypes
from ...sql import util as sql_util
+from ...sql.ddl import DDLBase
from ...types import BIGINT
from ...types import BOOLEAN
from ...types import CHAR
@@ -1499,10 +1500,7 @@ class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum):
if not bind.dialect.supports_native_enum:
return
- if not checkfirst or not bind.dialect.has_type(
- bind, self.name, schema=self.schema
- ):
- bind.execute(CreateEnumType(self))
+ bind._run_ddl_visitor(self.EnumGenerator, self, checkfirst=checkfirst)
def drop(self, bind=None, checkfirst=True):
"""Emit ``DROP TYPE`` for this
@@ -1522,10 +1520,49 @@ class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum):
if not bind.dialect.supports_native_enum:
return
- if not checkfirst or bind.dialect.has_type(
- bind, self.name, schema=self.schema
- ):
- bind.execute(DropEnumType(self))
+ bind._run_ddl_visitor(self.EnumDropper, self, checkfirst=checkfirst)
+
+ class EnumGenerator(DDLBase):
+ def __init__(self, dialect, connection, checkfirst=False, **kwargs):
+ super(ENUM.EnumGenerator, self).__init__(connection, **kwargs)
+ self.checkfirst = checkfirst
+
+ def _can_create_enum(self, enum):
+ if not self.checkfirst:
+ return True
+
+ effective_schema = self.connection.schema_for_object(enum)
+
+ return not self.connection.dialect.has_type(
+ self.connection, enum.name, schema=effective_schema
+ )
+
+ def visit_enum(self, enum):
+ if not self._can_create_enum(enum):
+ return
+
+ self.connection.execute(CreateEnumType(enum))
+
+ class EnumDropper(DDLBase):
+ def __init__(self, dialect, connection, checkfirst=False, **kwargs):
+ super(ENUM.EnumDropper, self).__init__(connection, **kwargs)
+ self.checkfirst = checkfirst
+
+ def _can_drop_enum(self, enum):
+ if not self.checkfirst:
+ return True
+
+ effective_schema = self.connection.schema_for_object(enum)
+
+ return self.connection.dialect.has_type(
+ self.connection, enum.name, schema=effective_schema
+ )
+
+ def visit_enum(self, enum):
+ if not self._can_drop_enum(enum):
+ return
+
+ self.connection.execute(DropEnumType(enum))
def _check_for_name_in_memos(self, checkfirst, kw):
"""Look in the 'ddl runner' for 'memos', then
@@ -1551,14 +1588,14 @@ class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum):
return False
def _on_table_create(self, target, bind, checkfirst=False, **kw):
+
if (
checkfirst
or (
not self.metadata
and not kw.get("_is_metadata_operation", False)
)
- and not self._check_for_name_in_memos(checkfirst, kw)
- ):
+ ) and not self._check_for_name_in_memos(checkfirst, kw):
self.create(bind=bind, checkfirst=checkfirst)
def _on_table_drop(self, target, bind, checkfirst=False, **kw):