Merge pull request #593 from robthew/master

Added database existence checks and clear_tables function

stix2/datastore/relational_db/relational_db.py

@@ -1,5 +1,6 @@
-from sqlalchemy import MetaData, create_engine, select
+from sqlalchemy import MetaData, create_engine, select, delete
 from sqlalchemy.schema import CreateSchema, CreateTable, Sequence
+from sqlalchemy_utils import database_exists, create_database, drop_database
 
 from stix2.base import _STIXBase
 from stix2.datastore import DataSink, DataSource, DataStoreMixin
@@ -55,7 +56,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, force_recreate=False, *stix_object_classes,
     ):
         """
         Initialize this store.
@@ -67,6 +68,8 @@ def __init__(
             version: TODO: unused so far
             instantiate_database: Whether tables, etc should be created in the
                 database (only necessary the first time)
+            force_recreate: Drops old database and creates new one (useful if
+                the schema has changed and the tables need to be updated)
             *stix_object_classes: STIX object classes to map into table schemas
                 (and ultimately database tables, if instantiation is desired).
                 This can be used to limit which table schemas are created, if
@@ -91,6 +94,7 @@ def __init__(
                 allow_custom=allow_custom,
                 version=version,
                 instantiate_database=instantiate_database,
+                force_recreate=force_recreate,
                 metadata=self.metadata,
             ),
         )
@@ -99,7 +103,7 @@ def __init__(
 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, force_recreate=False, *stix_object_classes, metadata=None,
     ):
         """
         Initialize this sink.  Only one of stix_object_classes and metadata
@@ -111,8 +115,10 @@ def __init__(
             allow_custom: Whether custom content is allowed when processing
                 dict content to be added to the sink
             version: TODO: unused so far
-            instantiate_database: Whether tables, etc should be created in the
-                database (only necessary the first time)
+            instantiate_database: Whether the database, tables, etc should be
+                created (only necessary the first time)
+            force_recreate: Drops old database and creates new one (useful if
+                the schema has changed and the tables need to be updated)
             *stix_object_classes: STIX object classes to map into table schemas
                 (and ultimately database tables, if instantiation is desired).
                 This can be used to limit which table schemas are created, if
@@ -132,6 +138,10 @@ def __init__(
         else:
             self.database_connection = database_connection_or_url
 
+        self.database_exists = database_exists(self.database_connection.url)
+        if force_recreate:
+            self._create_database()
+
         if metadata:
             self.metadata = metadata
         else:
@@ -148,6 +158,8 @@ def __init__(
             self.tables_dictionary[canonicalize_table_name(t.name, t.schema)] = t
 
         if instantiate_database:
+            if not self.database_exists:
+                self._create_database()
             self._create_schemas()
             self._instantiate_database()
 
@@ -161,6 +173,12 @@ def _create_schemas(self):
     def _instantiate_database(self):
         self.metadata.create_all(self.database_connection)
 
+    def _create_database(self):
+        if self.database_exists:
+            drop_database(self.database_connection.url)
+        create_database(self.database_connection.url)
+        self.database_exists = database_exists(self.database_connection.url)
+
     def generate_stix_schema(self):
         for t in self.metadata.tables.values():
             print(CreateTable(t).compile(self.database_connection))
@@ -178,6 +196,14 @@ def insert_object(self, stix_object):
                 trans.execute(stmt)
             trans.commit()
 
+    def clear_tables(self):
+        tables = list(reversed(self.metadata.sorted_tables))
+        with self.database_connection.begin() as trans:
+            for table in tables:
+                delete_stmt = delete(table)
+                print(f'delete_stmt: {delete_stmt}')
+                trans.execute(delete_stmt)
+
 
 class RelationalDBSource(DataSource):
     def __init__(

stix2/datastore/relational_db/relational_db_testing.py

@@ -22,13 +22,13 @@
 )
 
 s = stix2.v21.Software(
-        id="software--28897173-7314-4eec-b1cf-2c625b635bf6",
-        name="Word",
-        cpe="cpe:2.3:a:microsoft:word:2000:*:*:*:*:*:*:*",
-        swid="com.acme.rms-ce-v4-1-5-0",
-        version="2002",
-        languages=["c", "lisp"],
-        vendor="Microsoft",
+    id="software--28897173-7314-4eec-b1cf-2c625b635bf6",
+    name="Word",
+    cpe="cpe:2.3:a:microsoft:word:2000:*:*:*:*:*:*:*",
+    swid="com.acme.rms-ce-v4-1-5-0",
+    version="2002",
+    languages=["c", "lisp"],
+    vendor="Microsoft",
 )
 
 
@@ -101,14 +101,20 @@ def main():
         False,
         None,
         True,
+        False,
         stix2.Directory,
     )
-    store.sink.generate_stix_schema()
 
-    store.add(directory_stix_object)
+    if store.sink.database_exists:
+        store.sink.generate_stix_schema()
+        store.sink.clear_tables()
+
+        store.add(directory_stix_object)
 
-    read_obj = store.get(directory_stix_object.id)
-    print(read_obj)
+        read_obj = store.get(directory_stix_object.id)
+        print(read_obj)
+    else:
+        print("database does not exist")
 
 
 if __name__ == '__main__':