From 6beee679bb2b074cee3d37ab0c1eff1f0eff0142 Mon Sep 17 00:00:00 2001 From: Ming Fang Date: Wed, 4 Dec 2024 00:09:23 -0500 Subject: [PATCH 1/3] Initial changes to support SQLite backend. --- .../database_backend_base.py | 3 ++ .../database_backends/sqlite_backend.py | 44 +++++++++++++++++++ .../relational_db/relational_db_testing.py | 4 +- 3 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 stix2/datastore/relational_db/database_backends/sqlite_backend.py diff --git a/stix2/datastore/relational_db/database_backends/database_backend_base.py b/stix2/datastore/relational_db/database_backends/database_backend_base.py index e5082451..4cf46f1b 100644 --- a/stix2/datastore/relational_db/database_backends/database_backend_base.py +++ b/stix2/datastore/relational_db/database_backends/database_backend_base.py @@ -18,6 +18,9 @@ def __init__(self, database_connection_url, force_recreate=False, **kwargs: Any) self.database_connection = create_engine(database_connection_url) + def _fk_pragma_on_connect(self): + self.database_connnection.execute('pragma foreign_keys=ON') + def _create_database(self): if self.database_exists: drop_database(self.database_connection_url) diff --git a/stix2/datastore/relational_db/database_backends/sqlite_backend.py b/stix2/datastore/relational_db/database_backends/sqlite_backend.py new file mode 100644 index 00000000..ae5473bf --- /dev/null +++ b/stix2/datastore/relational_db/database_backends/sqlite_backend.py @@ -0,0 +1,44 @@ +import os +from typing import Any + +from sqlalchemy import TIMESTAMP, LargeBinary, Text +from sqlalchemy import event + +from stix2.base import ( + _DomainObject, _MetaObject, _Observable, _RelationshipObject, +) +from stix2.datastore.relational_db.utils import schema_for + +from .database_backend_base import DatabaseBackend + + +class SQLiteBackend(DatabaseBackend): + default_database_connection_url = f"sqlite:///stix-data-sink.db" + + def __init__(self, database_connection_url=default_database_connection_url, force_recreate=False, **kwargs: Any): + super().__init__(database_connection_url, force_recreate=force_recreate, **kwargs) + + event.listen(self.database_connection, 'connect', self._fk_pragma_on_connect) + + # ========================================================================= + # sql type methods (overrides) + + @staticmethod + def determine_sql_type_for_binary_property(): # noqa: F811 + return SQLiteBackend.determine_sql_type_for_string_property() + + @staticmethod + def determine_sql_type_for_hex_property(): # noqa: F811 + # return LargeBinary + return SQLiteBackend.determine_sql_type_for_string_property() + + @staticmethod + def determine_sql_type_for_timestamp_property(): # noqa: F811 + return TIMESTAMP(timezone=True) + + # ========================================================================= + # Other methods + + @staticmethod + def array_allowed(): + return False diff --git a/stix2/datastore/relational_db/relational_db_testing.py b/stix2/datastore/relational_db/relational_db_testing.py index 54bf0aae..acd0727a 100644 --- a/stix2/datastore/relational_db/relational_db_testing.py +++ b/stix2/datastore/relational_db/relational_db_testing.py @@ -1,6 +1,7 @@ import datetime as dt from database_backends.postgres_backend import PostgresBackend +from database_backends.sqlite_backend import SQLiteBackend import pytz import stix2 @@ -287,7 +288,8 @@ def test_dictionary(): def main(): store = RelationalDBStore( - PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True), + #PostgresBackend("postgresql://localhost/stix-data-sink", force_recreate=True), + SQLiteBackend("sqlite:///stix-data-sink.db", force_recreate=True), True, None, True, From 28603d73f901e339f5f63cc74eb048dc37d6323a Mon Sep 17 00:00:00 2001 From: Ming Fang Date: Wed, 4 Dec 2024 17:37:02 -0500 Subject: [PATCH 2/3] Enable foreign key constraint for SQLite directly from the SQLiteBackend subclass. --- .../database_backends/database_backend_base.py | 3 --- .../relational_db/database_backends/sqlite_backend.py | 7 ++++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/stix2/datastore/relational_db/database_backends/database_backend_base.py b/stix2/datastore/relational_db/database_backends/database_backend_base.py index 4cf46f1b..e5082451 100644 --- a/stix2/datastore/relational_db/database_backends/database_backend_base.py +++ b/stix2/datastore/relational_db/database_backends/database_backend_base.py @@ -18,9 +18,6 @@ def __init__(self, database_connection_url, force_recreate=False, **kwargs: Any) self.database_connection = create_engine(database_connection_url) - def _fk_pragma_on_connect(self): - self.database_connnection.execute('pragma foreign_keys=ON') - def _create_database(self): if self.database_exists: drop_database(self.database_connection_url) diff --git a/stix2/datastore/relational_db/database_backends/sqlite_backend.py b/stix2/datastore/relational_db/database_backends/sqlite_backend.py index ae5473bf..f8094a15 100644 --- a/stix2/datastore/relational_db/database_backends/sqlite_backend.py +++ b/stix2/datastore/relational_db/database_backends/sqlite_backend.py @@ -2,6 +2,7 @@ from typing import Any from sqlalchemy import TIMESTAMP, LargeBinary, Text +from sqlalchemy.engine import Engine from sqlalchemy import event from stix2.base import ( @@ -18,7 +19,11 @@ class SQLiteBackend(DatabaseBackend): def __init__(self, database_connection_url=default_database_connection_url, force_recreate=False, **kwargs: Any): super().__init__(database_connection_url, force_recreate=force_recreate, **kwargs) - event.listen(self.database_connection, 'connect', self._fk_pragma_on_connect) + set_sqlite_pragma(self) + + @event.listens_for(Engine, "connect") + def set_sqlite_pragma(self): + self.database_connection.execute("PRAGMA foreign_keys=ON") # ========================================================================= # sql type methods (overrides) From 2ac14aec75b61791a9a4244fb69b69ac7543084c Mon Sep 17 00:00:00 2001 From: Ming Fang Date: Thu, 5 Dec 2024 12:04:56 -0500 Subject: [PATCH 3/3] Use SQLAlchemy's event.listens_for() to ensure that foreign key constraint is enforced on new connections before use. --- .../database_backends/sqlite_backend.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/stix2/datastore/relational_db/database_backends/sqlite_backend.py b/stix2/datastore/relational_db/database_backends/sqlite_backend.py index f8094a15..545f928b 100644 --- a/stix2/datastore/relational_db/database_backends/sqlite_backend.py +++ b/stix2/datastore/relational_db/database_backends/sqlite_backend.py @@ -2,7 +2,6 @@ from typing import Any from sqlalchemy import TIMESTAMP, LargeBinary, Text -from sqlalchemy.engine import Engine from sqlalchemy import event from stix2.base import ( @@ -19,11 +18,14 @@ class SQLiteBackend(DatabaseBackend): def __init__(self, database_connection_url=default_database_connection_url, force_recreate=False, **kwargs: Any): super().__init__(database_connection_url, force_recreate=force_recreate, **kwargs) - set_sqlite_pragma(self) - - @event.listens_for(Engine, "connect") - def set_sqlite_pragma(self): - self.database_connection.execute("PRAGMA foreign_keys=ON") + @event.listens_for(self.database_connection, "connect") + def set_sqlite_pragma(dbapi_connection, connection_record): + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + result = cursor.execute("PRAGMA foreign_keys") + for row in result: + print('PRAGMA foreign_keys:', row) + cursor.close() # ========================================================================= # sql type methods (overrides)