flaky, put include all extensions

oasis-open · Apr 2, 2024 · b6c2276 · b6c2276
1 parent ac72888
commit b6c2276
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 20 deletions.
diff --git a/stix2/datastore/relational_db/relational_db.py b/stix2/datastore/relational_db/relational_db.py
@@ -55,7 +55,7 @@ def _add(store, stix_data, allow_custom=True, version="2.1"):
 class RelationalDBStore(DataStoreMixin):
     def __init__(
         self, database_connection_url, allow_custom=True, version=None,
-        instantiate_database=True, *stix_object_classes
+        instantiate_database=True, *stix_object_classes,
     ):
         """
         Initialize this store.
@@ -78,28 +78,28 @@ def __init__(
 
         self.metadata = MetaData()
         create_table_objects(
-            self.metadata, stix_object_classes
+            self.metadata, stix_object_classes,
         )
 
         super().__init__(
             source=RelationalDBSource(
                 database_connection,
-                metadata=self.metadata
+                metadata=self.metadata,
             ),
             sink=RelationalDBSink(
                 database_connection,
                 allow_custom=allow_custom,
                 version=version,
                 instantiate_database=instantiate_database,
-                metadata=self.metadata
+                metadata=self.metadata,
             ),
         )
 
 
 class RelationalDBSink(DataSink):
     def __init__(
         self, database_connection_or_url, allow_custom=True, version=None,
-        instantiate_database=True, *stix_object_classes, metadata=None
+        instantiate_database=True, *stix_object_classes, metadata=None,
     ):
         """
         Initialize this sink.  Only one of stix_object_classes and metadata
@@ -137,7 +137,7 @@ def __init__(
         else:
             self.metadata = MetaData()
             create_table_objects(
-                self.metadata, stix_object_classes
+                self.metadata, stix_object_classes,
             )
 
         self.allow_custom = allow_custom
@@ -164,7 +164,6 @@ def _instantiate_database(self):
     def generate_stix_schema(self):
         for t in self.metadata.tables.values():
             print(CreateTable(t).compile(self.database_connection))
-            print()
 
     def add(self, stix_data, version=None):
         _add(self, stix_data)
@@ -182,7 +181,7 @@ def insert_object(self, stix_object):
 
 class RelationalDBSource(DataSource):
     def __init__(
-        self, database_connection_or_url, *stix_object_classes, metadata=None
+        self, database_connection_or_url, *stix_object_classes, metadata=None,
     ):
         """
         Initialize this source.  Only one of stix_object_classes and metadata
@@ -215,15 +214,15 @@ def __init__(
         else:
             self.metadata = MetaData()
             create_table_objects(
-                self.metadata, stix_object_classes
+                self.metadata, stix_object_classes,
             )
 
     def get(self, stix_id, version=None, _composite_filters=None):
 
         stix_type = stix2.utils.get_type_from_id(stix_id)
         stix_class = stix2.registry.class_for_type(
             # TODO: give user control over STIX version used?
-            stix_type, stix_version=stix2.DEFAULT_VERSION
+            stix_type, stix_version=stix2.DEFAULT_VERSION,
         )
 
         # Info about the type-specific table

diff --git a/stix2/datastore/relational_db/relational_db_testing.py b/stix2/datastore/relational_db/relational_db_testing.py
@@ -3,9 +3,7 @@
 import pytz
 
 import stix2
-from stix2.datastore.relational_db.relational_db import (
-    RelationalDBSink, RelationalDBSource, RelationalDBStore
-)
+from stix2.datastore.relational_db.relational_db import RelationalDBStore
 import stix2.properties
 
 directory_stix_object = stix2.Directory(
@@ -103,7 +101,7 @@ def main():
         False,
         None,
         True,
-        stix2.Directory
+        stix2.Directory,
     )
     store.sink.generate_stix_schema()
 

diff --git a/stix2/datastore/relational_db/table_creation.py b/stix2/datastore/relational_db/table_creation.py
@@ -693,8 +693,8 @@ def create_table_objects(metadata, stix_object_classes):
                 stix_class,
                 metadata,
                 schema_name,
-                is_extension=is_extension
-            )
+                is_extension=is_extension,
+            ),
         )
 
     return tables
diff --git a/stix2/datastore/relational_db/utils.py b/stix2/datastore/relational_db/utils.py
@@ -1,10 +1,11 @@
 from collections.abc import Iterable, Mapping
+
 import inflection
+
 from stix2.v21.base import (
     _DomainObject, _Extension, _MetaObject, _Observable, _RelationshipObject,
 )
 
-
 # Helps us know which data goes in core, and which in a type-specific table.
 SCO_COMMON_PROPERTIES = {
     "id",
@@ -59,8 +60,8 @@ def get_stix_object_classes():
     yield from _get_all_subclasses(_MetaObject)
     # Non-object extensions (property or toplevel-property only)
     for ext_cls in _get_all_subclasses(_Extension):
-        if ext_cls.extension_type in (
-            "property-extension", "toplevel-property-extension"
+        if ext_cls.extension_type not in (
+            "new_sdo", "new_sco", "new_sro",
         ):
             yield ext_cls
 
@@ -104,7 +105,7 @@ def flat_classes(class_or_classes):
     if isinstance(class_or_classes, Iterable) and not isinstance(
         # Try to generically detect STIX objects, which are iterable, but we
         # don't want to iterate through those.
-        class_or_classes, Mapping
+        class_or_classes, Mapping,
     ):
         for class_ in class_or_classes:
             yield from flat_classes(class_)