|
import logging |
|
from logging import Logger |
|
from typing import Optional |
|
from uuid import uuid4 |
|
|
|
from google.cloud import datastore |
|
from google.cloud.datastore import Client, Entity, Query |
|
from slack_sdk.oauth import OAuthStateStore, InstallationStore |
|
from slack_sdk.oauth.installation_store import Installation, Bot |
|
|
|
|
|
class GoogleDatastoreInstallationStore(InstallationStore): |
|
datastore_client: Client |
|
|
|
def __init__( |
|
self, |
|
*, |
|
datastore_client: Client, |
|
logger: Logger, |
|
): |
|
self.datastore_client = datastore_client |
|
self._logger = logger |
|
|
|
@property |
|
def logger(self) -> Logger: |
|
if self._logger is None: |
|
self._logger = logging.getLogger(__name__) |
|
return self._logger |
|
|
|
def installation_key( |
|
self, |
|
*, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
user_id: Optional[str], |
|
suffix: Optional[str] = None, |
|
is_enterprise_install: Optional[bool] = None, |
|
): |
|
enterprise_id = enterprise_id or "none" |
|
team_id = "none" if is_enterprise_install else team_id or "none" |
|
name = ( |
|
f"{enterprise_id}-{team_id}-{user_id}" |
|
if user_id |
|
else f"{enterprise_id}-{team_id}" |
|
) |
|
if suffix is not None: |
|
name += "-" + suffix |
|
return self.datastore_client.key("installations", name) |
|
|
|
def bot_key( |
|
self, |
|
*, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
suffix: Optional[str] = None, |
|
is_enterprise_install: Optional[bool] = None, |
|
): |
|
enterprise_id = enterprise_id or "none" |
|
team_id = "none" if is_enterprise_install else team_id or "none" |
|
name = f"{enterprise_id}-{team_id}" |
|
if suffix is not None: |
|
name += "-" + suffix |
|
return self.datastore_client.key("bots", name) |
|
|
|
def save(self, i: Installation): |
|
# the latest installation in the workspace |
|
installation_entity: Entity = datastore.Entity( |
|
key=self.installation_key( |
|
enterprise_id=i.enterprise_id, |
|
team_id=i.team_id, |
|
user_id=None, # user_id is removed |
|
is_enterprise_install=i.is_enterprise_install, |
|
) |
|
) |
|
installation_entity.update(**i.to_dict()) |
|
self.datastore_client.put(installation_entity) |
|
|
|
# the latest installation associated with a user |
|
user_entity: Entity = datastore.Entity( |
|
key=self.installation_key( |
|
enterprise_id=i.enterprise_id, |
|
team_id=i.team_id, |
|
user_id=i.user_id, |
|
is_enterprise_install=i.is_enterprise_install, |
|
) |
|
) |
|
user_entity.update(**i.to_dict()) |
|
self.datastore_client.put(user_entity) |
|
# history data |
|
user_entity.key = self.installation_key( |
|
enterprise_id=i.enterprise_id, |
|
team_id=i.team_id, |
|
user_id=i.user_id, |
|
is_enterprise_install=i.is_enterprise_install, |
|
suffix=str(i.installed_at), |
|
) |
|
self.datastore_client.put(user_entity) |
|
|
|
# the latest bot authorization in the workspace |
|
bot = i.to_bot() |
|
bot_entity: Entity = datastore.Entity( |
|
key=self.bot_key( |
|
enterprise_id=i.enterprise_id, |
|
team_id=i.team_id, |
|
is_enterprise_install=i.is_enterprise_install, |
|
) |
|
) |
|
bot_entity.update(**bot.to_dict()) |
|
self.datastore_client.put(bot_entity) |
|
# history data |
|
bot_entity.key = self.bot_key( |
|
enterprise_id=i.enterprise_id, |
|
team_id=i.team_id, |
|
is_enterprise_install=i.is_enterprise_install, |
|
suffix=str(i.installed_at), |
|
) |
|
self.datastore_client.put(bot_entity) |
|
|
|
def find_bot( |
|
self, |
|
*, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
is_enterprise_install: Optional[bool] = False, |
|
) -> Optional[Bot]: |
|
entity: Entity = self.datastore_client.get( |
|
self.bot_key( |
|
enterprise_id=enterprise_id, |
|
team_id=team_id, |
|
is_enterprise_install=is_enterprise_install, |
|
) |
|
) |
|
if entity is not None: |
|
entity["installed_at"] = entity["installed_at"].timestamp() |
|
return Bot(**entity) |
|
return None |
|
|
|
def find_installation( |
|
self, |
|
*, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
user_id: Optional[str] = None, |
|
is_enterprise_install: Optional[bool] = False, |
|
) -> Optional[Installation]: |
|
entity: Entity = self.datastore_client.get( |
|
self.installation_key( |
|
enterprise_id=enterprise_id, |
|
team_id=team_id, |
|
user_id=user_id, |
|
is_enterprise_install=is_enterprise_install, |
|
) |
|
) |
|
if entity is not None: |
|
entity["installed_at"] = entity["installed_at"].timestamp() |
|
return Installation(**entity) |
|
return None |
|
|
|
def delete_installation( |
|
self, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
user_id: Optional[str], |
|
) -> None: |
|
installation_key = self.installation_key( |
|
enterprise_id=enterprise_id, |
|
team_id=team_id, |
|
user_id=user_id, |
|
) |
|
q: Query = self.datastore_client.query() |
|
q.key_filter(installation_key, ">=") |
|
for entity in q.fetch(): |
|
if entity.key.name.startswith(installation_key.name): |
|
self.datastore_client.delete(entity.key) |
|
else: |
|
break |
|
|
|
def delete_bot( |
|
self, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
) -> None: |
|
bot_key = self.bot_key( |
|
enterprise_id=enterprise_id, |
|
team_id=team_id, |
|
) |
|
q: Query = self.datastore_client.query() |
|
q.key_filter(bot_key, ">=") |
|
for entity in q.fetch(): |
|
if entity.key.name.startswith(bot_key.name): |
|
self.datastore_client.delete(entity.key) |
|
else: |
|
break |
|
|
|
def delete_all( |
|
self, |
|
enterprise_id: Optional[str], |
|
team_id: Optional[str], |
|
): |
|
self.delete_bot(enterprise_id=enterprise_id, team_id=team_id) |
|
self.delete_installation( |
|
enterprise_id=enterprise_id, team_id=team_id, user_id=None |
|
) |
|
|
|
|
|
class GoogleDatastoreOAuthStateStore(OAuthStateStore): |
|
logger: Logger |
|
datastore_client: Client |
|
collection_id: str |
|
|
|
def __init__( |
|
self, |
|
*, |
|
datastore_client: Client, |
|
logger: Logger, |
|
): |
|
self.datastore_client = datastore_client |
|
self._logger = logger |
|
self.collection_id = "oauth_state_values" |
|
|
|
@property |
|
def logger(self) -> Logger: |
|
if self._logger is None: |
|
self._logger = logging.getLogger(__name__) |
|
return self._logger |
|
|
|
def consume(self, state: str) -> bool: |
|
key = self.datastore_client.key(self.collection_id, state) |
|
entity = self.datastore_client.get(key) |
|
if entity is not None: |
|
self.datastore_client.delete(key) |
|
return True |
|
return False |
|
|
|
def issue(self, *args, **kwargs) -> str: |
|
state_value = str(uuid4()) |
|
entity: Entity = datastore.Entity( |
|
key=self.datastore_client.key(self.collection_id, state_value) |
|
) |
|
entity.update(value=state_value) |
|
self.datastore_client.put(entity) |
|
return state_value |
Is there any specific reason why this implementation does not use Datastore composite index?
Now it simulates the composite index by concatenating multiple values in the Key.
All other implementations (SQLite3InstallationStore, SQLAlchemyInstallationStore, DjangoInstallationStore) work with a composite index.
Also I tried to run the tests from TestSQLite3 (python-slack-sdk-main/tests/slack_sdk/oauth/installation_store/test_sqlite3.py)
against GoogleDatastoreInstallationStore and they do not pass.
I'm going to try to rework it to use Datastore composite index instead of concatenated values in the key.