fix: updates json array registration

Author: chalmerlowe

db_dtypes/json.py

@@ -277,5 +277,10 @@ def to_pandas_dtype(self):
 
 
 # Register the type to be included in RecordBatches, sent over IPC and received in
-# another Python process.
-pa.register_extension_type(JSONArrowType())
+# another Python process. Also handle potential pre-registration
+try:
+    pa.register_extension_type(JSONArrowType())
+except pa.ArrowKeyError:
+    # Type 'dbjson' might already be registered if the module is reloaded,
+    # which is okay.
+    pass

tests/unit/test_json.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import sys
 
 import numpy as np
 import pandas as pd
@@ -224,3 +225,70 @@ def test_json_arrow_record_batch():
         == '{"null_field":null,"order":{"address":{"city":"Anytown","street":"123 Main St"},"items":["book","pen","computer"],"total":15}}'
     )
     assert s[6] == "null"
+
+
+@pytest.fixture
+def cleanup_json_module_for_reload():
+    """
+    Fixture to ensure db_dtypes.json is registered and then removed
+    from sys.modules to allow testing the registration except block via reload.
+    """
+
+    json_module_name = "db_dtypes.json"
+    original_module = sys.modules.get(json_module_name)
+
+    # Ensure the type is registered initially (usually by the first import)
+    try:
+        # Make sure the module is loaded so the type exists
+        import db_dtypes.json
+
+        # Explicitly register just in case it wasn't, or was cleaned up elsewhere.
+        # This might raise ArrowKeyError itself if already registered, which is fine here.
+        pa.register_extension_type(db_dtypes.json.JSONArrowType())
+    except pa.ArrowKeyError:
+        pass  # Already registered is the state we want before the test runs
+    except ImportError:
+        pytest.skip("Could not import db_dtypes.json to set up test.")
+
+    # Remove the module from sys.modules so importlib.reload re-executes it
+    if json_module_name in sys.modules:
+        del sys.modules[json_module_name]
+
+    yield  # Run the test that uses this fixture
+
+    # Cleanup: Put the original module back if it existed
+    # This helps isolate from other tests that might import db_dtypes.json
+    if original_module:
+        sys.modules[json_module_name] = original_module
+    elif json_module_name in sys.modules:
+        # If the test re-imported it but it wasn't there originally, remove it
+        del sys.modules[json_module_name]
+
+    # Note: PyArrow doesn't have a public API to unregister types easily,
+    # thus we are using the testing pattern of module isolation/reloading.
+
+
+def test_json_arrow_type_reregistration_is_handled(cleanup_json_module_for_reload):
+    """
+    Verify that attempting to re-register JSONArrowType via module reload
+    is caught by the except block and does not raise an error.
+    """
+
+    try:
+        # Re-importing the module after the fixture removed it from sys.modules
+        # forces Python to execute the module's top-level code again.
+        # This includes the pa.register_extension_type call.
+        import db_dtypes.json  # noqa: F401
+
+        assert (
+            True
+        ), "Module re-import completed without error, except block likely worked."
+
+    except pa.ArrowKeyError:
+        # If this exception escapes, the except block in db_dtypes/json.py failed.
+        pytest.fail(
+            "pa.ArrowKeyError was raised during module reload, "
+            "indicating the except block failed."
+        )
+    except Exception as e:
+        pytest.fail(f"An unexpected exception occurred during module reload: {e}")