Skip to content

Commit

Permalink
flaky
Browse files Browse the repository at this point in the history
  • Loading branch information
@rpiazza
rpiazza committed Mar 26, 2024
1 parent b896339 commit 80c8734
Show file tree
Hide file tree
Showing 23 changed files with 317 additions and 277 deletions.
8 changes: 4 additions & 4 deletions .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,22 @@
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v3.4.0
rev: v4.5.0
hooks:
- id: trailing-whitespace
- id: check-merge-conflict
- repo: https://github.com/asottile/add-trailing-comma
rev: v2.0.2
rev: v3.1.0
hooks:
- id: add-trailing-comma
- repo: https://github.com/PyCQA/flake8
rev: 3.8.4
rev: 7.0.0
hooks:
- id: flake8
name: Check project styling
args:
- --max-line-length=160
- repo: https://github.com/PyCQA/isort
rev: 5.12.0
rev: 5.13.2
hooks:
- id: isort
name: Sort python imports (shows diff)
Expand Down
129 changes: 71 additions & 58 deletions stix2/datastore/relational_db/input_creation.py
View file
Original file line number Diff line number Diff line change
@@ -1,28 +1,25 @@
from collections import OrderedDict

from sqlalchemy import (
TIMESTAMP,
CheckConstraint,
Column,
ForeignKey,
Table,
Text,
create_engine,
insert,
)

from sqlalchemy import insert

from stix2.datastore.relational_db.utils import (
SCO_COMMON_PROPERTIES, SDO_COMMON_PROPERTIES, canonicalize_table_name,
)
from stix2.properties import (
DictionaryProperty, EmbeddedObjectProperty, EnumProperty, ExtensionsProperty, FloatProperty,
IntegerProperty, ListProperty, ReferenceProperty, StringProperty)

from stix2.datastore.relational_db.utils import SCO_COMMON_PROPERTIES, SDO_COMMON_PROPERTIES, canonicalize_table_name
DictionaryProperty, EmbeddedObjectProperty, EnumProperty,
ExtensionsProperty, FloatProperty, IntegerProperty, ListProperty,
ReferenceProperty, StringProperty,
)


def single_value(p):
return not(isinstance(p, (EmbeddedObjectProperty,
ListProperty,
DictionaryProperty)))
return not isinstance(
p, (
EmbeddedObjectProperty,
ListProperty,
DictionaryProperty,
),
)


def table_property(prop, name, core_properties):
Expand Down Expand Up @@ -59,10 +56,10 @@ def derive_column_name(prop):
return "value"


def generate_insert_for_array_in_table(table, property_name, values, prop, foreign_key_value):
def generate_insert_for_array_in_table(table, values, foreign_key_value):

bindings = {
"id": foreign_key_value
"id": foreign_key_value,
}

for idx, item in enumerate(values):
Expand All @@ -76,24 +73,20 @@ def generate_insert_for_array_in_table(table, property_name, values, prop, forei
def generate_single_values(stix_object, properties, core_properties=[]):
bindings = OrderedDict()
for name, prop in properties.items():
if (single_value(prop) and (name == 'id' or name not in core_properties) or
array_property(prop, name, core_properties)):
if (
single_value(prop) and (name == 'id' or name not in core_properties) or
array_property(prop, name, core_properties)
):
if name in stix_object and name != "type":
bindings[name] = stix_object[name] if not array_property(prop, name, core_properties) else "{" + ",".join(
['"' + x + '"' for x in stix_object[name]]) + "}"
['"' + x + '"' for x in stix_object[name]],
) + "}"
return bindings


def generate_insert_for_embedded_object(type_name, item, foreign_key_value):
bindings = generate_single_values(item, item._properties)
bindings["id"] = foreign_key_value
sql = f"INSERT INTO {canonicalize_table_name(type_name, item._type)}" \
f" ({','.join(bindings.keys())})" \
f" VALUES ({','.join(values)}, %(id)s )"

print("sql:", sql)
print("embedded:", bindings)
return [(sql, bindings)]


def generate_insert_for_dictionary(item, dictionary_table, foreign_key_value, value_types):
Expand Down Expand Up @@ -147,17 +140,21 @@ def generate_insert_for_external_references(data_sink, stix_object):

if "hashes" in er:
hashes_table = data_sink.tables_dictionary[canonicalize_table_name("external_references_hashes", "sdo")]
insert_statements.extend(generate_insert_for_hashes(er["hashes"],
hashes_table,
stix_object["id"]))
insert_statements.extend(
generate_insert_for_hashes(
er["hashes"],
hashes_table,
stix_object["id"],
),
)

return insert_statements


