Replace ndb "auto" types with unified String

This simplifies the ndb "AutoStringXYZ" series
of classes into a single oslo_db.sqlalchemy.types.String
type which features all necessary behaviors via
two new flags, mysql_ndb_length and mysql_ndb_type.

Change-Id: I7f9c830073bf9a30abce0aa4bb55b5c9cf661afe

3 files changed, 119 insertions, 62 deletions

diff --git a/oslo_db/sqlalchemy/ndb.py b/oslo_db/sqlalchemy/ndb.py
index c7a3de9..e48c569 100644
--- a/oslo_db/sqlalchemy/ndb.py
+++ b/oslo_db/sqlalchemy/ndb.py
@@ -15,9 +15,17 @@
 
 import re
 
-from sqlalchemy import String, event, schema
+import debtcollector.removals
+
+from oslo_db.sqlalchemy.types import String
+
+from sqlalchemy import event, schema
+from sqlalchemy.dialects.mysql import TEXT
+from sqlalchemy.dialects.mysql import TINYTEXT
 from sqlalchemy.ext.compiler import compiles
-from sqlalchemy.types import VARCHAR
+from sqlalchemy.types import String as _String
+from sqlalchemy.types import to_instance
+
 
 engine_regex = re.compile("engine=innodb", re.IGNORECASE)
 trans_regex = re.compile("savepoint|rollback|release savepoint", re.IGNORECASE)
@@ -78,60 +86,43 @@ def prefix_inserts(create_table, compiler, **kw):
     return existing
 
 
-class AutoStringTinyText(String):
-    """Class definition for AutoStringTinyText.
-
-    Class is used by compiler function _auto-string_tiny_text().
-    """
-
-    pass
-
-
-@compiles(AutoStringTinyText, 'mysql')
-def _auto_string_tiny_text(element, compiler, **kw):
-    if ndb_status(compiler):
-        return "TINYTEXT"
-    else:
-        return compiler.visit_string(element, **kw)
+@compiles(String, "mysql")
+def _compile_ndb_string(element, compiler, **kw):
+    """Process ndb specific overrides for String.
 
+    Function will intercept mysql_ndb_length and mysql_ndb_type
+    arguments to adjust columns automatically.
 
-class AutoStringText(String):
-    """Class definition for AutoStringText.
+    mysql_ndb_length argument will adjust the String length
+    to the requested value.
 
-    Class is used by compiler function _auto_string_text().
+    mysql_ndb_type will change the column type to the requested
+    data type.
     """
+    if not ndb_status(compiler):
+        return compiler.visit_string(element, **kw)
 
-    pass
-
-
-@compiles(AutoStringText, 'mysql')
-def _auto_string_text(element, compiler, **kw):
-    if ndb_status(compiler):
-        return "TEXT"
+    if element.mysql_ndb_length:
+        effective_type = element.adapt(
+            _String, length=element.mysql_ndb_length)
+        return compiler.visit_string(effective_type, **kw)
+    elif element.mysql_ndb_type:
+        effective_type = to_instance(element.mysql_ndb_type)
+        return compiler.process(effective_type, **kw)
     else:
         return compiler.visit_string(element, **kw)
 
 
-class AutoStringSize(String):
-    """Class definition for AutoStringSize.
-
-    Class is used by the compiler function _auto_string_size().
-    """
+@debtcollector.removals.remove
+def AutoStringTinyText(length, **kw):
+    return String(length, mysql_ndb_type=TINYTEXT, *kw)
 
-    def __init__(self, length, ndb_size, **kw):
-        """Initialize and extend the String arguments.
 
