more ARRAY

Author: rpiazza

stix2/datastore/relational_db/database_backends/postgres_backend.py

@@ -31,7 +31,7 @@ class PostgresBackend(DatabaseBackend):
         f"{os.getenv('POSTGRES_PORT', '5432')}/postgres"
 
     def __init__(self, database_connection_url=default_database_connection_url, force_recreate=False, **kwargs: Any):
-        super().__init__(database_connection_url, force_recreate=False, **kwargs)
+        super().__init__(database_connection_url, force_recreate=force_recreate, **kwargs)
 
     def _create_schemas(self):
         with self.database_connection.begin() as trans:

stix2/datastore/relational_db/input_creation.py

@@ -184,10 +184,10 @@ def generate_insert_information(self, name, stix_object, **kwargs):  # noqa: F81
 
 @add_method(ListProperty)
 def generate_insert_information(   # noqa: F811
-    self, name, stix_object, level=0, is_extension=False,
+    self, name, stix_object, data_sink=None, level=0, is_extension=False,
     foreign_key_value=None, schema_name=None, **kwargs,
 ):
-    data_sink = kwargs.get("data_sink")
+    db_backend = data_sink.db_backend
     table_name = kwargs.get("table_name")
     if isinstance(self.contained, ReferenceProperty):
         insert_statements = list()
