FB8-265: main.innodb_pk_extension_on fail in 8.0.23 after porting fro…

…m 8.0.20

https://jira.percona.com/browse/FB8-265

Failing MTR tests:
main.innodb_pk_extension_on
main.innodb_pk_extension_off

Reason:
8.0.23 reworked logic in in dd_table_share.cc::prepare_share().
In particular commit ff4d6b9 fixed
add_pk_parts_to_sk() to not add PK parts to SK if the key size exceeds
MAX_REF_PARTS limit.
However, commit ff4d6b9 moved the part
of setting initial value of field->part_of_key after the call to
add_pk_parts_to_sk(). This resulted in wrong field->part_of_key value.

Solution:
Partially revert fix done in commit
ff4d6b9.
This is not needed now, as we introduced 'goto restart' logic in commit
8dae15d

Failing MTR test:
rocksdb.rocksdb_parts

Reason:
For partitioned tables, the handler_file inside
dd_table_share.cc::prepare_share() points to ha_rockspart object which
inherits from Partition_base. The method init_with_fields() is not
overrode however, and the implementation provided by Partition_base does
nothing if handler is not initialized (underlying ha_rocksdb handlers
are not created/initialized).
ha_rockspart handler is opened after share is prepared, but it is too
late, as we need to know if HA_PRIMARY_KEY_IN_READ_INDEX is set inside
prepare_share()

Solution:
Override implementation of init_with_fields() in ha_rockspart.
We take the advantage of the fact that RocksDB sets
HA_PRIMARY_KEY_IN_READ_INDEX flag if primary key is present, so we can
make the decision in this scope before creation of underlying partitions
handlers.
This way the flow inside dd_table_share.cc::prepare_share() is able to
determine proper state of HA_PRIMARY_KEY_IN_READ_INDEX for ha_rockspart
handler as well.

sql/dd_table_share.cc