-        Function adds the innodb_size and ndb_size arguments to the
-        function String().
-        """
-        super(AutoStringSize, self).__init__(length=length, **kw)
-        self.ndb_size = ndb_size
-        self.length = length
+@debtcollector.removals.remove
+def AutoStringText(length, **kw):
+    return String(length, mysql_ndb_type=TEXT, **kw)
 
 
-@compiles(AutoStringSize, 'mysql')
-def _auto_string_size(element, compiler, **kw):
-    if ndb_status(compiler):
-        return compiler.process(VARCHAR(element.ndb_size), **kw)
-    else:
-        return compiler.visit_string(element, **kw)
+@debtcollector.removals.remove
+def AutoStringSize(length, ndb_size, **kw):
+    return String(length, mysql_ndb_length=ndb_size, **kw)
diff --git a/oslo_db/sqlalchemy/types.py b/oslo_db/sqlalchemy/types.py
index 2dabe4c..08b1f2b 100644
--- a/oslo_db/sqlalchemy/types.py
+++ b/oslo_db/sqlalchemy/types.py
@@ -12,16 +12,18 @@
 
 import json
 
-from sqlalchemy.types import Integer, TypeDecorator, Text
 from sqlalchemy.dialects import mysql
+from sqlalchemy.types import Integer, Text, TypeDecorator, String as _String
 
 
 class JsonEncodedType(TypeDecorator):
     """Base column type for data serialized as JSON-encoded string in db."""
+
     type = None
     impl = Text
 
     def __init__(self, mysql_as_long=False, mysql_as_medium=False):
+        """Initialize JSON-encoding type."""
         super(JsonEncodedType, self).__init__()
 
         if mysql_as_long and mysql_as_medium:
@@ -34,6 +36,7 @@ class JsonEncodedType(TypeDecorator):
             self.impl = Text().with_variant(mysql.MEDIUMTEXT(), 'mysql')
 
     def process_bind_param(self, value, dialect):
+        """Bind parameters to the process."""
         if value is None:
             if self.type is not None:
                 # Save default value according to current type to keep the
@@ -48,6 +51,7 @@ class JsonEncodedType(TypeDecorator):
         return serialized_value
 
     def process_result_value(self, value, dialect):
+        """Process result value."""
         if value is not None:
             value = json.loads(value)
         return value
@@ -61,6 +65,7 @@ class JsonEncodedDict(JsonEncodedType):
     back. See this page for more robust work around:
     http://docs.sqlalchemy.org/en/rel_1_0/orm/extensions/mutable.html
     """
+
     type = dict
 
 
