Add support for partitioning and sample pages on mysql reflected options

diff --git a/lib/sqlalchemy/dialects/mysql/reflection.py b/lib/sqlalchemy/dialects/mysql/reflection.py
index 5be6a010e9d056d35741773bb412c47ac0ed51f4..8aa33dafdb442816163c476b22d11bb2573ceacd 100644 (file)
--- a/lib/sqlalchemy/dialects/mysql/reflection.py
+++ b/lib/sqlalchemy/dialects/mysql/reflection.py
@@ -52,6 +52,8 @@ class MySQLTableDefinitionParser(object):
                 pass
             elif line.startswith("CREATE "):
                 self._parse_table_name(line, state)
+            elif "PARTITION" in line:
+                self._parse_partition_options(line, state)
             # Not present in real reflection, but may be if
             # loading from a file.
             elif not line:
@@ -160,6 +162,52 @@ class MySQLTableDefinitionParser(object):
         for opt, val in options.items():
             state.table_options["%s_%s" % (self.dialect.name, opt)] = val
 
+    def _parse_partition_options(self, line, state):
+        options = {}
+        rest_of_line = line[:]
+
+        while rest_of_line[0] == "(" or rest_of_line[0] == " ":
+            rest_of_line = rest_of_line[1:]
+
+        for regex, cleanup in self._pr_options:
+            m = regex.search(rest_of_line)
+            if not m or "PARTITION" not in regex.pattern:
+                continue
+
+            directive = m.group("directive")
+            directive = directive.replace(" ", "_").lower()
+            is_subpartition = directive == "subpartition"
+
+            if directive == "partition" or is_subpartition:
+                rest_of_line = rest_of_line.replace(") */", "")
+                rest_of_line = rest_of_line.replace(",", "")
+                if is_subpartition and rest_of_line[-1] == ")":
+                    rest_of_line = rest_of_line[:-1]
+                defs = "mysql_%s_definitions" % (directive)
+                options[defs] = rest_of_line
+            else:
+                value = m.group("val")
+                if cleanup:
+                    value = cleanup(value)
+                options[directive] = value
+            break
+
+        for opt, val in options.items():
+            if (
+                opt == "mysql_partition_definitions"
+                or opt == "mysql_subpartition_definitions"
+            ):
+                # builds a string of definitions
+                if opt not in state.table_options:
+                    state.table_options[opt] = val
+                else:
+                    state.table_options[opt] = "%s, %s" % (
+                        state.table_options[opt],
+                        val,
+                    )
+            else:
+                state.table_options["%s_%s" % (self.dialect.name, opt)] = val
+
     def _parse_column(self, line, state):
         """Extract column details.
 
@@ -487,9 +535,20 @@ class MySQLTableDefinitionParser(object):
             "PACK_KEYS",
             "ROW_FORMAT",
             "KEY_BLOCK_SIZE",
+            "STATS_SAMPLE_PAGES",
         ):
             self._add_option_word(option)
 
+        for option in (
+            "PARTITION BY",
+            "SUBPARTITION BY",
+            "PARTITIONS",
+            "SUBPARTITIONS",
+            "PARTITION",
+            "SUBPARTITION",
+        ):
+            self._add_partition_option_word(option)
+
         self._add_option_regex("UNION", r"\([^\)]+\)")
         self._add_option_regex("TABLESPACE", r".*? STORAGE DISK")
         self._add_option_regex(
@@ -517,6 +576,21 @@ class MySQLTableDefinitionParser(object):
         )
         self._pr_options.append(_pr_compile(regex))
 
+    def _add_partition_option_word(self, directive):
+        if directive == "PARTITION BY" or directive == "SUBPARTITION BY":
+            regex = r"(?<!\S)(?P<directive>%s)%s" r"(?P<val>\w+.*)" % (
+                re.escape(directive),
+                self._optional_equals,
+            )
+        elif directive == "SUBPARTITIONS" or directive == "PARTITIONS":
+            regex = r"(?<!\S)(?P<directive>%s)%s" r"(?P<val>\d+)" % (
+                re.escape(directive),
+                self._optional_equals,
+            )
+        else:
+            regex = r"(?<!\S)(?P<directive>%s)(?!\S)" % (re.escape(directive),)
+        self._pr_options.append(_pr_compile(regex))
+
     def _add_option_regex(self, directive, regex):
         regex = r"(?P<directive>%s)%s" r"(?P<val>%s)" % (
             re.escape(directive),

diff --git a/test/dialect/mysql/test_reflection.py b/test/dialect/mysql/test_reflection.py
index 026025a88fb2c583d708c8d0d83a8e967acdb989..d285453c7002e6e4d7ee2e810d11b12e47774f34 100644 (file)
--- a/test/dialect/mysql/test_reflection.py
+++ b/test/dialect/mysql/test_reflection.py
@@ -321,6 +321,7 @@ class ReflectionTest(fixtures.TestBase, AssertsCompiledSQL):
                 mysql_avg_row_length="3",
                 mysql_password="secret",
                 mysql_connection="fish",
+                mysql_stats_sample_pages="4",
             )
 
         def_table = Table(
@@ -367,6 +368,7 @@ class ReflectionTest(fixtures.TestBase, AssertsCompiledSQL):
             assert def_table.kwargs["mysql_avg_row_length"] == "3"
             assert def_table.kwargs["mysql_password"] == "secret"
             assert def_table.kwargs["mysql_connection"] == "fish"
+            assert def_table.kwargs["mysql_stats_sample_pages"] == "4"
 
             assert reflected.kwargs["mysql_engine"] == "MEMORY"
 
@@ -375,6 +377,7 @@ class ReflectionTest(fixtures.TestBase, AssertsCompiledSQL):
             assert reflected.kwargs["mysql_default charset"] == "utf8"
             assert reflected.kwargs["mysql_avg_row_length"] == "3"
             assert reflected.kwargs["mysql_connection"] == "fish"
+            assert reflected.kwargs["mysql_stats_sample_pages"] == "4"
 
             # This field doesn't seem to be returned by mysql itself.
             # assert reflected.kwargs['mysql_password'] == 'secret'
@@ -382,6 +385,139 @@ class ReflectionTest(fixtures.TestBase, AssertsCompiledSQL):
             # This is explicitly ignored when reflecting schema.
             # assert reflected.kwargs['mysql_auto_increment'] == '5'
 
+    def test_reflection_with_partition_options(self):
+        kwargs = dict(
+            mysql_engine="InnoDB",
+            mysql_default_charset="utf8",
+            mysql_partition_by="HASH(MONTH(c2))",
+            mysql_partitions="6",
+        )
+
+        def_table = Table(
+            "mysql_def",
+            MetaData(testing.db),
+            Column("c1", Integer()),
+            Column("c2", DateTime),
+            **kwargs
+        )
+
+        def_table.create()
+        try:
+            reflected = Table("mysql_def", MetaData(testing.db), autoload=True)
+        finally:
+            def_table.drop()
+
+        assert def_table.kwargs["mysql_partition_by"] == "HASH(MONTH(c2))"
+        assert def_table.kwargs["mysql_partitions"] == "6"
+
+        assert reflected.kwargs["mysql_partition_by"] == "HASH (month(`c2`))"
+        assert reflected.kwargs["mysql_partitions"] == "6"
+
+    def test_reflection_with_subpartition_options(self):
+
+        subpartititon_text = """HASH (TO_DAYS (c2))
+                                SUBPARTITIONS 2(
+                                 PARTITION p0 VALUES LESS THAN (1990),
+                                 PARTITION p1 VALUES LESS THAN (2000),
+                                 PARTITION p2 VALUES LESS THAN MAXVALUE
+                             );"""
+        kwargs = dict(
+            mysql_engine="InnoDB",
+            mysql_default_charset="utf8",
+            mysql_partition_by="RANGE(YEAR(c2))",
+            mysql_subpartition_by=subpartititon_text,
+        )
+
+        def_table = Table(
+            "mysql_def",
+            MetaData(testing.db),
+            Column("c1", Integer()),
+            Column("c2", DateTime),
+            **kwargs
+        )
+
+        def_table.create()
+        try:
+            reflected = Table("mysql_def", MetaData(testing.db), autoload=True)
+        finally:
+            def_table.drop()
+
+        assert def_table.kwargs["mysql_partition_by"] == "RANGE(YEAR(c2))"
+
+        ref_kwargs = reflected.kwargs
+        assert ref_kwargs["mysql_partition_by"] == "RANGE (year(`c2`))"
+        assert ref_kwargs["mysql_subpartition_by"] == "HASH (to_days(`c2`))"
+        assert ref_kwargs["mysql_subpartitions"] == "2"
+
+        opts = reflected.dialect_options["mysql"]
+        expected = (
+            "PARTITION p0 VALUES LESS THAN (1990) ENGINE = InnoDB,"
+            " PARTITION p1 VALUES LESS THAN (2000) ENGINE = InnoDB,"
+            " PARTITION p2 VALUES LESS THAN MAXVALUE ENGINE = InnoDB"
+        )
+        assert opts["partition_definitions"] == expected
+
+    def test_reflection_with_subpartition_options_two(self):
+        partititon_text = """RANGE (YEAR (c2))
+                                SUBPARTITION BY HASH( TO_DAYS(c2))(
+                                 PARTITION p0 VALUES LESS THAN (1990)(
+                                 SUBPARTITION s0,
+                                 SUBPARTITION s1
+                                 ),
+                                 PARTITION p1 VALUES LESS THAN (2000)(
+                                 SUBPARTITION s2,
+                                 SUBPARTITION s3
+                                 ),
+                                PARTITION p2 VALUES LESS THAN MAXVALUE(
+                                 SUBPARTITION s4,
+                                 SUBPARTITION s5
+                                 )
+                             );"""
+
+        kwargs = dict(
+            mysql_engine="InnoDB",
+            mysql_default_charset="utf8",
+            mysql_partition_by=partititon_text,
+        )
+
+        def_table = Table(
+            "mysql_def",
+            MetaData(testing.db),
+            Column("c1", Integer()),
+            Column("c2", DateTime),
+            **kwargs
+        )
+
+        def_table.create()
+        try:
+            reflected = Table("mysql_def", MetaData(testing.db), autoload=True)
+        finally:
+            def_table.drop()
+
+        assert def_table.kwargs["mysql_partition_by"] == partititon_text
+
+        ref_kwargs = reflected.kwargs
+        assert ref_kwargs["mysql_partition_by"] == "RANGE (year(`c2`))"
+        assert ref_kwargs["mysql_subpartition_by"] == "HASH (to_days(`c2`))"
+
+        opts = reflected.dialect_options["mysql"]
+        part_expected = (
+            "PARTITION p0 VALUES LESS THAN (1990),"
+            " PARTITION p1 VALUES LESS THAN (2000),"
+            " PARTITION p2 VALUES LESS THAN MAXVALUE"
+        )
+        subpart_expected = (
+            "SUBPARTITION s0 ENGINE = InnoDB,"
+            " SUBPARTITION s1 ENGINE = InnoDB,"
+            " SUBPARTITION s2 ENGINE = InnoDB,"
+            " SUBPARTITION s3 ENGINE = InnoDB,"
+            " SUBPARTITION s4 ENGINE = InnoDB,"
+            " SUBPARTITION s5 ENGINE = InnoDB"
+        )
+
+        assert opts["partition_definitions"] == part_expected
+        assert opts["subpartition_definitions"] == subpart_expected
+
     def test_reflection_on_include_columns(self):
         """Test reflection of include_columns to be sure they respect case."""