def generate_insert_for_granular_markings(data_sink, stix_object, granular_markings_table):
def generate_insert_for_granular_markings(granular_markings_table, stix_object):
granular_markings = stix_object["granular_markings"]
bindings = {
"id": stix_object["id"]
"id": stix_object["id"],
}
for idx, granular_marking in enumerate(granular_markings):
lang_binding_name = f"lang{idx}"
Expand All @@ -171,15 +168,19 @@ def generate_insert_for_granular_markings(data_sink, stix_object, granular_marki
return [insert(granular_markings_table).values(bindings)]


def generate_insert_for_extensions(extensions, foreign_key_value, type_name, core_properties):
sql_bindings_tuples = list()
for name, ex in extensions.items():
sql_bindings_tuples.extend(generate_insert_for_subtype_extension(name,
ex,
foreign_key_value,
type_name,
core_properties))
return sql_bindings_tuples
# def generate_insert_for_extensions(extensions, foreign_key_value, type_name, core_properties):
# sql_bindings_tuples = list()
# for name, ex in extensions.items():
# sql_bindings_tuples.extend(
# generate_insert_for_subtype_extension(
# name,
# ex,
# foreign_key_value,
# type_name,
# core_properties,
# ),
# )
# return sql_bindings_tuples


def generate_insert_for_core(data_sink, stix_object, core_properties, schema_name):
Expand All @@ -205,20 +206,26 @@ def generate_insert_for_core(data_sink, stix_object, core_properties, schema_nam
object_markings_ref_table = data_sink.tables_dictionary["common.object_marking_refs_sdo"]
else:
object_markings_ref_table = data_sink.tables_dictionary["common.object_marking_refs_sco"]
insert_statements.extend(generate_insert_for_array_in_table(data_sink, stix_object, object_markings_ref_table))
insert_statements.extend(
generate_insert_for_array_in_table(
object_markings_ref_table,
stix_object["object_marking_refs"],
stix_object["id"],
),
)

# Granular markings
if "granular_markings" in stix_object:
if schema_name == "sdo":
granular_marking_table = data_sink.tables_dictionary["common.granular_marking_sdo"]
else:
granular_marking_table = data_sink.tables_dictionary["common.granular_marking_sco"]
granular_input_statements = generate_insert_for_granular_markings(data_sink,
stix_object.granular_markings,
granular_marking_table)
granular_input_statements = generate_insert_for_granular_markings(
granular_marking_table,
stix_object.granular_markings,
)
insert_statements.extend(granular_input_statements)


return insert_statements


Expand Down Expand Up @@ -255,16 +262,22 @@ def generate_insert_for_object(data_sink, stix_object, schema_name, foreign_key_
if table_property(prop, name, core_properties):
if name in stix_object:
if embedded_object_list_property(prop, name, core_properties):
insert_statements.extend(generate_insert_for_embedded_objects(name,
stix_object[name],
stix_object["id"]))
insert_statements.extend(
generate_insert_for_embedded_objects(
name,
stix_object[name],
stix_object["id"],
),
)
elif isinstance(prop, ExtensionsProperty):
pass
else:
insert_statements.extend(generate_insert_for_array_in_table(stix_object["type"],
name,
stix_object[name],
properties[name],
stix_object["id"] ))
insert_statements.extend(
generate_insert_for_array_in_table(
stix_object["type"],
name,
stix_object[name],
properties[name],
stix_object["id"], ),
)
return insert_statements

24 changes: 19 additions & 5 deletions stix2/datastore/relational_db/relational_db.py
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,18 @@
from sqlalchemy import MetaData, create_engine
from sqlalchemy.schema import CreateTable
from sqlalchemy.schema import CreateSchema, CreateTable

from stix2.base import _STIXBase
from stix2.datastore import DataSink
from stix2.datastore.relational_db.table_creation import create_core_tables, generate_object_table
from stix2.datastore.relational_db.input_creation import generate_insert_for_object

from stix2.datastore.relational_db.input_creation import (
generate_insert_for_object,
)
from stix2.datastore.relational_db.table_creation import (
create_core_tables, generate_object_table,
)
from stix2.parsing import parse
from stix2.v21.base import (_DomainObject, _Extension, _Observable, _RelationshipObject,)
from stix2.v21.base import (
_DomainObject, _Extension, _Observable, _RelationshipObject,
)


def _get_all_subclasses(cls):
Expand Down Expand Up @@ -89,6 +94,8 @@ def __init__(
self.metadata = MetaData()
self.database_connection = create_engine(database_connection_url)

self._create_schemas()

self.tables = self._create_table_objects()
self.tables_dictionary = dict()
for t in self.tables:
Expand All @@ -97,6 +104,13 @@ def __init__(
if instantiate_database:
self._instantiate_database()

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 _create_table_objects(self):
tables = create_core_tables(self.metadata)
for stix_class in _get_all_subclasses(_DomainObject):
Expand Down
Loading

0 comments on commit 80c8734

Please sign in to comment.