@@ -433,9 +433,20 @@ static bool prepare_share(THD *thd, TABLE_SHARE *share,
         if (key == primary_key) {
           field->set_flag(PRI_KEY_FLAG);
           /*
-            "if (ha_option & HA_PRIMARY_KEY_IN_READ_INDEX)" ... was moved below
-            for MyRocks
+            We need to set field->part_of_key before add_pk_parts_to_sk()
+            as calling it may clean key maps that exceed sizes.
           */
+          /*
+             If this field is part of the primary key and all keys contains
+             the primary key, then we can use any key to find this column
+             */
+          if (ha_option & HA_PRIMARY_KEY_IN_READ_INDEX) {
+            if (field->key_length() == key_part->length &&
+                !field->is_flag_set(BLOB_FLAG))
+              field->part_of_key = share->keys_in_use;
+            if (field->part_of_sortkey.is_set(key))
+              field->part_of_sortkey = share->keys_in_use;
+          }
         }
         if (field->key_length() != key_part->length) {
 #ifndef TO_BE_DELETED_ON_PRODUCTION
@@ -509,24 +520,6 @@ static bool prepare_share(THD *thd, TABLE_SHARE *share,
       ++idx_it;
     }
 
-    if (primary_key < MAX_KEY &&
-        (handler_file->ha_table_flags() & HA_PRIMARY_KEY_IN_READ_INDEX)) {
-      keyinfo = &share->key_info[primary_key];
-      key_part = keyinfo->key_part;
-      for (uint i = 0; i < keyinfo->user_defined_key_parts; key_part++, i++) {
-        Field *field = key_part->field;
-        /*
-          If this field is part of the primary key and all keys contains
-          the primary key, then we can use any key to find this column
-        */
-        if (field->key_length() == key_part->length &&
-            !field->is_flag_set(BLOB_FLAG))
-          field->part_of_key = share->keys_in_use;
-        if (field->part_of_sortkey.is_set(primary_key))
-          field->part_of_sortkey = share->keys_in_use;
-      }
-    }
-
     if (share->primary_key != MAX_KEY) {
       /*
          If we are using an integer as the primary key then allow the user to

sql/table.cc

@@ -2143,9 +2143,20 @@ static int open_binary_frm(THD *thd, TABLE_SHARE *share,
         if (key == primary_key) {
           field->set_flag(PRI_KEY_FLAG);
           /*
-            "if (ha_option & HA_PRIMARY_KEY_IN_READ_INDEX)" ... was moved below
-            for MyRocks
+            We need to set field->part_of_key before add_pk_parts_to_sk()
+            as calling it may clean key maps that exceed sizes.
           */
+          /*
+            If this field is part of the primary key and all keys contains
+            the primary key, then we can use any key to find this column
+          */
+          if (ha_option & HA_PRIMARY_KEY_IN_READ_INDEX) {
+            if (field->key_length() == key_part->length &&
+                !field->is_flag_set(BLOB_FLAG))
+              field->part_of_key = share->keys_in_use;
+            if (field->part_of_sortkey.is_set(key))
+              field->part_of_sortkey = share->keys_in_use;
+          }
         }
         if (field->key_length() != key_part->length) {
           if (field->type() == MYSQL_TYPE_NEWDECIMAL) {
@@ -2200,24 +2211,6 @@ static int open_binary_frm(THD *thd, TABLE_SHARE *share,
             std::max(share->max_unique_length, keyinfo->key_length);
     }
 
-    if (primary_key < MAX_KEY &&
-        (handler_file->ha_table_flags() & HA_PRIMARY_KEY_IN_READ_INDEX)) {
-      keyinfo = &share->key_info[primary_key];
-      key_part = keyinfo->key_part;
-      for (i = 0; i < keyinfo->user_defined_key_parts; key_part++, i++) {
-        Field *field = key_part->field;
-        /*
-          If this field is part of the primary key and all keys contains
-          the primary key, then we can use any key to find this column
-        */
-        if (field->key_length() == key_part->length &&
-            !field->is_flag_set(BLOB_FLAG))
-          field->part_of_key = share->keys_in_use;
-        if (field->part_of_sortkey.is_set(primary_key))
-          field->part_of_sortkey = share->keys_in_use;
-      }
-    }
-
     if (share->primary_key != MAX_KEY) {
       /*
         If we are using an integer as the primary key then allow the user to

storage/rocksdb/ha_rockspart.cc

@@ -21,6 +21,24 @@ namespace myrocks {
 extern handlerton *rocksdb_hton;
 }
 
+bool ha_rockspart::init_with_fields() {
+  /*
+    MySql layer calls this before initializing actual partitions and then makes
+    decisions basing on the result of updated table_flags.
+    As RocksDB enables HA_PRIMARY_KEY_IN_READ_INDEX flag if PK is present and
+    disables it otherwise, we can cache this flag earlier, even if partitions
+    handlers are not created yet.
+  */
+  const uint pk = table_share->primary_key;
+  if (pk != MAX_KEY) {
+    cached_table_flags |= HA_PRIMARY_KEY_IN_READ_INDEX;
+  } else {
+    cached_table_flags &= ~HA_PRIMARY_KEY_IN_READ_INDEX;
+  }
+
+  return Partition_base::init_with_fields();
+}
+
 handler *ha_rockspart::get_file_handler(TABLE_SHARE *share,
                                         MEM_ROOT *alloc) const {
   ha_rocksdb *file = new (alloc) ha_rocksdb(myrocks::rocksdb_hton, share);

storage/rocksdb/ha_rockspart.h

@@ -42,6 +42,8 @@ class ha_rockspart : public native_part::Partition_base {
     return ROW_TYPE_NOT_USED;
   }
 
+  bool init_with_fields() override;
+
  private:
   handler *get_file_handler(TABLE_SHARE *share, MEM_ROOT *alloc) const override;
   handler *clone(const char *name, MEM_ROOT *mem_root) override;