minor edits

Author: rpiazza

stix2/datastore/relational_db/demo.py

@@ -0,0 +1,36 @@
+import datetime as dt
+
+from database_backends.postgres_backend import PostgresBackend
+import sys
+import json
+
+import stix2
+from stix2.datastore.relational_db.relational_db import RelationalDBStore
+import stix2.properties
+
+# needed so the relational db code knows to create tables for this
+from incident import incident, event, task, impact
+from identity_contact_information import identity_contact_information
+from observed_string import observed_string
+
+
+def main():
+    with open(sys.argv[1], "r") as f:
+        bundle = stix2.parse(json.load(f), allow_custom=True)
+    store = RelationalDBStore(
+        PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True),
+        True,
+        None,
+        True,
+        print_sql=True,
+    )
+
+    if store.sink.db_backend.database_exists:
+        for obj in bundle.objects:
+            store.add(obj)
+    else:
+        print("database does not exist")
+
+
+if __name__ == '__main__':
+    main()

stix2/datastore/relational_db/relational_db.py

@@ -12,7 +12,7 @@
 from stix2.parsing import parse
 
 
-def _add(store, stix_data, allow_custom=True, version="2.1"):
+def _add(sink, stix_data, allow_custom=True, version="2.1"):
     """Add STIX objects to MemoryStore/Sink.
 
     Adds STIX objects to an in-memory dictionary for fast lookup.
@@ -32,12 +32,12 @@ def _add(store, stix_data, allow_custom=True, version="2.1"):
     if isinstance(stix_data, list):
         # STIX objects are in a list- recurse on each object
         for stix_obj in stix_data:
-            _add(store, stix_obj, allow_custom, version)
+            _add(sink, stix_obj, allow_custom, version)
 
     elif stix_data["type"] == "bundle":
         # adding a json bundle - so just grab STIX objects
         for stix_obj in stix_data.get("objects", []):
-            _add(store, stix_obj, allow_custom, version)
+            _add(sink, stix_obj, allow_custom, version)
 
     else:
         # Adding a single non-bundle object
@@ -46,7 +46,7 @@ def _add(store, stix_data, allow_custom=True, version="2.1"):
         else:
             stix_obj = parse(stix_data, allow_custom, version)
 
-        store.insert_object(stix_obj)
+        sink.insert_object(stix_obj)
 
 
 class RelationalDBStore(DataStoreMixin):