Check with fallback descriptorDB for FindExtensionByNumber()/FindAllExtensions in UPB python pool.

PiperOrigin-RevId: 740001906

Author: anandolee

python/descriptor_pool.c

@@ -146,6 +146,61 @@ static bool PyUpb_DescriptorPool_TryLoadFilename(PyUpb_DescriptorPool* self,
   return ret;
 }
 
+static bool PyUpb_DescriptorPool_TryLoadExtension(PyUpb_DescriptorPool* self,
+                                                  const upb_MessageDef* m,
+                                                  int32_t field_number) {
+  if (!self->db) return false;
+  const char* full_name = upb_MessageDef_FullName(m);
+  PyObject* py_name = PyUnicode_FromStringAndSize(full_name, strlen(full_name));
+  PyObject* py_descriptor = PyObject_CallMethod(
+      self->db, "FindFileContainingExtension", "Oi", py_name, field_number);
+  Py_DECREF(py_name);
+  if (!py_descriptor) {
+    PyErr_Clear();
+    return false;
+  }
+  bool ret = PyUpb_DescriptorPool_TryLoadFileProto(self, py_descriptor);
+  Py_DECREF(py_descriptor);
+  return ret;
+}
+
+static void PyUpb_DescriptorPool_TryLoadAllExtensions(
+    PyUpb_DescriptorPool* self, const upb_MessageDef* m) {
+  if (!self->db) return;
+  const char* full_name = upb_MessageDef_FullName(m);
+  PyObject* py_name = PyUnicode_FromStringAndSize(full_name, strlen(full_name));
+  PyObject* py_list =
+      PyObject_CallMethod(self->db, "FindAllExtensionNumbers", "O", py_name);
+  Py_DECREF(py_name);
+  if (!py_list) {
+    PyErr_Clear();
+    return;
+  }
+  Py_ssize_t size = PyList_Size(py_list);
+  if (size == -1) {
+    PyErr_Format(
+        PyExc_RuntimeError,
+        "FindAllExtensionNumbers() on fall back DB must return a list, not %S",
+        py_list);
+    PyErr_Print();
+    Py_DECREF(py_list);
+    return;
+  }
+  int64_t field_number;
+  const upb_ExtensionRegistry* reg =
+      upb_DefPool_ExtensionRegistry(self->symtab);
+  const upb_MiniTable* t = upb_MessageDef_MiniTable(m);
+  for (Py_ssize_t i = 0; i < size; ++i) {
+    PyObject* item = PySequence_GetItem(py_list, i);
+    field_number = PyLong_AsLong(item);
+    Py_DECREF(item);
+    if (!upb_ExtensionRegistry_Lookup(reg, t, field_number)) {
+      PyUpb_DescriptorPool_TryLoadExtension(self, m, field_number);
+    }
+  }
+  Py_DECREF(py_list);
+}
+
 bool PyUpb_DescriptorPool_CheckNoDatabase(PyObject* _self) { return true; }
 
 static bool PyUpb_DescriptorPool_LoadDependentFiles(
@@ -582,8 +637,15 @@ static PyObject* PyUpb_DescriptorPool_FindExtensionByNumber(PyObject* _self,
     return NULL;
   }
 
-  const upb_FieldDef* f = upb_DefPool_FindExtensionByNumber(
-      self->symtab, PyUpb_Descriptor_GetDef(message_descriptor), number);
+  const upb_MessageDef* message_def =
+      PyUpb_Descriptor_GetDef(message_descriptor);
+  const upb_FieldDef* f =
+      upb_DefPool_FindExtensionByNumber(self->symtab, message_def, number);
+  if (f == NULL && self->db) {
+    if (PyUpb_DescriptorPool_TryLoadExtension(self, message_def, number)) {
+      f = upb_DefPool_FindExtensionByNumber(self->symtab, message_def, number);
+    }
+  }
   if (f == NULL) {
     return PyErr_Format(PyExc_KeyError, "Couldn't find Extension %d", number);
   }
@@ -595,6 +657,9 @@ static PyObject* PyUpb_DescriptorPool_FindAllExtensions(PyObject* _self,
                                                         PyObject* msg_desc) {
   PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
   const upb_MessageDef* m = PyUpb_Descriptor_GetDef(msg_desc);
+  if (self->db) {
+    PyUpb_DescriptorPool_TryLoadAllExtensions(self, m);
+  }
   size_t n;
   const upb_FieldDef** ext = upb_DefPool_GetAllExtensions(self->symtab, m, &n);
   PyObject* ret = PyList_New(n);

python/google/protobuf/descriptor_pool.py

@@ -580,11 +580,21 @@ def FindAllExtensions(self, message_descriptor):
     if self._descriptor_db and hasattr(
         self._descriptor_db, 'FindAllExtensionNumbers'):
       full_name = message_descriptor.full_name
-      all_numbers = self._descriptor_db.FindAllExtensionNumbers(full_name)
-      for number in all_numbers:
-        if number in self._extensions_by_number[message_descriptor]:
-          continue
-        self._TryLoadExtensionFromDB(message_descriptor, number)
+      try:
+        all_numbers = self._descriptor_db.FindAllExtensionNumbers(full_name)
+      except:
+        pass
+      else:
+        if isinstance(all_numbers, list):
+          for number in all_numbers:
+            if number in self._extensions_by_number[message_descriptor]:
+              continue
+            self._TryLoadExtensionFromDB(message_descriptor, number)
+        else:
+          warnings.warn(
+              'FindAllExtensionNumbers() on fall back DB must return a list,'
+              ' not {0}'.format(type(all_numbers))
+          )
 
     return list(self._extensions_by_number[message_descriptor].values())
 
@@ -603,8 +613,13 @@ def _TryLoadExtensionFromDB(self, message_descriptor, number):
       return
 
     full_name = message_descriptor.full_name
-    file_proto = self._descriptor_db.FindFileContainingExtension(
-        full_name, number)
+    file_proto = None
+    try:
+      file_proto = self._descriptor_db.FindFileContainingExtension(
+          full_name, number
+      )
+    except:
+      return
 
     if file_proto is None:
       return
@@ -668,7 +683,6 @@ def SetFeatureSetDefaults(self, defaults):
     if not isinstance(defaults, descriptor_pb2.FeatureSetDefaults):
       raise TypeError('SetFeatureSetDefaults called with invalid type')
 
-
     if defaults.minimum_edition > defaults.maximum_edition:
       raise ValueError(
           'Invalid edition range %s to %s'

python/google/protobuf/internal/descriptor_pool_test.py

@@ -1531,5 +1531,91 @@ def testAboveMaximum(self):
         'google/protobuf/internal/more_messages.proto'])
 
 
+class LocalFakeDB(descriptor_database.DescriptorDatabase):
+
+  def FindFileContainingExtension(self, extendee_name, extension_number):
+    return descriptor_pb2.FileDescriptorProto.FromString(
+        factory_test2_pb2.DESCRIPTOR.serialized_pb
+    )
+
+  def FindAllExtensionNumbers(self, extendee_name):
+    return [1001, 1002]
+
+
+class BadDB(descriptor_database.DescriptorDatabase):
+
+  def FindFileContainingExtension(self, extendee_name, extension_number):
+    raise RuntimeError('just ignore it')
+
+  def FindAllExtensionNumbers(self, extendee_name):
+    raise RuntimeError('just ignore it')
+
+
+class BadDB2(descriptor_database.DescriptorDatabase):
+
+  # Returns a none list value.
+  def FindAllExtensionNumbers(self, extendee_name):
+    return 1.2
+
+
+@testing_refleaks.TestCase
+class FallBackDBTest(unittest.TestCase):
+
+  def setUp(self):
+    self.factory_test1_fd = descriptor_pb2.FileDescriptorProto.FromString(
+        factory_test1_pb2.DESCRIPTOR.serialized_pb
+    )
+    factory_test2_fd = descriptor_pb2.FileDescriptorProto.FromString(
+        factory_test2_pb2.DESCRIPTOR.serialized_pb
+    )
+    db = LocalFakeDB()
+    db.Add(self.factory_test1_fd)
+    db.Add(factory_test2_fd)
+    self.pool = descriptor_pool.DescriptorPool(db)
+    file_desc = self.pool.FindFileByName(
+        'google/protobuf/internal/factory_test1.proto'
+    )
+    self.message_desc = file_desc.message_types_by_name['Factory1Message']
+
+    bad_db = BadDB()
+    bad_db.Add(self.factory_test1_fd)
+    self.bad_pool = descriptor_pool.DescriptorPool(bad_db)
+
+  def testFindExtensionByNumber(self):
+    ext = self.pool.FindExtensionByNumber(self.message_desc, 1001)
+    self.assertEqual(ext.name, 'one_more_field')
+
+  def testFindAllExtensions(self):
+    extensions = self.pool.FindAllExtensions(self.message_desc)
+    self.assertEqual(len(extensions), 2)
+
+  def testIgnoreBadFindExtensionByNumber(self):
+    file_desc = self.bad_pool.FindFileByName(
+        'google/protobuf/internal/factory_test1.proto'
+    )
+    message_desc = file_desc.message_types_by_name['Factory1Message']
+    with self.assertRaises(KeyError):
+      ext = self.bad_pool.FindExtensionByNumber(message_desc, 1001)
+
+  def testIgnoreBadFindAllExtensions(self):
+    file_desc = self.bad_pool.FindFileByName(
+        'google/protobuf/internal/factory_test1.proto'
+    )
+    message_desc = file_desc.message_types_by_name['Factory1Message']
+    extensions = self.bad_pool.FindAllExtensions(message_desc)
+    self.assertEqual(len(extensions), 0)
+
+  def testFindAllExtensionsReturnsNoneList(self):
+    db = BadDB2()
+    db.Add(self.factory_test1_fd)
+    pool = descriptor_pool.DescriptorPool(db)
+    file_desc = pool.FindFileByName(
+        'google/protobuf/internal/factory_test1.proto'
+    )
+    message_desc = file_desc.message_types_by_name['Factory1Message']
+    extensions = self.bad_pool.FindAllExtensions(message_desc)
+    self.assertEqual(len(extensions), 0)
+
+
 if __name__ == '__main__':
   unittest.main()

python/google/protobuf/pyext/descriptor_database.cc

@@ -146,6 +146,14 @@ bool PyDescriptorDatabase::FindAllExtensionNumbers(
     return false;
   }
   Py_ssize_t size = PyList_Size(py_list.get());
+  if (size == -1) {
+    PyErr_Format(
+        PyExc_RuntimeError,
+        "FindAllExtensionNumbers() on fall back DB must return a list, not %S",
+        py_list.get());
+    PyErr_Print();
+    return false;
+  }
   int64_t item_value;
   for (Py_ssize_t i = 0 ; i < size; ++i) {
     ScopedPyObjectPtr item(PySequence_GetItem(py_list.get(), i));