@@ -226,7 +226,7 @@ def generate_insert_information(   # noqa: F811
     elif isinstance(self.contained, EmbeddedObjectProperty):
         insert_statements = list()
         for value in stix_object[name]:
-            with data_sink.database_connection.begin() as trans:
+            with db_backend.database_connection.begin() as trans:
                 next_id = trans.execute(data_sink.sequence)
             table = data_sink.tables_dictionary[canonicalize_table_name(table_name + "_" + name, schema_name)]
             bindings = {
@@ -248,7 +248,22 @@ def generate_insert_information(   # noqa: F811
             )
         return insert_statements
     else:
-        return {name: stix_object[name]}
+        if db_backend.array_allowed():
+            return {name: stix_object[name]}
+        else:
+            insert_statements = list()
+            table = data_sink.tables_dictionary[
+                canonicalize_table_name(
+                    table_name + "_" + name,
+                    schema_name,
+                )
+            ]
+            for elem in stix_object[name]:
+                bindings = {"id": stix_object["id"], name: elem}
+                insert_statements.append(insert(table).values(bindings))
+            return insert_statements
+
+
 
 
 @add_method(ReferenceProperty)
@@ -286,6 +301,7 @@ def generate_insert_for_array_in_table(table, values, foreign_key_value):
 
 
 def generate_insert_for_external_references(data_sink, stix_object):
+    db_backend = data_sink.db_backend
     insert_statements = list()
     next_id = None
     object_table = data_sink.tables_dictionary["common.external_references"]
@@ -295,7 +311,7 @@ def generate_insert_for_external_references(data_sink, stix_object):
             if prop in er:
                 bindings[prop] = er[prop]
         if "hashes" in er:
-            with data_sink.database_connection.begin() as trans:
+            with db_backend.database_connection.begin() as trans:
                 next_id = trans.execute(data_sink.sequence)
             bindings["hash_ref_id"] = next_id
         else:
@@ -326,17 +342,22 @@ def generate_insert_for_granular_markings(data_sink, granular_markings_table, st
             bindings["marking_ref"] = marking_ref_value
         if db_backend.array_allowed():
             bindings["selectors"] = granular_marking.get("selectors")
+            insert_statements.append(insert(granular_markings_table).values(bindings))
         else:
+            with db_backend.database_connection.begin() as trans:
+                next_id = trans.execute(data_sink.sequence)
+            bindings["selectors"] = next_id
+            insert_statements.append(insert(granular_markings_table).values(bindings))
             table = data_sink.tables_dictionary[
                 canonicalize_table_name(
-                    granular_markings_table + "_selector",
+                    granular_markings_table.name + "_selector",
                     db_backend.schema_for_core(),
                 )
             ]
             for sel in granular_marking.get("selectors"):
-                selector_bindings = {"id": stix_object["id"], "selector": sel}
+                selector_bindings = {"id": next_id, "selector": sel}
                 insert_statements.append(insert(table).values(selector_bindings))
-        insert_statements.append(insert(granular_markings_table).values(bindings))
+
     return insert_statements
 
 

stix2/datastore/relational_db/relational_db.py

@@ -57,7 +57,7 @@ def _add(store, stix_data, allow_custom=True, version="2.1"):
 class RelationalDBStore(DataStoreMixin):
     def __init__(
         self, db_backend, allow_custom=True, version=None,
-        instantiate_database=True, *stix_object_classes,
+        instantiate_database=True, print_sql=False, *stix_object_classes,
     ):
         """
         Initialize this store.
@@ -91,6 +91,7 @@ def __init__(
             ),
             sink=RelationalDBSink(
                 db_backend,
+                print_sql=print_sql,
                 allow_custom=allow_custom,
                 version=version,
                 instantiate_database=instantiate_database,
@@ -102,7 +103,7 @@ def __init__(
 class RelationalDBSink(DataSink):
     def __init__(
         self, db_backend, allow_custom=True, version=None,
-        instantiate_database=True, *stix_object_classes, metadata=None,
+        instantiate_database=True, print_sql=False, *stix_object_classes, metadata=None,
     ):
         """
         Initialize this sink.  Only one of stix_object_classes and metadata
@@ -134,15 +135,6 @@ def __init__(
 
         self.db_backend = db_backend
 
-        # if isinstance(database_connection_or_url, str):
-        #     self.database_connection = create_engine(database_connection_or_url)
-        # 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:
@@ -161,44 +153,21 @@ def __init__(
         if instantiate_database:
             if not self.db_backend.database_exists:
                 self.db_backend._create_database()
+            # else:
+            #     self.clear_tables()
             self.db_backend._create_schemas()
-            self._instantiate_database()
+            self._instantiate_database(print_sql)
 
-    # def _create_schemas(self):
-    #     with self.database_connection.begin() as trans:
-    #         trans.execute(CreateSchema("common", if_not_exists=True))
-    #         trans.execute(CreateSchema("sdo", if_not_exists=True))
-    #         trans.execute(CreateSchema("sco", if_not_exists=True))
-    #         trans.execute(CreateSchema("sro", if_not_exists=True))
-
-    def _instantiate_database(self):
+    def _instantiate_database(self, print_sql=False):
         self.metadata.create_all(self.db_backend.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.db_backend.database_connection))
+        if print_sql:
+            for t in self.metadata.tables.values():
+                print(CreateTable(t).compile(self.db_backend.database_connection))
 
     def add(self, stix_data, version=None):
         _add(self, stix_data)
     add.__doc__ = _add.__doc__
 
-    # @staticmethod
-    # def _determine_schema_name(stix_object):
-    #     if isinstance(stix_object, _DomainObject):
-    #         return "sdo"
-    #     elif isinstance(stix_object, _Observable):
-    #         return "sco"
-    #     elif isinstance(stix_object, _RelationshipObject):
-    #         return "sro"
-    #     elif isinstance(stix_object, _MetaObject):
-    #         return "common"
-
     def insert_object(self, stix_object):
         schema_name = self.db_backend.determine_schema_name(stix_object)
         with self.db_backend.database_connection.begin() as trans:

stix2/datastore/relational_db/relational_db_testing.py

@@ -254,15 +254,16 @@ def custom_obj():
 
 def main():
     store = RelationalDBStore(
-        PostgresBackend("postgresql://localhost/stix-data-sink"),
+        PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True),
         False,
         None,
         True,
+        print_sql=True,
     )
 
     if store.sink.db_backend.database_exists:
-        store.sink.generate_stix_schema()
-        store.sink.clear_tables()
+        # store.sink.generate_stix_schema()
+        # store.sink.clear_tables()
 
         # td = test_dictionary()
 
@@ -295,8 +296,8 @@ def main():
         dict_example = dictionary_test()
         store.add(dict_example)
 
-        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")
 

stix2/datastore/relational_db/table_creation.py

@@ -32,12 +32,12 @@ def aux_table_property(prop, name, core_properties):
         return False
 
 
-def create_array_column(property_name, contained_sql_type):
+def create_array_column(property_name, contained_sql_type, optional):
     return Column(
         property_name,
         ARRAY(contained_sql_type),
-        CheckConstraint(f"array_length({property_name}, 1) IS NOT NULL"),
-        nullable=False,
+        CheckConstraint(f"{property_name} IS NULL or array_length({property_name}, 1) IS NOT NULL"),
+        nullable=optional
     )
 
 
@@ -47,12 +47,12 @@ def create_array_child_table(metadata, db_backend, table_name, property_name, co
         Column(
             "id",
             db_backend.determine_sql_type_for_key_as_id(),
-            ForeignKey(
-                canonicalize_table_name(table_name, schema_name) + ".id",
-                ondelete="CASCADE",
-            ),
+            # ForeignKey(
+            #     canonicalize_table_name(table_name, schema_name) + ".id",
+            #     ondelete="CASCADE",
+            # ),
             nullable=False,
-            primary_key=True,
+            #primary_key=True,
         ),
         Column(
             property_name,
@@ -178,7 +178,7 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo):
             db_backend.determine_sql_type_for_key_as_id(),
             ForeignKey(canonicalize_table_name("core_" + sco_or_sdo, schema_name) + ".id", ondelete="CASCADE"),
             nullable=False,
-            primary_key=True,
+            # primary_key=not db_backend.array_allowed(),
         ),
         Column("lang", db_backend.determine_sql_type_for_string_property()),
         Column(
@@ -191,35 +191,61 @@ def create_granular_markings_table(metadata, db_backend, sco_or_sdo):
         ),
     ]
     if db_backend.array_allowed():
-        child_table = False
-        columns.append(create_array_column("selectors", db_backend.determine_sql_type_for_string_property()))
-    else:
-        child_table = True
-
-    tables = [
-        Table(
-            canonicalize_table_name("granular_marking_" + sco_or_sdo),
-            metadata,
-            *columns,
-            CheckConstraint(
-            """(lang IS NULL AND marking_ref IS NOT NULL)
-                  OR
-                  (lang IS NOT NULL AND marking_ref IS NULL)""",
-            ),
-            schema=schema_name
-        ),
-    ]
-    if child_table:
-        tables.append(
-            create_array_child_table(
+        columns.append(create_array_column("selectors", db_backend.determine_sql_type_for_string_property(), False))
+        return [
+            Table(
+                "granular_marking_" + sco_or_sdo,
                 metadata,
-                db_backend,
+                *columns,
+                CheckConstraint(
+                """(lang IS NULL AND marking_ref IS NOT NULL)
+                      OR
+                      (lang IS NOT NULL AND marking_ref IS NULL)""",
+                ),
+                schema=schema_name
+            ),
+        ]
+    else:
+        columns.append(Column(
+                "selectors",
+                db_backend.determine_sql_type_for_key_as_int(),
+                unique=True
+            )
+        )
+        tables = [
+            Table(
                 "granular_marking_" + sco_or_sdo,
+                metadata,
+                *columns,
+                CheckConstraint(
+                    """(lang IS NULL AND marking_ref IS NOT NULL)
+                          OR
+                          (lang IS NOT NULL AND marking_ref IS NULL)""",
+                ),
+                schema=schema_name
+            ),
+        ]
+        schema_name = db_backend.schema_for_core()
+        columns = [
+            Column(
+                "id",
+                db_backend.determine_sql_type_for_key_as_int(),
+                ForeignKey(
+                    canonicalize_table_name("granular_marking_" + sco_or_sdo, schema_name) + ".selectors",
+                    ondelete="CASCADE",
+                ),
+                # primary_key=True,
+                nullable=False,
+            ),
+            Column(
                 "selector",
                 db_backend.determine_sql_type_for_string_property(),
+                nullable=False,
             ),
-        )
-    return tables
+        ]
+        tables.append(Table(canonicalize_table_name("granular_marking_" + sco_or_sdo + "_" + "selector"),
+                            metadata, *columns, schema=schema_name))
+        return tables
 
 
 def create_external_references_tables(metadata, db_backend):
@@ -275,7 +301,7 @@ def create_core_table(metadata, db_backend, schema_name):
         ]
         columns.extend(sdo_columns)
         if db_backend.array_allowed():
-            columns.append(create_array_column("labels", db_backend.determine_sql_type_for_string_property())),
+            columns.append(create_array_column("labels", db_backend.determine_sql_type_for_string_property(), True))
     else:
         columns.append(Column("defanged", db_backend.determine_sql_type_for_boolean_property(), default=False))
 

stix2/datastore/relational_db/utils.py

@@ -148,7 +148,7 @@ def determine_sql_type_from_stix(cls_or_inst, db_backend):  # noqa: F811
     elif is_class_or_instance(cls_or_inst, IntegerProperty):
         return db_backend.determine_sql_type_for_integer_property()
     elif is_class_or_instance(cls_or_inst, StringProperty):
-        return db_backend.determine_sql_type_for_integer_property()
+        return db_backend.determine_sql_type_for_string_property()
     elif is_class_or_instance(cls_or_inst, ReferenceProperty):
         db_backend.determine_sql_type_for_reference_property()
     elif is_class_or_instance(cls_or_inst, TimestampProperty):