Repository: oasis-open/cti-python-stix2
Pull request #599 (2024-07-11 14:19:07)

feat: add 'pretty' parameter to optimize JSON serialization performance
This pull request introduces an optimization to the `add` method of the `FileSystemSink` class by allowing the specification of the `pretty` parameter. The default setting `pretty=True` in the `fp_serialize` function negatively impacts the insertion performance of large STIX objects. In some cases, such as `XMitreCollection Enterprise ATT&CK`, the addition may fail, crashing the thread. Setting `pretty=False` mitigates this problem, significantly improving performance.

### Rationale

By allowing the `pretty` parameter to be specified in the `add` method, users can choose to disable "pretty" formatting when saving STIX objects, resulting in significant performance improvements, especially when dealing with large volumes of data.

### Performance Tests

1. **Performance**: Insertion with `pretty=False` is approximately twice as fast compared to `pretty=True`.
```bash
Tue Jul  9 20:34:07 2024    ./pretty_true.txt

         267530643 function calls (261758598 primitive calls) in 395.561 seconds

   Ordered by: internal time
   List reduced from 1060 to 10 due to restriction <10>

   ncalls       tottime  percall  cumtime  percall filename:lineno(function)
   157612       118.332    0.001  118.332    0.001 {method 'read' of '_ssl._SSLSocket' objects}
   615930       24.570    0.000   25.648    0.000 {built-in method io.open}
557607/185998   18.579    0.000  108.664    0.001 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/base.py:115(__init__)
  1231019       13.523    0.000   15.540    0.000 {built-in method posix.stat}
   932494       10.714    0.000   10.714    0.000 {method 'write' of '_io.TextIOWrapper' objects}
   186047       10.701    0.000   10.701    0.000 {built-in method posix.mkdir}
   859476       10.171    0.000   10.719    0.000 {method '__exit__' of '_io._IOBase' objects}
   185998       6.705    0.000   22.124    0.000 /openstix-python/.venv/default/lib/python3.10/site-packages/simplejson/encoder.py:306(iterencode)
   371996       5.351    0.000    9.555    0.000 /usr/lib/python3.10/_strptime.py:309(_strptime)
      391       5.151    0.013    5.151    0.013 {method 'do_handshake' of '_ssl._SSLSocket' objects}
```

**Results with `pretty=False`**:
```bash
Tue Jul  9 20:34:15 2024    ./pretty_false.txt

         125050663 function calls (118632644 primitive calls) in 117.531 seconds

   Ordered by: internal time
   List reduced from 1015 to 10 due to restriction <10>

   ncalls       tottime  percall  cumtime  percall filename:lineno(function)
   327985       13.913    0.000   14.160    0.000 {built-in method io.open}
   147785       13.138    0.000   13.138    0.000 {method 'read' of '_ssl._SSLSocket' objects}
   571531       5.123    0.000    5.692    0.000 {method '__exit__' of '_io._IOBase' objects}
   655506       5.101    0.000    5.879    0.000 {built-in method posix.stat}
126659/42220    4.598    0.000   26.854    0.001 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/base.py:115(__init__)
    42264       2.481    0.000    2.481    0.000 {built-in method posix.mkdir}
5826222/2237607 2.266    0.000   14.854    0.000 /openstix-python/.venv/default/lib/python3.10/site-packages/simplejson/encoder.py:677(_iterencode)
   243585       1.851    0.000   31.258    0.000 /usr/lib/python3.10/zipfile.py:1664(_extract_member)
   243856       1.676    0.000    1.676    0.000 {method 'decompress' of 'zlib.Decompress' objects}
1393227/464409  1.654    0.000    4.674    0.000 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/serialization.py:159(find_property_index)
```