@@ -72,6 +77,7 @@ class JsonEncodedList(JsonEncodedType):
     back. See this page for more robust work around:
     http://docs.sqlalchemy.org/en/rel_1_0/orm/extensions/mutable.html
     """
+
     type = list
 
 
@@ -97,7 +103,26 @@ class SoftDeleteInteger(TypeDecorator):
     impl = Integer
 
     def process_bind_param(self, value, dialect):
+        """Return the binding parameter."""
         if value is None:
             return None
         else:
             return int(value)
+
+
+class String(_String):
+    """String subclass that implements oslo_db specific options.
+
+    Initial goal is to support ndb-specific flags.
+
+    mysql_ndb_type is used to override the String with another data type.
+    mysql_ndb_size is used to adjust the length of the String.
+
+    """
+
+    def __init__(
+            self, length, mysql_ndb_length=None, mysql_ndb_type=None, **kw):
+        """Initialize options."""
+        super(String, self).__init__(length, **kw)
+        self.mysql_ndb_type = mysql_ndb_type
+        self.mysql_ndb_length = mysql_ndb_length
diff --git a/oslo_db/tests/sqlalchemy/test_ndb.py b/oslo_db/tests/sqlalchemy/test_ndb.py
index dc5dc62..ea2c643 100644
--- a/oslo_db/tests/sqlalchemy/test_ndb.py
+++ b/oslo_db/tests/sqlalchemy/test_ndb.py
@@ -24,18 +24,20 @@ from oslo_db.sqlalchemy import ndb
 from oslo_db.sqlalchemy import test_fixtures
 from oslo_db.sqlalchemy import utils
 
+from oslo_db.sqlalchemy.types import String
+
 from oslotest import base as test_base
 
 from sqlalchemy import Column
 from sqlalchemy import Integer
 from sqlalchemy import MetaData
-from sqlalchemy import String
 from sqlalchemy import Table
 from sqlalchemy import Text
 
 from sqlalchemy import create_engine
 from sqlalchemy import schema
 
+from sqlalchemy.dialects.mysql import TEXT
 from sqlalchemy.dialects.mysql import TINYTEXT
 
 LOG = logging.getLogger(__name__)
@@ -43,9 +45,9 @@ LOG = logging.getLogger(__name__)
 _MOCK_CONNECTION = 'mysql+pymysql://'
 _TEST_TABLE = Table("test_ndb", MetaData(),
                     Column('id', Integer, primary_key=True),
-                    Column('test1', ndb.AutoStringTinyText(255)),
-                    Column('test2', ndb.AutoStringText(4096)),
-                    Column('test3', ndb.AutoStringSize(255, 64)),
+                    Column('test1', String(255, mysql_ndb_type=TEXT)),
+                    Column('test2', String(4096, mysql_ndb_type=TEXT)),
+                    Column('test3', String(255, mysql_ndb_length=64)),
                     mysql_engine='InnoDB')
 
 
@@ -56,6 +58,10 @@ class NDBMockTestBase(test_base.BaseTestCase):
         self.test_engine = test_engine = create_engine(
             _MOCK_CONNECTION, module=mock_dbapi)
         test_engine.dialect._oslodb_enable_ndb_support = True
+
+        self.addCleanup(
+            setattr, test_engine.dialect, "_oslodb_enable_ndb_support", False
+        )
         ndb.init_ndb_events(test_engine)
 
 
@@ -67,7 +73,6 @@ class NDBEventTestCase(NDBMockTestBase):
             str(schema.CreateTable(_TEST_TABLE).compile(
                 dialect=test_engine.dialect)),
             "ENGINE=NDBCLUSTER")
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
     def test_ndb_engine_override(self):
         test_engine = self.test_engine
@@ -76,7 +81,6 @@ class NDBEventTestCase(NDBMockTestBase):
             statement, dialect = fn(
                 mock.Mock(), mock.Mock(), statement, {}, mock.Mock(), False)
         self.assertEqual(statement, "ENGINE=NDBCLUSTER")
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
     def test_ndb_savepoint_override(self):
         test_engine = self.test_engine
@@ -86,7 +90,6 @@ class NDBEventTestCase(NDBMockTestBase):
                 mock.Mock(), mock.Mock(), statement, {}, mock.Mock(), False)
         self.assertEqual(statement,
                          "SET @oslo_db_ndb_savepoint_rollback_disabled = 0;")
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
     def test_ndb_rollback_override(self):
         test_engine = self.test_engine
@@ -96,7 +99,6 @@ class NDBEventTestCase(NDBMockTestBase):
                 mock.Mock(), mock.Mock(), statement, {}, mock.Mock(), False)
         self.assertEqual(statement,
                          "SET @oslo_db_ndb_savepoint_rollback_disabled = 0;")
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
     def test_ndb_rollback_release_override(self):
         test_engine = self.test_engine
@@ -106,30 +108,69 @@ class NDBEventTestCase(NDBMockTestBase):
                 mock.Mock(), mock.Mock(), statement, {}, mock.Mock(), False)
         self.assertEqual(statement,
                          "SET @oslo_db_ndb_savepoint_rollback_disabled = 0;")
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
 
 class NDBDatatypesTestCase(NDBMockTestBase):
-    def test_ndb_autostringtinytext(self):
+    def test_ndb_deprecated_autostringtinytext(self):
         test_engine = self.test_engine
         self.assertEqual("TINYTEXT",
                          str(ndb.AutoStringTinyText(255).compile(
                              dialect=test_engine.dialect)))
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
-    def test_ndb_autostringtext(self):
+    def test_ndb_deprecated_autostringtext(self):
         test_engine = self.test_engine
         self.assertEqual("TEXT",
                          str(ndb.AutoStringText(4096).compile(
                              dialect=test_engine.dialect)))
-        test_engine.dialect._oslodb_enable_ndb_support = False
 
-    def test_ndb_autostringsize(self):
+    def test_ndb_deprecated_autostringsize(self):
         test_engine = self.test_engine
         self.assertEqual('VARCHAR(64)',
                          str(ndb.AutoStringSize(255, 64).compile(
                              dialect=test_engine.dialect)))
-        test_engine.dialect._oslodb_enable_ndb_support = False
+
+    def test_ndb_string_to_tinytext(self):
+        test_engine = self.test_engine
+        self.assertEqual("TINYTEXT",
+                         str(String(255, mysql_ndb_type=TINYTEXT).compile(
+                             dialect=test_engine.dialect)))
+
+    def test_ndb_string_to_text(self):
+        test_engine = self.test_engine
+        self.assertEqual("TEXT",
+                         str(String(4096, mysql_ndb_type=TEXT).compile(
+                             dialect=test_engine.dialect)))
+
+    def test_ndb_string_length(self):
+        test_engine = self.test_engine
+        self.assertEqual('VARCHAR(64)',
+                         str(String(255, mysql_ndb_length=64).compile(
+                             dialect=test_engine.dialect)))
+
+
+class NDBDatatypesDefaultTestCase(NDBMockTestBase):
+    def setUp(self):
+        super(NDBMockTestBase, self).setUp()
+        mock_dbapi = mock.Mock()
+        self.test_engine = create_engine(_MOCK_CONNECTION, module=mock_dbapi)
+
+    def test_non_ndb_string_to_text(self):
+        test_engine = self.test_engine
+        self.assertEqual("VARCHAR(255)",
+                         str(String(255, mysql_ndb_type=TINYTEXT).compile(
+                             dialect=test_engine.dialect)))
+
+    def test_non_ndb_autostringtext(self):
+        test_engine = self.test_engine
+        self.assertEqual("VARCHAR(4096)",
+                         str(String(4096, mysql_ndb_type=TEXT).compile(
+                             dialect=test_engine.dialect)))
+
+    def test_non_ndb_autostringsize(self):
+        test_engine = self.test_engine
+        self.assertEqual('VARCHAR(255)',
+                         str(String(255, mysql_ndb_length=64).compile(
+                             dialect=test_engine.dialect)))
 
 
 class NDBOpportunisticTestCase(