add better implementation for sqlite sequence

Author: rpiazza

stix2/datastore/relational_db/database_backends/database_backend_base.py

@@ -1,7 +1,7 @@
 from typing import Any
 
 from sqlalchemy import (
-    Boolean, CheckConstraint, Float, Integer, String, Text, create_engine,
+    Boolean, CheckConstraint, Float, Integer, Sequence, String, Text, create_engine,
 )
 from sqlalchemy_utils import create_database, database_exists, drop_database
 
@@ -147,3 +147,6 @@ def process_value_for_insert(self, stix_type, value):
     def next_id(self, data_sink):
         with self.database_connection.begin() as trans:
             return trans.execute(data_sink.sequence)
+
+    def create_sequence(self, metadata):
+        return Sequence("my_general_seq", metadata=metadata, start=1, schema=self.schema_for_core())

stix2/datastore/relational_db/database_backends/sqlite_backend.py

@@ -1,6 +1,7 @@
 from typing import Any
 
-from sqlalchemy import Text, event
+from sqlalchemy import Table, Column, Text, event, insert, select, update
+from sqlalchemy.schema import CreateTable
 
 from .database_backend_base import DatabaseBackend
 
@@ -10,6 +11,8 @@ class SQLiteBackend(DatabaseBackend):
 
     temp_sequence_count = 0
 
+    select_stmt_for_sequence = None
+
     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)
 
@@ -60,8 +63,35 @@ def array_allowed():
     def create_regex_constraint_clause(self, column_name, pattern):
         return f"{column_name} REGEXP {pattern}"
 
-    @staticmethod
-    def next_id(data_sink):
+    def next_id(self, data_sink):
         # hack, which is not reliable, must look for a better solution
-        SQLiteBackend.temp_sequence_count += 1
-        return SQLiteBackend.temp_sequence_count
+        # SQLiteBackend.temp_sequence_count += 1
+        # return SQLiteBackend.temp_sequence_count
+        # ALWAYS CALL WITHIN A TRANSACTION?????
+
+        # better solution, but probably not best
+        value = 0
+        conn = self.database_connection.connect()
+        for row in conn.execute(self.select_stmt_for_sequence):
+            value = row[0]
+        if value == 0:
+            stmt = insert(data_sink.sequence).values({"id": 1, "value": 0})
+            conn = self.database_connection.connect()
+            conn.execute(stmt)
+            conn.commit()
+        value += 1
+        conn.execute(update(data_sink.sequence).where(data_sink.sequence.c.id == 1).values({"value": value}))
+        conn.commit()
+        return value
+
+    def create_sequence(self, metadata):
+        # need id column, so update has something to work with (see above)
+        t = Table("my_general_seq",
+                  metadata,
+                  Column("id", SQLiteBackend.determine_sql_type_for_key_as_int()),
+                  Column("value", SQLiteBackend.determine_sql_type_for_integer_property()),
+                  schema=self.schema_for_core())
+        CreateTable(t).compile(self.database_connection)
+        self.select_stmt_for_sequence = select(t.c.value)
+        return t
+

stix2/datastore/relational_db/relational_db.py

@@ -1,5 +1,5 @@
 from sqlalchemy import MetaData, delete
-from sqlalchemy.schema import CreateTable, Sequence
+from sqlalchemy.schema import CreateTable
 
 from stix2.base import _STIXBase
 from stix2.datastore import DataSink, DataSource, DataStoreMixin
@@ -138,7 +138,6 @@ def __init__(
             create_table_objects(
                 self.metadata, stix_object_classes,
             )
-        self.sequence = Sequence("my_general_seq", metadata=self.metadata, start=1, schema=db_backend.schema_for_core())
 
         self.allow_custom = allow_custom
 
@@ -155,6 +154,7 @@ def __init__(
             self._instantiate_database(print_sql)
 
     def _instantiate_database(self, print_sql=False):
+        self.sequence = self.db_backend.create_sequence(self.metadata)
         self.metadata.create_all(self.db_backend.database_connection)
         if print_sql:
             for t in self.metadata.tables.values():