2. **Testing with a large dataset: Using the dataset [mitre-atlas](https://raw.githubusercontent.com/mitre-atlas/atlas-navigator-data/main/dist/stix-atlas-attack-enterprise.json):

With `pretty=True`, the process was stopped after 301 seconds, generating 596K of incomplete data.
```bash
Wed Jul 10 20:49:26 2024    ./pretty_true_atlas.txt

         1644951185 function calls (1349693667 primitive calls) in 301.613 seconds

   Ordered by: internal time
   List reduced from 971 to 10 due to restriction <10>

   ncalls           tottime  percall  cumtime  percall filename:lineno(function)
220871233/18172     124.048    0.000  286.845    0.016 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/serialization.py:159(find_property_index)
590821359/590821355 52.875    0.000  105.371    0.000 {built-in method builtins.isinstance}
73587421/18160      42.840    0.000  286.805    0.016 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/serialization.py:137(_find_property_in_seq)
221211256           26.982    0.000   52.497    0.000 /usr/lib/python3.10/abc.py:117(__instancecheck__)
221211256           25.511    0.000   25.514    0.000 {built-in method _abc._abc_instancecheck}
220871233           10.677    0.000   10.677    0.000 {method 'isdigit' of 'str' objects}
 73551105           4.198    0.000    4.198    0.000 {method 'values' of 'dict' objects}
      924           3.480    0.004    3.480    0.004 {method 'read' of '_ssl._SSLSocket' objects}
  52483/1           1.613    0.000    9.912    9.912 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/base.py:115(__init__)
  1375284           0.605    0.000    1.170    0.000 /usr/lib/python3.10/collections/__init__.py:1000(__contains__)
```

With pretty=False, the process completed in 17 seconds, generating 90M of data.
```bash
Wed Jul 10 21:21:29 2024    ./pretty_false_atlas.txt

         31057698 function calls (30303084 primitive calls) in 17.029 seconds

   Ordered by: internal time
   List reduced from 979 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
  52925/2    1.599    0.000    9.891    4.946 /openstix-python/.venv/default/lib/python3.10/site-packages/stix2/base.py:115(__init__)
    21995    0.725    0.000    0.736    0.000 {built-in method io.open}
    22006    0.643    0.000    0.643    0.000 {built-in method posix.mkdir}
  1385596    0.607    0.000    1.174    0.000 /usr/lib/python3.10/collections/__init__.py:1000(__contains__)
    21991    0.602    0.000    2.389    0.000 /openstix-python/.venv/default/lib/python3.10/site-packages/simplejson/encoder.py:306(iterencode)
    66011    0.535    0.000    0.535    0.000 {built-in method posix.stat}
      959    0.523    0.001    0.523    0.001 {method 'read' of '_ssl._SSLSocket' objects}
   884643    0.465    0.000    0.796    0.000 /openstix-python/.venv/default/lib/python3.10/site-packages/simplejson/encoder.py:39(encode_basestring)
  1385596    0.438    0.000    1.807    0.000 /usr/lib/python3.10/collections/__init__.py:988(get)
    44034    0.373    0.000    0.635    0.000 /usr/lib/python3.10/_strptime.py:309(_strptime)
```

### Checklist

- [X] I have signed the Individual Contributor License Agreement (CLA)
- [X] I have read the [Contribution Guide](https://raw.githubusercontent.com/oasis-open/cti-python-stix2/50fd81fd6ba4f26824a864319305bc298e89bb45/CONTRIBUTING.md)
- [X] I ran the unit tests and they all passed
- [X] Added/updated documentation as necessary


Patch:
diff --git a/stix2/datastore/filesystem.py b/stix2/datastore/filesystem.py
index 2209116..d6f8327 100644
--- a/stix2/datastore/filesystem.py
+++ b/stix2/datastore/filesystem.py
@@ -554,7 +554,7 @@ class FileSystemSink(DataSink):
     def stix_dir(self):
         return self._stix_dir
 
-    def _check_path_and_write(self, stix_obj, encoding='utf-8'):
+    def _check_path_and_write(self, stix_obj, encoding='utf-8', pretty=True):
         """Write the given STIX object to a file in the STIX file directory.
         """
         type_dir = os.path.join(self._stix_dir, stix_obj["type"])
@@ -585,9 +585,9 @@ class FileSystemSink(DataSink):
             raise DataSourceError("Attempted to overwrite file (!) at: {}".format(file_path))
 
         with io.open(file_path, mode='w', encoding=encoding) as f:
-            fp_serialize(stix_obj, f, pretty=True, encoding=encoding, ensure_ascii=False)
+            fp_serialize(stix_obj, f, pretty=pretty, encoding=encoding, ensure_ascii=False)
 
-    def add(self, stix_data=None, version=None):
+    def add(self, stix_data=None, version=None, pretty=True):
         """Add STIX objects to file directory.
 
         Args:
@@ -597,6 +597,7 @@ class FileSystemSink(DataSink):
             version (str): If present, it forces the parser to use the version
                 provided. Otherwise, the library will make the best effort based
                 on checking the "spec_version" property.
+            pretty (bool): If True, the resulting JSON will be "pretty printed"
 
         Note:
             ``stix_data`` can be a Bundle object, but each object in it will be
@@ -607,24 +608,24 @@ class FileSystemSink(DataSink):
         if isinstance(stix_data, (v20.Bundle, v21.Bundle)):
             # recursively add individual STIX objects
             for stix_obj in stix_data.get("objects", []):
-                self.add(stix_obj, version=version)
+                self.add(stix_obj, version=version, pretty=pretty)
 
         elif isinstance(stix_data, _STIXBase):
             # adding python STIX object
-            self._check_path_and_write(stix_data)
+            self._check_path_and_write(stix_data, pretty=pretty)
 
         elif isinstance(stix_data, (str, dict)):
             parsed_data = parse(stix_data, allow_custom=self.allow_custom, version=version)
             if isinstance(parsed_data, _STIXBase):
-                self.add(parsed_data, version=version)
+                self.add(parsed_data, version=version, pretty=pretty)
             else:
                 # custom unregistered object type
-                self._check_path_and_write(parsed_data)
+                self._check_path_and_write(parsed_data, pretty=pretty)
 
         elif isinstance(stix_data, list):
             # recursively add individual STIX objects
             for stix_obj in stix_data:
-                self.add(stix_obj)
+                self.add(stix_obj, version=version, pretty=pretty)
 
         else:
             raise TypeError(


Test patch:
diff --git a/stix2/test/v21/test_datastore_filesystem.py b/stix2/test/v21/test_datastore_filesystem.py
index 3eb20b5..39b4557 100644
--- a/stix2/test/v21/test_datastore_filesystem.py
+++ b/stix2/test/v21/test_datastore_filesystem.py
@@ -151,6 +151,42 @@ def test_filesystem_source_bad_stix_file(fs_source, bad_stix_files):
     except STIXError as e:
         assert "Can't parse object with no 'type' property" in str(e)
 
+def test_filesystem_sink_add_pretty_true(fs_sink, fs_source):
+    """Test adding a STIX object with pretty=True."""
+    camp1 = stix2.v21.Campaign(
+        name="Hannibal",
+        objective="Targeting Italian and Spanish Diplomat internet accounts",
+        aliases=["War Elephant"],
+    )
+    fs_sink.add(camp1, pretty=True)
+    filepath = os.path.join(
+        FS_PATH, "campaign", camp1.id, _timestamp2filename(camp1.modified) + ".json",
+    )
+    assert os.path.exists(filepath)
+    with open(filepath, 'r', encoding='utf-8') as f:
+        content = f.read()
+        assert '\n' in content  # Check for pretty-printed output
+
+    os.remove(filepath)
+
+def test_filesystem_sink_add_pretty_false(fs_sink, fs_source):
+    """Test adding a STIX object with pretty=False."""
+    camp1 = stix2.v21.Campaign(
+        name="Hannibal",
+        objective="Targeting Italian and Spanish Diplomat internet accounts",
+        aliases=["War Elephant"],
+    )
+    fs_sink.add(camp1, pretty=False)
+    filepath = os.path.join(
+        FS_PATH, "campaign", camp1.id, _timestamp2filename(camp1.modified) + ".json",
+    )
+    assert os.path.exists(filepath)
+    with open(filepath, 'r', encoding='utf-8') as f:
+        content = f.read()
+        assert '\n' not in content  # Check for non-pretty-printed output
+
+    os.remove(filepath)
+
 
 def test_filesystem_source_get_object(fs_source):
     # get (latest) object