From 5eff6d4ead105faaec62aa257953715deb89e480 Mon Sep 17 00:00:00 2001
From: jbirddog <100367399+jbirddog@users.noreply.github.com>
Date: Fri, 9 Jun 2023 15:28:59 -0400
Subject: [PATCH] Typeahead datastore (#321)

* Wiring up the datastore

* Writes into the data store

* Bulk save needs the timestamps

* Prep to do the local query

* Local typeahead working

* Pre pr cleanup

* ignore migrations dir in pre-commit for ruff w/ burnettk

* Getting ./bin/pyl to pass

---------

Co-authored-by: jasquat <jasquat@users.noreply.github.com>
---
 .pre-commit-config.yaml                       |  1 +
 .../migrations/versions/377be1608b45_.py      | 44 +++++++++++++
 .../data_stores/typeahead.py                  | 65 +++++++++++++++++++
 .../load_database_models.py                   |  3 +
 .../spiffworkflow_backend/models/typeahead.py | 16 +++++
 .../routes/connector_proxy_controller.py      | 35 ++++++++++
 .../services/custom_parser.py                 |  7 ++
 .../services/process_instance_processor.py    |  2 +
 8 files changed, 173 insertions(+)
 create mode 100644 spiffworkflow-backend/migrations/versions/377be1608b45_.py
 create mode 100644 spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py
 create mode 100644 spiffworkflow-backend/src/spiffworkflow_backend/models/typeahead.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7641f492..6a15caa2 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -51,6 +51,7 @@ repos:
         language: system
         types: [python]
         require_serial: true
+        exclude: "/migrations/"
       - id: trailing-whitespace
         files: ^spiffworkflow-backend/
         name: Trim Trailing Whitespace
diff --git a/spiffworkflow-backend/migrations/versions/377be1608b45_.py b/spiffworkflow-backend/migrations/versions/377be1608b45_.py
new file mode 100644
index 00000000..5c88bb04
--- /dev/null
+++ b/spiffworkflow-backend/migrations/versions/377be1608b45_.py
@@ -0,0 +1,44 @@
+"""empty message
+
+Revision ID: 377be1608b45
+Revises: e4b6bbf83a3e
+Create Date: 2023-06-07 12:39:19.989484
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '377be1608b45'
+down_revision = 'e4b6bbf83a3e'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('typeahead',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('category', sa.String(length=255), nullable=True),
+    sa.Column('search_term', sa.String(length=255), nullable=True),
+    sa.Column('result', sa.JSON(), nullable=True),
+    sa.Column('updated_at_in_seconds', sa.Integer(), nullable=True),
+    sa.Column('created_at_in_seconds', sa.Integer(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    with op.batch_alter_table('typeahead', schema=None) as batch_op:
+        batch_op.create_index(batch_op.f('ix_typeahead_category'), ['category'], unique=False)
+        batch_op.create_index(batch_op.f('ix_typeahead_search_term'), ['search_term'], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('typeahead', schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f('ix_typeahead_search_term'))
+        batch_op.drop_index(batch_op.f('ix_typeahead_category'))
+
+    op.drop_table('typeahead')
+    # ### end Alembic commands ###
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py
new file mode 100644
index 00000000..4ea8c53b
--- /dev/null
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py
@@ -0,0 +1,65 @@
+from time import time
+from typing import Any
+
+from SpiffWorkflow.bpmn.serializer.helpers.spec import BpmnSpecConverter  # type: ignore
+from SpiffWorkflow.bpmn.specs.data_spec import BpmnDataStoreSpecification  # type: ignore
+from SpiffWorkflow.task import Task as SpiffTask  # type: ignore
+from spiffworkflow_backend.models.db import db
+from spiffworkflow_backend.models.typeahead import TypeaheadModel
+
+
+class TypeaheadDataStore(BpmnDataStoreSpecification):  # type: ignore
+    """TypeaheadDataStore."""
+
+    def get(self, my_task: SpiffTask) -> None:
+        """get."""
+        raise Exception("This is a write only data store.")
+
+    def set(self, my_task: SpiffTask) -> None:
+        """set."""
+        typeahead_data_by_category = my_task.data[self.bpmn_id]
+        for category, items in typeahead_data_by_category.items():
+            db.session.query(TypeaheadModel).filter_by(category=category).delete()
+            objects = [self._make_object(category, item) for item in items]
+            db.session.bulk_save_objects(objects)
+        db.session.commit()
+        del my_task.data[self.bpmn_id]
+
+    def _make_object(self, category: str, item: dict[str, Any]) -> TypeaheadModel:
+        now = round(time())
+        return TypeaheadModel(
+            category=category,
+            search_term=item["search_term"],
+            result=item["result"],
+            created_at_in_seconds=now,
+            updated_at_in_seconds=now,
+        )
+
+    @staticmethod
+    def register_converter(spec_config: dict[str, Any]) -> None:
+        spec_config["task_specs"].append(TypeaheadDataStoreConverter)
+
+    @staticmethod
+    def register_data_store_class(data_store_classes: dict[str, Any]) -> None:
+        data_store_classes["TypeaheadDataStore"] = TypeaheadDataStore
+
+
+class TypeaheadDataStoreConverter(BpmnSpecConverter):  # type: ignore
+    """TypeaheadDataStoreConverter."""
+
+    def __init__(self, registry):  # type: ignore
+        """__init__."""
+        super().__init__(TypeaheadDataStore, registry)
+
+    def to_dict(self, spec: Any) -> dict[str, Any]:
+        """to_dict."""
+        return {
+            "bpmn_id": spec.bpmn_id,
+            "bpmn_name": spec.bpmn_name,
+            "capacity": spec.capacity,
+            "is_unlimited": spec.is_unlimited,
+        }
+
+    def from_dict(self, dct: dict[str, Any]) -> TypeaheadDataStore:
+        """from_dict."""
+        return TypeaheadDataStore(**dct)
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/load_database_models.py b/spiffworkflow-backend/src/spiffworkflow_backend/load_database_models.py
index 22e570e1..0553b1c2 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/load_database_models.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/load_database_models.py
@@ -79,5 +79,8 @@ from spiffworkflow_backend.models.active_user import (
 from spiffworkflow_backend.models.process_model_cycle import (
     ProcessModelCycleModel,
 )  # noqa: F401
+from spiffworkflow_backend.models.typeahead import (
+    TypeaheadModel,
+)  # noqa: F401
 
 add_listeners()
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/models/typeahead.py b/spiffworkflow-backend/src/spiffworkflow_backend/models/typeahead.py
new file mode 100644
index 00000000..538a485c
--- /dev/null
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/models/typeahead.py
@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+from spiffworkflow_backend.models.db import SpiffworkflowBaseDBModel
+from spiffworkflow_backend.models.db import db
+
+
+@dataclass
+class TypeaheadModel(SpiffworkflowBaseDBModel):
+    __tablename__ = "typeahead"
+
+    id: int = db.Column(db.Integer, primary_key=True)
+    category: str = db.Column(db.String(255), index=True)
+    search_term: str = db.Column(db.String(255), index=True)
+    result: dict = db.Column(db.JSON)
+    updated_at_in_seconds: int = db.Column(db.Integer)
+    created_at_in_seconds: int = db.Column(db.Integer)
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/routes/connector_proxy_controller.py b/spiffworkflow-backend/src/spiffworkflow_backend/routes/connector_proxy_controller.py
index a37abd24..f9f05ff8 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/routes/connector_proxy_controller.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/routes/connector_proxy_controller.py
@@ -1,3 +1,4 @@
+import json
 from typing import Any
 
 import flask.wrappers
@@ -6,6 +7,8 @@ from flask import current_app
 from flask.wrappers import Response
 
 from spiffworkflow_backend.config import HTTP_REQUEST_TIMEOUT_SECONDS
+from spiffworkflow_backend.models.db import db
+from spiffworkflow_backend.models.typeahead import TypeaheadModel
 
 
 def connector_proxy_typeahead_url() -> Any:
@@ -14,6 +17,34 @@ def connector_proxy_typeahead_url() -> Any:
 
 
 def typeahead(category: str, prefix: str, limit: int) -> flask.wrappers.Response:
+    if _has_local_data(category):
+        return _local_typeahead(category, prefix, limit)
+
+    return _remote_typeahead(category, prefix, limit)
+
+
+def _local_typeahead(category: str, prefix: str, limit: int) -> flask.wrappers.Response:
+    results = (
+        db.session.query(TypeaheadModel.result)
+        .filter(
+            TypeaheadModel.category == category,
+            TypeaheadModel.search_term.ilike(f"{prefix}%"),  # type: ignore
+        )
+        .order_by(TypeaheadModel.search_term)
+        .limit(limit)
+        .all()
+        or []
+    )
+
+    # this is a bummer but sqlalchemy returns a tuple of one field for each result
+    results = [result[0] for result in results]
+
+    response = json.dumps(results)
+
+    return Response(response, status=200, mimetype="application/json")
+
+
+def _remote_typeahead(category: str, prefix: str, limit: int) -> flask.wrappers.Response:
     url = f"{connector_proxy_typeahead_url()}/v1/typeahead/{category}?prefix={prefix}&limit={limit}"
 
     proxy_response = requests.get(url, timeout=HTTP_REQUEST_TIMEOUT_SECONDS)
@@ -21,3 +52,7 @@ def typeahead(category: str, prefix: str, limit: int) -> flask.wrappers.Response
     response = proxy_response.text
 
     return Response(response, status=status, mimetype="application/json")
+
+
+def _has_local_data(category: str) -> bool:
+    return db.session.query(TypeaheadModel.category).filter_by(category=category).first() is not None
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/services/custom_parser.py b/spiffworkflow-backend/src/spiffworkflow_backend/services/custom_parser.py
index a65482f7..c7f649cd 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/services/custom_parser.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/services/custom_parser.py
@@ -1,5 +1,8 @@
+from typing import Any
+
 from SpiffWorkflow.dmn.parser.BpmnDmnParser import BpmnDmnParser  # type: ignore
 from SpiffWorkflow.spiff.parser.process import SpiffBpmnParser  # type: ignore
+from spiffworkflow_backend.data_stores.typeahead import TypeaheadDataStore
 from spiffworkflow_backend.specs.start_event import StartEvent
 
 
@@ -10,3 +13,7 @@ class MyCustomParser(BpmnDmnParser):  # type: ignore
     OVERRIDE_PARSER_CLASSES.update(SpiffBpmnParser.OVERRIDE_PARSER_CLASSES)
 
     StartEvent.register_parser_class(OVERRIDE_PARSER_CLASSES)
+
+    DATA_STORE_CLASSES: dict[str, Any] = {}
+
+    TypeaheadDataStore.register_data_store_class(DATA_STORE_CLASSES)
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py b/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py
index c491b192..1fda2967 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py
@@ -41,6 +41,7 @@ from SpiffWorkflow.spiff.serializer.config import SPIFF_SPEC_CONFIG  # type: ign
 from SpiffWorkflow.task import Task as SpiffTask  # type: ignore
 from SpiffWorkflow.task import TaskState
 from SpiffWorkflow.util.deep_merge import DeepMerge  # type: ignore
+from spiffworkflow_backend.data_stores.typeahead import TypeaheadDataStore
 from spiffworkflow_backend.exceptions.api_error import ApiError
 from spiffworkflow_backend.models.bpmn_process import BpmnProcessModel
 from spiffworkflow_backend.models.bpmn_process_definition import BpmnProcessDefinitionModel
@@ -84,6 +85,7 @@ from spiffworkflow_backend.specs.start_event import StartEvent
 from sqlalchemy import and_
 
 StartEvent.register_converter(SPIFF_SPEC_CONFIG)
+TypeaheadDataStore.register_converter(SPIFF_SPEC_CONFIG)
 
 # Sorry about all this crap.  I wanted to move this thing to another file, but
 # importing a bunch of types causes circular imports.