Build reusable bot framework
This commit is contained in:
1
core/__init__.py
Normal file
1
core/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Framework services and persistence helpers."""
|
||||
351
core/api_keys.py
Normal file
351
core/api_keys.py
Normal file
@@ -0,0 +1,351 @@
|
||||
import datetime
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import uuid
|
||||
|
||||
import core.postgres as postgres
|
||||
import core.users as users
|
||||
|
||||
|
||||
USER_KEY_TYPE = "user"
|
||||
SERVICE_KEY_TYPE = "service"
|
||||
DEFAULT_SERVICE_SCOPES = (
|
||||
"discord:session",
|
||||
"outbox:claim",
|
||||
"outbox:deliver",
|
||||
)
|
||||
_TOKEN_PREFIX_LENGTH = 20
|
||||
_MIN_TOKEN_LENGTH = 32
|
||||
|
||||
|
||||
def _utcNow():
|
||||
return datetime.datetime.now(datetime.timezone.utc)
|
||||
|
||||
|
||||
def _hashToken(token):
|
||||
return hashlib.sha256(token.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _getPrefix(token):
|
||||
return token[:_TOKEN_PREFIX_LENGTH]
|
||||
|
||||
|
||||
def _normalizeScopes(scopes):
|
||||
if scopes is None:
|
||||
return []
|
||||
if isinstance(scopes, str):
|
||||
scopes = re.split(r"[\s,]+", scopes)
|
||||
if not isinstance(scopes, (list, tuple, set)):
|
||||
raise ValueError("scopes must be a list or comma-separated string")
|
||||
|
||||
normalized = []
|
||||
for scope in scopes:
|
||||
if not isinstance(scope, str) or not scope.strip():
|
||||
continue
|
||||
scope = scope.strip()
|
||||
if len(scope) > 100:
|
||||
raise ValueError("service scopes must be at most 100 characters")
|
||||
if scope not in normalized:
|
||||
normalized.append(scope)
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalizeExpiry(expiresAt):
|
||||
if expiresAt is None:
|
||||
return None
|
||||
if isinstance(expiresAt, str):
|
||||
try:
|
||||
expiresAt = datetime.datetime.fromisoformat(expiresAt.replace("Z", "+00:00"))
|
||||
except ValueError as error:
|
||||
raise ValueError("expires_at must be an ISO-8601 datetime") from error
|
||||
if not isinstance(expiresAt, datetime.datetime):
|
||||
raise ValueError("expires_at must be a datetime")
|
||||
if expiresAt.tzinfo is None:
|
||||
expiresAt = expiresAt.replace(tzinfo=datetime.timezone.utc)
|
||||
expiresAt = expiresAt.astimezone(datetime.timezone.utc)
|
||||
if expiresAt <= _utcNow():
|
||||
raise ValueError("expires_at must be in the future")
|
||||
return expiresAt
|
||||
|
||||
|
||||
def _validateSecret(token):
|
||||
if not isinstance(token, str) or len(token) < _MIN_TOKEN_LENGTH:
|
||||
raise ValueError(f"API keys must be at least {_MIN_TOKEN_LENGTH} characters")
|
||||
return token
|
||||
|
||||
|
||||
def _keyID(value):
|
||||
try:
|
||||
return str(uuid.UUID(str(value)))
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
return None
|
||||
|
||||
|
||||
def _publicKey(record, includeSecret=None):
|
||||
if not record:
|
||||
return None
|
||||
public = dict(record)
|
||||
public.pop("key_hash", None)
|
||||
scopes = public.get("scopes")
|
||||
if isinstance(scopes, str):
|
||||
public["scopes"] = json.loads(scopes)
|
||||
if includeSecret is not None:
|
||||
public["key"] = includeSecret
|
||||
return public
|
||||
|
||||
|
||||
def createApiKey(
|
||||
name,
|
||||
keyType,
|
||||
userUUID=None,
|
||||
serviceName=None,
|
||||
scopes=None,
|
||||
expiresAt=None,
|
||||
secret=None,
|
||||
):
|
||||
if not isinstance(name, str) or not name.strip():
|
||||
raise ValueError("API key name is required")
|
||||
name = name.strip()
|
||||
if len(name) > 255:
|
||||
raise ValueError("API key name must be at most 255 characters")
|
||||
if keyType not in {USER_KEY_TYPE, SERVICE_KEY_TYPE}:
|
||||
raise ValueError("key_type must be user or service")
|
||||
|
||||
normalizedScopes = _normalizeScopes(scopes)
|
||||
if keyType == USER_KEY_TYPE:
|
||||
if not userUUID or not users.doesUserUUIDExist(userUUID):
|
||||
raise ValueError("user does not exist")
|
||||
if serviceName is not None:
|
||||
raise ValueError("user API keys cannot have a service name")
|
||||
if normalizedScopes:
|
||||
raise ValueError("user API keys cannot have service scopes")
|
||||
else:
|
||||
if userUUID is not None:
|
||||
raise ValueError("service API keys cannot have a user")
|
||||
if not isinstance(serviceName, str) or not serviceName.strip():
|
||||
raise ValueError("service name is required")
|
||||
serviceName = serviceName.strip()
|
||||
if len(serviceName) > 255:
|
||||
raise ValueError("service name must be at most 255 characters")
|
||||
|
||||
expiresAt = _normalizeExpiry(expiresAt)
|
||||
if secret is None:
|
||||
secret = f"llmbot_{keyType}_{secrets.token_urlsafe(32)}"
|
||||
secret = _validateSecret(secret)
|
||||
|
||||
keyData = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"name": name,
|
||||
"key_type": keyType,
|
||||
"user_uuid": str(userUUID) if userUUID else None,
|
||||
"service_name": serviceName,
|
||||
"key_prefix": _getPrefix(secret),
|
||||
"key_hash": _hashToken(secret),
|
||||
"scopes": json.dumps(normalizedScopes),
|
||||
"expires_at": expiresAt,
|
||||
}
|
||||
with postgres.get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO api_keys (
|
||||
id, name, key_type, user_uuid, service_name,
|
||||
key_prefix, key_hash, scopes, expires_at
|
||||
) VALUES (
|
||||
%(id)s, %(name)s, %(key_type)s, %(user_uuid)s, %(service_name)s,
|
||||
%(key_prefix)s, %(key_hash)s, %(scopes)s::jsonb, %(expires_at)s
|
||||
)
|
||||
RETURNING *
|
||||
""",
|
||||
keyData,
|
||||
)
|
||||
record = dict(cursor.fetchone())
|
||||
return _publicKey(record, includeSecret=secret)
|
||||
|
||||
|
||||
def createUserApiKey(userUUID, name, expiresAt=None):
|
||||
return createApiKey(
|
||||
name,
|
||||
USER_KEY_TYPE,
|
||||
userUUID=userUUID,
|
||||
expiresAt=expiresAt,
|
||||
)
|
||||
|
||||
|
||||
def createServiceApiKey(serviceName, name, scopes, expiresAt=None):
|
||||
return createApiKey(
|
||||
name,
|
||||
SERVICE_KEY_TYPE,
|
||||
serviceName=serviceName,
|
||||
scopes=scopes,
|
||||
expiresAt=expiresAt,
|
||||
)
|
||||
|
||||
|
||||
def getApiKey(keyID):
|
||||
keyID = _keyID(keyID)
|
||||
if not keyID:
|
||||
return None
|
||||
return _publicKey(postgres.select_one("api_keys", {"id": keyID}))
|
||||
|
||||
|
||||
def listApiKeys(userUUID=None, serviceName=None, includeRevoked=False):
|
||||
if userUUID is not None and serviceName is not None:
|
||||
raise ValueError("filter by either user or service, not both")
|
||||
|
||||
clauses = []
|
||||
params = {}
|
||||
if userUUID is not None:
|
||||
clauses.append("user_uuid = %(user_uuid)s")
|
||||
params["user_uuid"] = str(userUUID)
|
||||
if serviceName is not None:
|
||||
clauses.append("service_name = %(service_name)s")
|
||||
params["service_name"] = serviceName
|
||||
if not includeRevoked:
|
||||
clauses.append("revoked_at IS NULL")
|
||||
|
||||
query = "SELECT * FROM api_keys"
|
||||
if clauses:
|
||||
query += " WHERE " + " AND ".join(clauses)
|
||||
query += " ORDER BY created_at DESC"
|
||||
return [_publicKey(record) for record in postgres.execute(query, params)]
|
||||
|
||||
|
||||
def listUserApiKeys(userUUID, includeRevoked=False):
|
||||
return listApiKeys(userUUID=userUUID, includeRevoked=includeRevoked)
|
||||
|
||||
|
||||
def revokeApiKey(keyID, userUUID=None, serviceName=None):
|
||||
keyID = _keyID(keyID)
|
||||
if not keyID:
|
||||
return False
|
||||
clauses = ["id = %(id)s", "revoked_at IS NULL"]
|
||||
params = {"id": keyID}
|
||||
if userUUID is not None:
|
||||
clauses.append("user_uuid = %(user_uuid)s")
|
||||
params["user_uuid"] = str(userUUID)
|
||||
if serviceName is not None:
|
||||
clauses.append("service_name = %(service_name)s")
|
||||
params["service_name"] = serviceName
|
||||
|
||||
records = postgres.execute(
|
||||
"UPDATE api_keys SET revoked_at = CURRENT_TIMESTAMP "
|
||||
f"WHERE {' AND '.join(clauses)} RETURNING id",
|
||||
params,
|
||||
)
|
||||
return bool(records)
|
||||
|
||||
|
||||
def revokeUserApiKey(userUUID, keyID):
|
||||
return revokeApiKey(keyID, userUUID=userUUID)
|
||||
|
||||
|
||||
def authenticateApiKey(token, requiredScopes=None):
|
||||
if not isinstance(token, str) or len(token) < _MIN_TOKEN_LENGTH:
|
||||
return None
|
||||
|
||||
records = postgres.select(
|
||||
"api_keys",
|
||||
{
|
||||
"key_prefix": _getPrefix(token),
|
||||
"key_hash": _hashToken(token),
|
||||
"revoked_at": None,
|
||||
},
|
||||
)
|
||||
if not records:
|
||||
return None
|
||||
|
||||
record = records[0]
|
||||
if not hmac.compare_digest(record["key_hash"], _hashToken(token)):
|
||||
return None
|
||||
expiresAt = record.get("expires_at")
|
||||
if expiresAt is not None:
|
||||
if expiresAt.tzinfo is None:
|
||||
expiresAt = expiresAt.replace(tzinfo=datetime.timezone.utc)
|
||||
if expiresAt <= _utcNow():
|
||||
return None
|
||||
|
||||
scopes = record.get("scopes") or []
|
||||
if isinstance(scopes, str):
|
||||
scopes = json.loads(scopes)
|
||||
requiredScopes = _normalizeScopes(requiredScopes)
|
||||
if requiredScopes and "*" not in scopes:
|
||||
if not set(requiredScopes).issubset(set(scopes)):
|
||||
return None
|
||||
|
||||
postgres.update(
|
||||
"api_keys",
|
||||
{"last_used_at": _utcNow()},
|
||||
{"id": record["id"]},
|
||||
)
|
||||
return {
|
||||
"type": record["key_type"],
|
||||
"authentication": "api_key",
|
||||
"user_uuid": record.get("user_uuid"),
|
||||
"service_name": record.get("service_name"),
|
||||
"api_key_id": record["id"],
|
||||
"scopes": scopes,
|
||||
"can_manage_api_keys": False,
|
||||
}
|
||||
|
||||
|
||||
def bootstrapServiceApiKey():
|
||||
secret = os.getenv("BOT_API_KEY")
|
||||
if not secret:
|
||||
return None
|
||||
_validateSecret(secret)
|
||||
|
||||
serviceName = "discord-bot"
|
||||
keyName = "BOT_API_KEY"
|
||||
scopes = _normalizeScopes(os.getenv("BOT_API_KEY_SCOPES"))
|
||||
if not scopes:
|
||||
scopes = list(DEFAULT_SERVICE_SCOPES)
|
||||
tokenHash = _hashToken(secret)
|
||||
|
||||
with postgres.get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE api_keys
|
||||
SET revoked_at = CURRENT_TIMESTAMP
|
||||
WHERE key_type = 'service'
|
||||
AND service_name = %(service_name)s
|
||||
AND name = %(name)s
|
||||
AND key_hash != %(key_hash)s
|
||||
AND revoked_at IS NULL
|
||||
""",
|
||||
{
|
||||
"service_name": serviceName,
|
||||
"name": keyName,
|
||||
"key_hash": tokenHash,
|
||||
},
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO api_keys (
|
||||
id, name, key_type, service_name,
|
||||
key_prefix, key_hash, scopes
|
||||
) VALUES (
|
||||
%(id)s, %(name)s, 'service', %(service_name)s,
|
||||
%(key_prefix)s, %(key_hash)s, %(scopes)s::jsonb
|
||||
)
|
||||
ON CONFLICT (key_hash) DO UPDATE SET
|
||||
name = EXCLUDED.name,
|
||||
service_name = EXCLUDED.service_name,
|
||||
scopes = EXCLUDED.scopes,
|
||||
expires_at = NULL,
|
||||
revoked_at = NULL
|
||||
RETURNING *
|
||||
""",
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"name": keyName,
|
||||
"service_name": serviceName,
|
||||
"key_prefix": _getPrefix(secret),
|
||||
"key_hash": tokenHash,
|
||||
"scopes": json.dumps(scopes),
|
||||
},
|
||||
)
|
||||
return _publicKey(dict(cursor.fetchone()))
|
||||
180
core/auth.py
Normal file
180
core/auth.py
Normal file
@@ -0,0 +1,180 @@
|
||||
import datetime
|
||||
import os
|
||||
|
||||
import bcrypt
|
||||
import jwt
|
||||
from jwt.exceptions import ExpiredSignatureError, InvalidTokenError
|
||||
|
||||
import core.api_keys as apiKeys
|
||||
import core.postgres as postgres
|
||||
import core.users as users
|
||||
|
||||
|
||||
DEFAULT_TOKEN_LIFETIME = datetime.timedelta(hours=1)
|
||||
|
||||
|
||||
def decodeLoginToken(loginToken):
|
||||
secret = os.getenv("JWT_SECRET")
|
||||
if not secret or not isinstance(loginToken, str):
|
||||
return None
|
||||
try:
|
||||
payload = jwt.decode(loginToken, secret, algorithms=["HS256"])
|
||||
except (ExpiredSignatureError, InvalidTokenError):
|
||||
return None
|
||||
|
||||
userUUID = payload.get("sub")
|
||||
if not isinstance(userUUID, str) or not users.doesUserUUIDExist(userUUID):
|
||||
return None
|
||||
return payload
|
||||
|
||||
|
||||
def createLoginToken(userUUID, name=None, expiresIn=None, extraClaims=None):
|
||||
secret = os.getenv("JWT_SECRET")
|
||||
if not secret:
|
||||
raise RuntimeError("JWT_SECRET is required")
|
||||
if not userUUID or not users.doesUserUUIDExist(userUUID):
|
||||
raise ValueError("user does not exist")
|
||||
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
if expiresIn is None:
|
||||
expiresIn = DEFAULT_TOKEN_LIFETIME
|
||||
elif isinstance(expiresIn, (int, float)):
|
||||
expiresIn = datetime.timedelta(seconds=expiresIn)
|
||||
if not isinstance(expiresIn, datetime.timedelta) or expiresIn.total_seconds() <= 0:
|
||||
raise ValueError("token lifetime must be positive")
|
||||
|
||||
payload = {
|
||||
"sub": str(userUUID),
|
||||
"name": name if name is not None else users.getUserFirstName(userUUID),
|
||||
"iat": now,
|
||||
"exp": now + expiresIn,
|
||||
}
|
||||
if extraClaims is not None:
|
||||
if not isinstance(extraClaims, dict):
|
||||
raise ValueError("extra claims must be an object")
|
||||
protectedClaims = {"sub", "iat", "exp"}
|
||||
payload.update(
|
||||
{
|
||||
key: value
|
||||
for key, value in extraClaims.items()
|
||||
if key not in protectedClaims
|
||||
}
|
||||
)
|
||||
return jwt.encode(payload, secret, algorithm="HS256")
|
||||
|
||||
|
||||
def verifyLoginToken(login_token, username=False, userUUID=False):
|
||||
if username:
|
||||
userUUID = users.getUserUUID(username)
|
||||
|
||||
if userUUID:
|
||||
decodedToken = decodeLoginToken(login_token)
|
||||
if decodedToken and decodedToken.get("sub") == str(userUUID):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def getUserpasswordHash(userUUID):
|
||||
user = postgres.select_one("users", {"id": userUUID})
|
||||
if user:
|
||||
pw_hash = user.get("password_hashed")
|
||||
if isinstance(pw_hash, memoryview):
|
||||
return bytes(pw_hash)
|
||||
return pw_hash
|
||||
return None
|
||||
|
||||
|
||||
def getLoginToken(username, password):
|
||||
if not isinstance(username, str) or not isinstance(password, str):
|
||||
return False
|
||||
userUUID = users.getUserUUID(username)
|
||||
if userUUID:
|
||||
formattedPass = password.encode("utf-8")
|
||||
if not formattedPass or len(formattedPass) > users.MAX_PASSWORD_BYTES:
|
||||
return False
|
||||
usersHashedPassword = getUserpasswordHash(userUUID)
|
||||
if not usersHashedPassword:
|
||||
return False
|
||||
try:
|
||||
if bcrypt.checkpw(formattedPass, usersHashedPassword):
|
||||
return createLoginToken(userUUID)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def unregisterUser(userUUID, password):
|
||||
pw_hash = getUserpasswordHash(userUUID)
|
||||
if not pw_hash or not isinstance(password, str):
|
||||
return False
|
||||
formattedPassword = password.encode("utf-8")
|
||||
if not formattedPassword or len(formattedPassword) > users.MAX_PASSWORD_BYTES:
|
||||
return False
|
||||
try:
|
||||
if bcrypt.checkpw(formattedPassword, pw_hash):
|
||||
return users.deleteUser(userUUID)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def getBearerToken(authorization):
|
||||
if not isinstance(authorization, str):
|
||||
return None
|
||||
scheme, separator, token = authorization.strip().partition(" ")
|
||||
if not separator or scheme.lower() != "bearer" or not token.strip():
|
||||
return None
|
||||
return token.strip()
|
||||
|
||||
|
||||
def authenticateBearerToken(
|
||||
authorization,
|
||||
requiredScopes=None,
|
||||
allowUser=True,
|
||||
allowService=True,
|
||||
):
|
||||
token = getBearerToken(authorization)
|
||||
if not token:
|
||||
return None
|
||||
|
||||
if token.count(".") == 2:
|
||||
payload = decodeLoginToken(token)
|
||||
if payload:
|
||||
if not allowUser or requiredScopes:
|
||||
return None
|
||||
return {
|
||||
"type": "user",
|
||||
"authentication": "jwt",
|
||||
"user_uuid": payload["sub"],
|
||||
"service_name": None,
|
||||
"api_key_id": None,
|
||||
"scopes": [],
|
||||
"can_manage_api_keys": True,
|
||||
"claims": payload,
|
||||
}
|
||||
|
||||
principal = apiKeys.authenticateApiKey(token, requiredScopes=requiredScopes)
|
||||
if not principal:
|
||||
return None
|
||||
if principal["type"] == apiKeys.USER_KEY_TYPE and not allowUser:
|
||||
return None
|
||||
if principal["type"] == apiKeys.SERVICE_KEY_TYPE and not allowService:
|
||||
return None
|
||||
return principal
|
||||
|
||||
|
||||
def isUserPrincipal(principal, userUUID=None, requireLogin=False):
|
||||
if not principal or principal.get("type") != "user":
|
||||
return False
|
||||
if userUUID is not None and str(principal.get("user_uuid")) != str(userUUID):
|
||||
return False
|
||||
if requireLogin and principal.get("authentication") != "jwt":
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def hasServiceScope(principal, scope):
|
||||
if not principal or principal.get("type") != "service":
|
||||
return False
|
||||
scopes = principal.get("scopes") or []
|
||||
return "*" in scopes or scope in scopes
|
||||
224
core/identity.py
Normal file
224
core/identity.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
|
||||
import core.postgres as postgres
|
||||
import core.users as users
|
||||
|
||||
|
||||
DISCORD_PROVIDER = "discord"
|
||||
DEFAULT_ENROLLMENT_MODE = "allowlist"
|
||||
_PROVIDER_PATTERN = re.compile(r"^[a-z0-9][a-z0-9_.-]{0,49}$")
|
||||
|
||||
|
||||
def _normalizeProvider(provider):
|
||||
provider = provider.strip().lower() if isinstance(provider, str) else ""
|
||||
if not _PROVIDER_PATTERN.fullmatch(provider):
|
||||
raise ValueError("provider must be a valid provider name")
|
||||
return provider
|
||||
|
||||
|
||||
def _normalizeProviderUserID(providerUserID):
|
||||
if providerUserID is None:
|
||||
raise ValueError("provider user ID is required")
|
||||
providerUserID = str(providerUserID).strip()
|
||||
if not providerUserID or len(providerUserID) > 255:
|
||||
raise ValueError("provider user ID must be between 1 and 255 characters")
|
||||
return providerUserID
|
||||
|
||||
|
||||
def _normalizeDisplayName(displayName):
|
||||
if displayName is None:
|
||||
return None
|
||||
if not isinstance(displayName, str):
|
||||
raise ValueError("display name must be a string")
|
||||
displayName = displayName.strip()
|
||||
if not displayName:
|
||||
return None
|
||||
if len(displayName) > 255:
|
||||
raise ValueError("display name must be at most 255 characters")
|
||||
return displayName
|
||||
|
||||
|
||||
def getDiscordEnrollmentMode():
|
||||
mode = os.getenv("DISCORD_ENROLLMENT_MODE", DEFAULT_ENROLLMENT_MODE)
|
||||
mode = mode.strip().lower()
|
||||
return "open" if mode == "open" else DEFAULT_ENROLLMENT_MODE
|
||||
|
||||
|
||||
def getDiscordAllowlist():
|
||||
configured = os.getenv("DISCORD_ALLOWLIST")
|
||||
if configured is None:
|
||||
configured = os.getenv("DISCORD_ALLOWED_USER_IDS", "")
|
||||
return {
|
||||
value
|
||||
for value in re.split(r"[\s,]+", configured)
|
||||
if value
|
||||
}
|
||||
|
||||
|
||||
def isDiscordEnrollmentAllowed(discordID):
|
||||
discordID = _normalizeProviderUserID(discordID)
|
||||
if getDiscordEnrollmentMode() == "open":
|
||||
return True
|
||||
return discordID in getDiscordAllowlist()
|
||||
|
||||
|
||||
def getProviderIdentity(provider, providerUserID):
|
||||
provider = _normalizeProvider(provider)
|
||||
providerUserID = _normalizeProviderUserID(providerUserID)
|
||||
return postgres.select_one(
|
||||
"provider_identities",
|
||||
{"provider": provider, "provider_user_id": providerUserID},
|
||||
)
|
||||
|
||||
|
||||
def getUserForIdentity(provider, providerUserID):
|
||||
identity = getProviderIdentity(provider, providerUserID)
|
||||
if not identity:
|
||||
return None
|
||||
return users.getUser(identity["user_uuid"])
|
||||
|
||||
|
||||
def listUserIdentities(userUUID):
|
||||
return postgres.select("provider_identities", {"user_uuid": userUUID})
|
||||
|
||||
|
||||
def linkProviderIdentity(userUUID, provider, providerUserID, displayName=None):
|
||||
if not users.doesUserUUIDExist(userUUID):
|
||||
raise ValueError("user does not exist")
|
||||
provider = _normalizeProvider(provider)
|
||||
providerUserID = _normalizeProviderUserID(providerUserID)
|
||||
displayName = _normalizeDisplayName(displayName)
|
||||
|
||||
existing = getProviderIdentity(provider, providerUserID)
|
||||
if existing:
|
||||
if str(existing["user_uuid"]) != str(userUUID):
|
||||
raise ValueError("provider identity is already linked to another user")
|
||||
if existing.get("display_name") != displayName:
|
||||
updated = postgres.update(
|
||||
"provider_identities",
|
||||
{"display_name": displayName},
|
||||
{"id": existing["id"]},
|
||||
)
|
||||
return updated[0]
|
||||
return existing
|
||||
|
||||
return postgres.insert(
|
||||
"provider_identities",
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_uuid": str(userUUID),
|
||||
"provider": provider,
|
||||
"provider_user_id": providerUserID,
|
||||
"display_name": displayName,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def unlinkProviderIdentity(userUUID, provider, providerUserID):
|
||||
provider = _normalizeProvider(provider)
|
||||
providerUserID = _normalizeProviderUserID(providerUserID)
|
||||
deleted = postgres.delete(
|
||||
"provider_identities",
|
||||
{
|
||||
"user_uuid": userUUID,
|
||||
"provider": provider,
|
||||
"provider_user_id": providerUserID,
|
||||
},
|
||||
)
|
||||
return bool(deleted)
|
||||
|
||||
|
||||
def linkDiscordIdentity(userUUID, discordID, displayName=None):
|
||||
return linkProviderIdentity(
|
||||
userUUID,
|
||||
DISCORD_PROVIDER,
|
||||
discordID,
|
||||
displayName=displayName,
|
||||
)
|
||||
|
||||
|
||||
def linkDiscordUser(discordID, userUUID=None, username=None, displayName=None):
|
||||
if bool(userUUID) == bool(username):
|
||||
raise ValueError("provide either user UUID or username")
|
||||
if username:
|
||||
userUUID = users.getUserUUID(username)
|
||||
if not userUUID or not users.doesUserUUIDExist(userUUID):
|
||||
raise ValueError("user does not exist")
|
||||
return linkDiscordIdentity(userUUID, discordID, displayName=displayName)
|
||||
|
||||
|
||||
def getDiscordUser(discordID):
|
||||
return getUserForIdentity(DISCORD_PROVIDER, discordID)
|
||||
|
||||
|
||||
def getOrCreateDiscordUser(discordID, displayName=None, timezoneName=None):
|
||||
discordID = _normalizeProviderUserID(discordID)
|
||||
displayName = _normalizeDisplayName(displayName)
|
||||
if timezoneName is None:
|
||||
timezoneName = users.getDefaultTimezone()
|
||||
else:
|
||||
timezoneName = users.normalizeTimezone(timezoneName)
|
||||
|
||||
existing = getProviderIdentity(DISCORD_PROVIDER, discordID)
|
||||
if existing:
|
||||
if displayName is not None and existing.get("display_name") != displayName:
|
||||
postgres.update(
|
||||
"provider_identities",
|
||||
{"display_name": displayName},
|
||||
{"id": existing["id"]},
|
||||
)
|
||||
return users.getUser(existing["user_uuid"])
|
||||
if not isDiscordEnrollmentAllowed(discordID):
|
||||
return None
|
||||
|
||||
with postgres.get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT pg_advisory_xact_lock(hashtext(%(lock_key)s))",
|
||||
{"lock_key": f"discord:{discordID}"},
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT users.*
|
||||
FROM provider_identities
|
||||
JOIN users ON users.id = provider_identities.user_uuid
|
||||
WHERE provider = %(provider)s AND provider_user_id = %(provider_user_id)s
|
||||
""",
|
||||
{
|
||||
"provider": DISCORD_PROVIDER,
|
||||
"provider_user_id": discordID,
|
||||
},
|
||||
)
|
||||
existingUser = cursor.fetchone()
|
||||
if existingUser:
|
||||
return dict(existingUser)
|
||||
|
||||
userUUID = str(uuid.uuid4())
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO users (id, username, password_hashed, timezone)
|
||||
VALUES (%(id)s, NULL, NULL, %(timezone)s)
|
||||
RETURNING *
|
||||
""",
|
||||
{"id": userUUID, "timezone": timezoneName},
|
||||
)
|
||||
userRecord = dict(cursor.fetchone())
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO provider_identities (
|
||||
id, user_uuid, provider, provider_user_id, display_name
|
||||
) VALUES (
|
||||
%(id)s, %(user_uuid)s, %(provider)s,
|
||||
%(provider_user_id)s, %(display_name)s
|
||||
)
|
||||
""",
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_uuid": userUUID,
|
||||
"provider": DISCORD_PROVIDER,
|
||||
"provider_user_id": discordID,
|
||||
"display_name": displayName,
|
||||
},
|
||||
)
|
||||
return userRecord
|
||||
343
core/jobs.py
Normal file
343
core/jobs.py
Normal file
@@ -0,0 +1,343 @@
|
||||
"""PostgreSQL-backed scheduled job operations."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
import uuid
|
||||
|
||||
from psycopg2.extras import Json
|
||||
|
||||
from core import postgres
|
||||
|
||||
|
||||
DEFAULT_MAX_ATTEMPTS = 3
|
||||
DEFAULT_LEASE_SECONDS = 300
|
||||
DEFAULT_RETRY_SECONDS = 30
|
||||
MAX_RETRY_SECONDS = 900
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _using_cursor(cursor=None):
|
||||
if cursor is not None:
|
||||
yield cursor
|
||||
return
|
||||
with postgres.get_cursor() as owned_cursor:
|
||||
yield owned_cursor
|
||||
|
||||
|
||||
def _timestamp(value, field="timestamp"):
|
||||
if isinstance(value, str):
|
||||
value = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
if not isinstance(value, datetime):
|
||||
raise ValueError(f"{field} must be a datetime or ISO-8601 string")
|
||||
if value.tzinfo is None:
|
||||
value = value.replace(tzinfo=timezone.utc)
|
||||
return value.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _positive(value, field, maximum=None):
|
||||
try:
|
||||
value = int(value)
|
||||
except (TypeError, ValueError) as error:
|
||||
raise ValueError(f"{field} must be a whole number") from error
|
||||
if value < 1 or (maximum is not None and value > maximum):
|
||||
suffix = f" and at most {maximum}" if maximum is not None else ""
|
||||
raise ValueError(f"{field} must be at least 1{suffix}")
|
||||
return value
|
||||
|
||||
|
||||
def _row(cursor):
|
||||
record = cursor.fetchone()
|
||||
return dict(record) if record else None
|
||||
|
||||
|
||||
def create_job(
|
||||
job_type,
|
||||
payload,
|
||||
run_at,
|
||||
user_uuid=None,
|
||||
max_attempts=DEFAULT_MAX_ATTEMPTS,
|
||||
idempotency_key=None,
|
||||
job_id=None,
|
||||
cursor=None,
|
||||
):
|
||||
"""Create a job, returning the existing row for a repeated idempotency key."""
|
||||
if not isinstance(job_type, str) or not job_type.strip():
|
||||
raise ValueError("job_type is required")
|
||||
max_attempts = _positive(max_attempts, "max_attempts", 100)
|
||||
values = {
|
||||
"id": str(job_id or uuid.uuid4()),
|
||||
"user_uuid": user_uuid,
|
||||
"job_type": job_type.strip(),
|
||||
"payload": Json(payload if payload is not None else {}),
|
||||
"run_at": _timestamp(run_at, "run_at"),
|
||||
"max_attempts": max_attempts,
|
||||
"idempotency_key": idempotency_key,
|
||||
}
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
INSERT INTO scheduled_jobs (
|
||||
id, user_uuid, job_type, payload, run_at,
|
||||
max_attempts, idempotency_key
|
||||
) VALUES (
|
||||
%(id)s, %(user_uuid)s, %(job_type)s, %(payload)s, %(run_at)s,
|
||||
%(max_attempts)s, %(idempotency_key)s
|
||||
)
|
||||
ON CONFLICT (job_type, idempotency_key)
|
||||
DO UPDATE SET idempotency_key = EXCLUDED.idempotency_key
|
||||
WHERE scheduled_jobs.user_uuid IS NOT DISTINCT FROM EXCLUDED.user_uuid
|
||||
RETURNING *
|
||||
""",
|
||||
values,
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def get_job(job_id, cursor=None):
|
||||
"""Return one job by UUID."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute("SELECT * FROM scheduled_jobs WHERE id = %s", (job_id,))
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def list_jobs(user_uuid=None, status=None, job_type=None, limit=100, cursor=None):
|
||||
"""List jobs newest first, optionally filtered by owner, status, or type."""
|
||||
limit = _positive(limit, "limit", 500)
|
||||
clauses = []
|
||||
params = []
|
||||
if user_uuid is not None:
|
||||
clauses.append("user_uuid = %s")
|
||||
params.append(user_uuid)
|
||||
if status is not None:
|
||||
clauses.append("status = %s")
|
||||
params.append(status)
|
||||
if job_type is not None:
|
||||
clauses.append("job_type = %s")
|
||||
params.append(job_type)
|
||||
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
|
||||
params.append(limit)
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
f"""
|
||||
SELECT * FROM scheduled_jobs
|
||||
{where}
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
""",
|
||||
params,
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
|
||||
|
||||
def claim_due_jobs(
|
||||
worker_id,
|
||||
limit=10,
|
||||
lease_seconds=DEFAULT_LEASE_SECONDS,
|
||||
job_types=None,
|
||||
cursor=None,
|
||||
):
|
||||
"""Atomically lease due jobs using row locks that skip other workers."""
|
||||
if not worker_id:
|
||||
raise ValueError("worker_id is required")
|
||||
limit = _positive(limit, "limit", 100)
|
||||
lease_seconds = _positive(lease_seconds, "lease_seconds")
|
||||
if isinstance(job_types, str):
|
||||
job_types = [job_types]
|
||||
elif job_types is not None:
|
||||
job_types = list(job_types)
|
||||
if not job_types:
|
||||
return []
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
WITH exhausted AS (
|
||||
SELECT id
|
||||
FROM scheduled_jobs
|
||||
WHERE status = 'running'
|
||||
AND (lease_until IS NULL OR lease_until <= NOW())
|
||||
AND attempts >= max_attempts
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE scheduled_jobs AS job
|
||||
SET status = 'failed',
|
||||
lease_until = NULL,
|
||||
leased_by = NULL,
|
||||
last_error = COALESCE(last_error, 'lease expired'),
|
||||
updated_at = NOW()
|
||||
FROM exhausted
|
||||
WHERE job.id = exhausted.id
|
||||
"""
|
||||
)
|
||||
active_cursor.execute(
|
||||
"""
|
||||
WITH candidates AS (
|
||||
SELECT id
|
||||
FROM scheduled_jobs
|
||||
WHERE run_at <= NOW()
|
||||
AND attempts < max_attempts
|
||||
AND (
|
||||
status = 'pending'
|
||||
OR (
|
||||
status = 'running'
|
||||
AND (lease_until IS NULL OR lease_until <= NOW())
|
||||
)
|
||||
)
|
||||
AND (%(job_types)s IS NULL OR job_type = ANY(%(job_types)s))
|
||||
ORDER BY run_at, created_at
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT %(limit)s
|
||||
)
|
||||
UPDATE scheduled_jobs AS job
|
||||
SET status = 'running',
|
||||
attempts = job.attempts + 1,
|
||||
leased_by = %(worker_id)s,
|
||||
lease_until = NOW() + (%(lease_seconds)s * INTERVAL '1 second'),
|
||||
updated_at = NOW()
|
||||
FROM candidates
|
||||
WHERE job.id = candidates.id
|
||||
RETURNING job.*
|
||||
""",
|
||||
{
|
||||
"job_types": job_types,
|
||||
"limit": limit,
|
||||
"worker_id": worker_id,
|
||||
"lease_seconds": lease_seconds,
|
||||
},
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
|
||||
|
||||
def renew_job_lease(job_id, worker_id, lease_seconds=DEFAULT_LEASE_SECONDS, cursor=None):
|
||||
"""Extend a lease only while it is owned by the requesting worker."""
|
||||
lease_seconds = _positive(lease_seconds, "lease_seconds")
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET lease_until = NOW() + (%s * INTERVAL '1 second'), updated_at = NOW()
|
||||
WHERE id = %s AND status = 'running' AND leased_by = %s
|
||||
RETURNING *
|
||||
""",
|
||||
(lease_seconds, job_id, worker_id),
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def complete_job(job_id, worker_id, cursor=None):
|
||||
"""Mark a job complete when its lease is still owned by the worker."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'completed', completed_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL, last_error = NULL
|
||||
WHERE id = %s AND status = 'running' AND leased_by = %s
|
||||
RETURNING *
|
||||
""",
|
||||
(job_id, worker_id),
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def fail_job(
|
||||
job_id,
|
||||
worker_id,
|
||||
error,
|
||||
retry_seconds=DEFAULT_RETRY_SECONDS,
|
||||
max_retry_seconds=MAX_RETRY_SECONDS,
|
||||
cursor=None,
|
||||
):
|
||||
"""Fail or reschedule an owned job using bounded exponential backoff."""
|
||||
retry_seconds = _positive(retry_seconds, "retry_seconds")
|
||||
max_retry_seconds = _positive(max_retry_seconds, "max_retry_seconds")
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
SELECT attempts, max_attempts
|
||||
FROM scheduled_jobs
|
||||
WHERE id = %s AND status = 'running' AND leased_by = %s
|
||||
FOR UPDATE
|
||||
""",
|
||||
(job_id, worker_id),
|
||||
)
|
||||
current = active_cursor.fetchone()
|
||||
if not current:
|
||||
return None
|
||||
|
||||
exhausted = current["attempts"] >= current["max_attempts"]
|
||||
delay = min(
|
||||
max_retry_seconds,
|
||||
retry_seconds * (2 ** min(max(current["attempts"] - 1, 0), 30)),
|
||||
)
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = %(status)s,
|
||||
run_at = CASE
|
||||
WHEN %(exhausted)s THEN run_at
|
||||
ELSE NOW() + (%(delay)s * INTERVAL '1 second')
|
||||
END,
|
||||
leased_by = NULL,
|
||||
lease_until = NULL,
|
||||
last_error = %(error)s,
|
||||
updated_at = NOW()
|
||||
WHERE id = %(job_id)s
|
||||
RETURNING *
|
||||
""",
|
||||
{
|
||||
"status": "failed" if exhausted else "pending",
|
||||
"exhausted": exhausted,
|
||||
"delay": delay,
|
||||
"error": str(error)[:4000],
|
||||
"job_id": job_id,
|
||||
},
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def cancel_job(job_id, user_uuid=None, cursor=None):
|
||||
"""Cancel one unfinished job, optionally enforcing its owner."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE id = %(job_id)s
|
||||
AND status IN ('pending', 'running')
|
||||
AND (%(user_uuid)s IS NULL OR user_uuid = %(user_uuid)s)
|
||||
RETURNING *
|
||||
""",
|
||||
{"job_id": job_id, "user_uuid": user_uuid},
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def cancel_jobs(
|
||||
user_uuid=None, job_type=None, idempotency_key=None, cursor=None
|
||||
):
|
||||
"""Cancel matching unfinished jobs; at least one filter is required."""
|
||||
filters = {
|
||||
"user_uuid": user_uuid,
|
||||
"job_type": job_type,
|
||||
"idempotency_key": idempotency_key,
|
||||
}
|
||||
clauses = [f"{name} = %({name})s" for name, value in filters.items() if value is not None]
|
||||
if not clauses:
|
||||
raise ValueError("at least one cancellation filter is required")
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
f"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE status IN ('pending', 'running')
|
||||
AND {' AND '.join(clauses)}
|
||||
RETURNING *
|
||||
""",
|
||||
filters,
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
76
core/manage.py
Normal file
76
core/manage.py
Normal file
@@ -0,0 +1,76 @@
|
||||
"""Small administrative command-line tools for framework operators."""
|
||||
|
||||
import argparse
|
||||
import uuid
|
||||
|
||||
import psycopg2
|
||||
|
||||
from core import identity
|
||||
|
||||
|
||||
def _discordID(value):
|
||||
value = value.strip()
|
||||
if not value or not value.isdigit():
|
||||
raise argparse.ArgumentTypeError("Discord ID must contain only digits")
|
||||
return value
|
||||
|
||||
|
||||
def _userUUID(value):
|
||||
try:
|
||||
return str(uuid.UUID(value))
|
||||
except (AttributeError, TypeError, ValueError) as error:
|
||||
raise argparse.ArgumentTypeError("user UUID must be a valid UUID") from error
|
||||
|
||||
|
||||
def _username(value):
|
||||
value = value.strip()
|
||||
if not value:
|
||||
raise argparse.ArgumentTypeError("username cannot be empty")
|
||||
return value
|
||||
|
||||
|
||||
def _linkDiscord(args):
|
||||
try:
|
||||
linked = identity.linkDiscordUser(
|
||||
args.discord_id,
|
||||
userUUID=args.user_uuid,
|
||||
username=args.username,
|
||||
)
|
||||
except ValueError as error:
|
||||
args.command_parser.error(str(error))
|
||||
except psycopg2.Error:
|
||||
args.command_parser.exit(
|
||||
1,
|
||||
"link-discord failed: database operation failed\n",
|
||||
)
|
||||
|
||||
print(
|
||||
"linked Discord user "
|
||||
f"{linked['provider_user_id']} to user {linked['user_uuid']}"
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def buildParser():
|
||||
parser = argparse.ArgumentParser(description="Manage framework data")
|
||||
commands = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
linkParser = commands.add_parser(
|
||||
"link-discord",
|
||||
help="link an existing user to a Discord account",
|
||||
)
|
||||
linkParser.add_argument("--discord-id", required=True, type=_discordID)
|
||||
userSelector = linkParser.add_mutually_exclusive_group(required=True)
|
||||
userSelector.add_argument("--user-uuid", type=_userUUID)
|
||||
userSelector.add_argument("--username", type=_username)
|
||||
linkParser.set_defaults(handler=_linkDiscord, command_parser=linkParser)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
args = buildParser().parse_args(argv)
|
||||
return args.handler(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
202
core/migrations/__init__.py
Normal file
202
core/migrations/__init__.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""Versioned PostgreSQL migrations for the framework and feature modules."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
import re
|
||||
|
||||
import psycopg2.extras
|
||||
|
||||
from core import postgres
|
||||
|
||||
|
||||
MIGRATION_NAME = re.compile(r"^(?P<version>\d+)(?:[_-].*)?\.sql$")
|
||||
LOCK_NAME = "llm-bot-framework:schema-migrations"
|
||||
|
||||
|
||||
class MigrationError(RuntimeError):
|
||||
"""Raised when migration history is inconsistent or cannot be applied."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Migration:
|
||||
namespace: str
|
||||
version: int
|
||||
name: str
|
||||
path: Path
|
||||
checksum: str
|
||||
sql: str
|
||||
|
||||
|
||||
def _project_root(project_root=None):
|
||||
if project_root is not None:
|
||||
return Path(project_root).resolve()
|
||||
return Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def _read_namespace(namespace, directory):
|
||||
migrations = []
|
||||
if not directory.is_dir():
|
||||
return migrations
|
||||
|
||||
for path in sorted(directory.glob("*.sql")):
|
||||
match = MIGRATION_NAME.match(path.name)
|
||||
if not match:
|
||||
raise MigrationError(f"Invalid migration filename: {path}")
|
||||
sql = path.read_text(encoding="utf-8")
|
||||
migrations.append(
|
||||
Migration(
|
||||
namespace=namespace,
|
||||
version=int(match.group("version")),
|
||||
name=path.name,
|
||||
path=path,
|
||||
checksum=sha256(sql.encode("utf-8")).hexdigest(),
|
||||
sql=sql,
|
||||
)
|
||||
)
|
||||
return migrations
|
||||
|
||||
|
||||
def discover_migrations(project_root=None):
|
||||
"""Return core and feature migrations in deterministic application order."""
|
||||
root = _project_root(project_root)
|
||||
migrations = _read_namespace("core", root / "config" / "migrations")
|
||||
|
||||
modules_root = root / "modules"
|
||||
if modules_root.is_dir():
|
||||
for module_path in sorted(modules_root.iterdir(), key=lambda path: path.name):
|
||||
if module_path.is_dir() and not module_path.name.startswith("_"):
|
||||
migrations.extend(
|
||||
_read_namespace(module_path.name, module_path / "migrations")
|
||||
)
|
||||
|
||||
seen = set()
|
||||
for migration in migrations:
|
||||
key = (migration.namespace, migration.version)
|
||||
if key in seen:
|
||||
raise MigrationError(
|
||||
f"Duplicate migration {migration.namespace}:{migration.version}"
|
||||
)
|
||||
seen.add(key)
|
||||
|
||||
return sorted(
|
||||
migrations,
|
||||
key=lambda item: (
|
||||
item.namespace != "core",
|
||||
item.namespace,
|
||||
item.version,
|
||||
item.name,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _ensure_history_table(cursor):
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
namespace VARCHAR(255) NOT NULL,
|
||||
version INTEGER NOT NULL,
|
||||
checksum CHAR(64) NOT NULL,
|
||||
applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
PRIMARY KEY (namespace, version)
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _applied_migrations(cursor):
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT namespace, version, checksum, applied_at
|
||||
FROM schema_migrations
|
||||
ORDER BY namespace, version
|
||||
"""
|
||||
)
|
||||
return {
|
||||
(row["namespace"], row["version"]): dict(row) for row in cursor.fetchall()
|
||||
}
|
||||
|
||||
|
||||
def upgrade(project_root=None):
|
||||
"""Apply pending migrations transactionally under a PostgreSQL advisory lock."""
|
||||
migrations = discover_migrations(project_root)
|
||||
applied_now = []
|
||||
|
||||
with postgres.get_connection() as connection:
|
||||
with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
cursor.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (LOCK_NAME,))
|
||||
_ensure_history_table(cursor)
|
||||
applied = _applied_migrations(cursor)
|
||||
|
||||
for migration in migrations:
|
||||
key = (migration.namespace, migration.version)
|
||||
previous = applied.get(key)
|
||||
if previous:
|
||||
if previous["checksum"] != migration.checksum:
|
||||
raise MigrationError(
|
||||
"Applied migration checksum changed: "
|
||||
f"{migration.namespace}:{migration.version}"
|
||||
)
|
||||
continue
|
||||
|
||||
cursor.execute(migration.sql)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO schema_migrations (namespace, version, checksum)
|
||||
VALUES (%s, %s, %s)
|
||||
""",
|
||||
(migration.namespace, migration.version, migration.checksum),
|
||||
)
|
||||
applied_now.append(migration)
|
||||
|
||||
return applied_now
|
||||
|
||||
|
||||
def migration_status(project_root=None):
|
||||
"""Return applied, pending, changed, and source-missing migration records."""
|
||||
migrations = discover_migrations(project_root)
|
||||
with postgres.get_connection() as connection:
|
||||
with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
|
||||
cursor.execute(
|
||||
"SELECT to_regclass('public.schema_migrations') AS table_name"
|
||||
)
|
||||
if cursor.fetchone()["table_name"] is None:
|
||||
applied = {}
|
||||
else:
|
||||
applied = _applied_migrations(cursor)
|
||||
|
||||
status = []
|
||||
source_keys = set()
|
||||
for migration in migrations:
|
||||
key = (migration.namespace, migration.version)
|
||||
source_keys.add(key)
|
||||
record = applied.get(key)
|
||||
state = "pending"
|
||||
applied_at = None
|
||||
if record:
|
||||
state = "applied" if record["checksum"] == migration.checksum else "changed"
|
||||
applied_at = record["applied_at"]
|
||||
status.append(
|
||||
{
|
||||
"namespace": migration.namespace,
|
||||
"version": migration.version,
|
||||
"name": migration.name,
|
||||
"checksum": migration.checksum,
|
||||
"applied_at": applied_at,
|
||||
"state": state,
|
||||
}
|
||||
)
|
||||
|
||||
for key, record in sorted(applied.items()):
|
||||
if key not in source_keys:
|
||||
status.append(
|
||||
{
|
||||
"namespace": key[0],
|
||||
"version": key[1],
|
||||
"name": None,
|
||||
"checksum": record["checksum"],
|
||||
"applied_at": record["applied_at"],
|
||||
"state": "missing",
|
||||
}
|
||||
)
|
||||
return status
|
||||
37
core/migrations/__main__.py
Normal file
37
core/migrations/__main__.py
Normal file
@@ -0,0 +1,37 @@
|
||||
"""Command-line entry point for database migrations."""
|
||||
|
||||
import argparse
|
||||
|
||||
from core.migrations import MigrationError, migration_status, upgrade
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description="Manage database migrations")
|
||||
parser.add_argument("command", choices=("upgrade", "status"))
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
try:
|
||||
if args.command == "upgrade":
|
||||
applied = upgrade()
|
||||
for migration in applied:
|
||||
print(f"applied {migration.namespace}:{migration.version} {migration.name}")
|
||||
if not applied:
|
||||
print("database is up to date")
|
||||
return 0
|
||||
|
||||
records = migration_status()
|
||||
if not records:
|
||||
print("no migrations found")
|
||||
for record in records:
|
||||
name = record["name"] or "<source missing>"
|
||||
print(
|
||||
f"{record['state']:<8} "
|
||||
f"{record['namespace']}:{record['version']} {name}"
|
||||
)
|
||||
return 1 if any(row["state"] == "changed" for row in records) else 0
|
||||
except MigrationError as error:
|
||||
parser.exit(1, f"migration error: {error}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
105
core/notifications.py
Normal file
105
core/notifications.py
Normal file
@@ -0,0 +1,105 @@
|
||||
"""Optional Discord-webhook and ntfy notification helpers."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from urllib.parse import quote, urlparse
|
||||
import uuid
|
||||
|
||||
import requests
|
||||
|
||||
import core.postgres as postgres
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
REQUEST_TIMEOUT = float(os.getenv("NOTIFICATION_TIMEOUT", 10))
|
||||
|
||||
|
||||
def _sendToEnabledChannels(notif_settings, message):
|
||||
"""Send to enabled channels and return True when at least one succeeds."""
|
||||
if not isinstance(notif_settings, dict):
|
||||
return False
|
||||
sent = False
|
||||
if notif_settings.get("discord_enabled") and notif_settings.get(
|
||||
"discord_webhook"
|
||||
):
|
||||
sent = discord.send(notif_settings["discord_webhook"], message) or sent
|
||||
if notif_settings.get("ntfy_enabled") and notif_settings.get("ntfy_topic"):
|
||||
sent = ntfy.send(notif_settings["ntfy_topic"], message) or sent
|
||||
return sent
|
||||
|
||||
|
||||
def getNotificationSettings(userUUID):
|
||||
return postgres.select_one("notifications", {"user_uuid": userUUID}) or False
|
||||
|
||||
|
||||
def setNotificationSettings(userUUID, data_dict):
|
||||
if not isinstance(data_dict, dict):
|
||||
return False
|
||||
allowed = {
|
||||
"discord_webhook",
|
||||
"discord_enabled",
|
||||
"ntfy_topic",
|
||||
"ntfy_enabled",
|
||||
}
|
||||
updates = {key: value for key, value in data_dict.items() if key in allowed}
|
||||
if not updates:
|
||||
return False
|
||||
if updates.get("discord_webhook"):
|
||||
_validateDiscordWebhook(updates["discord_webhook"])
|
||||
|
||||
existing = postgres.select_one("notifications", {"user_uuid": userUUID})
|
||||
if existing:
|
||||
postgres.update("notifications", updates, {"user_uuid": userUUID})
|
||||
else:
|
||||
updates["id"] = str(uuid.uuid4())
|
||||
updates["user_uuid"] = userUUID
|
||||
postgres.insert("notifications", updates)
|
||||
return True
|
||||
|
||||
|
||||
def _validateDiscordWebhook(webhookURL):
|
||||
parsed = urlparse(webhookURL)
|
||||
allowedHosts = {"discord.com", "canary.discord.com", "ptb.discord.com"}
|
||||
if parsed.scheme != "https" or parsed.hostname not in allowedHosts:
|
||||
raise ValueError("Discord webhook must use an official HTTPS Discord host")
|
||||
if not parsed.path.startswith("/api/webhooks/"):
|
||||
raise ValueError("Invalid Discord webhook path")
|
||||
return webhookURL
|
||||
|
||||
|
||||
class discord:
|
||||
@staticmethod
|
||||
def send(webhook_url, message):
|
||||
try:
|
||||
_validateDiscordWebhook(webhook_url)
|
||||
response = requests.post(
|
||||
webhook_url,
|
||||
json={"content": str(message)},
|
||||
timeout=REQUEST_TIMEOUT,
|
||||
)
|
||||
return response.status_code in {200, 204}
|
||||
except (ValueError, requests.RequestException) as error:
|
||||
logger.warning("Discord webhook delivery failed: %s", error)
|
||||
return False
|
||||
|
||||
|
||||
class ntfy:
|
||||
@staticmethod
|
||||
def send(topic, message):
|
||||
if not isinstance(topic, str) or not topic.strip():
|
||||
return False
|
||||
baseURL = os.getenv("NTFY_BASE_URL", "https://ntfy.sh").rstrip("/")
|
||||
headers = {}
|
||||
if os.getenv("NTFY_TOKEN"):
|
||||
headers["Authorization"] = f"Bearer {os.environ['NTFY_TOKEN']}"
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{baseURL}/{quote(topic.strip(), safe='')}",
|
||||
data=str(message).encode("utf-8"),
|
||||
headers=headers,
|
||||
timeout=REQUEST_TIMEOUT,
|
||||
)
|
||||
return 200 <= response.status_code < 300
|
||||
except requests.RequestException as error:
|
||||
logger.warning("ntfy delivery failed: %s", error)
|
||||
return False
|
||||
357
core/outbox.py
Normal file
357
core/outbox.py
Normal file
@@ -0,0 +1,357 @@
|
||||
"""Durable, lease-based outbound message operations."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
import uuid
|
||||
|
||||
from psycopg2.extras import Json
|
||||
|
||||
from core import postgres
|
||||
|
||||
|
||||
DEFAULT_MAX_ATTEMPTS = 3
|
||||
DEFAULT_LEASE_SECONDS = 300
|
||||
DEFAULT_RETRY_SECONDS = 30
|
||||
MAX_RETRY_SECONDS = 900
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _using_cursor(cursor=None):
|
||||
if cursor is not None:
|
||||
yield cursor
|
||||
return
|
||||
with postgres.get_cursor() as owned_cursor:
|
||||
yield owned_cursor
|
||||
|
||||
|
||||
def _timestamp(value, field="timestamp"):
|
||||
if isinstance(value, str):
|
||||
value = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
if not isinstance(value, datetime):
|
||||
raise ValueError(f"{field} must be a datetime or ISO-8601 string")
|
||||
if value.tzinfo is None:
|
||||
value = value.replace(tzinfo=timezone.utc)
|
||||
return value.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _positive(value, field, maximum=None):
|
||||
try:
|
||||
value = int(value)
|
||||
except (TypeError, ValueError) as error:
|
||||
raise ValueError(f"{field} must be a whole number") from error
|
||||
if value < 1 or (maximum is not None and value > maximum):
|
||||
suffix = f" and at most {maximum}" if maximum is not None else ""
|
||||
raise ValueError(f"{field} must be at least 1{suffix}")
|
||||
return value
|
||||
|
||||
|
||||
def _row(cursor):
|
||||
record = cursor.fetchone()
|
||||
return dict(record) if record else None
|
||||
|
||||
|
||||
def enqueue_message(
|
||||
user_uuid,
|
||||
channel,
|
||||
payload,
|
||||
idempotency_key,
|
||||
available_at=None,
|
||||
max_attempts=DEFAULT_MAX_ATTEMPTS,
|
||||
message_id=None,
|
||||
cursor=None,
|
||||
):
|
||||
"""Queue a message, returning the existing row for a repeated key."""
|
||||
if not user_uuid:
|
||||
raise ValueError("user_uuid is required")
|
||||
if not isinstance(channel, str) or not channel.strip():
|
||||
raise ValueError("channel is required")
|
||||
if not idempotency_key:
|
||||
raise ValueError("idempotency_key is required")
|
||||
if payload is None:
|
||||
raise ValueError("payload is required")
|
||||
max_attempts = _positive(max_attempts, "max_attempts", 100)
|
||||
available_at = available_at or datetime.now(timezone.utc)
|
||||
values = {
|
||||
"id": str(message_id or uuid.uuid4()),
|
||||
"user_uuid": user_uuid,
|
||||
"channel": channel.strip(),
|
||||
"payload": Json(payload),
|
||||
"idempotency_key": idempotency_key,
|
||||
"available_at": _timestamp(available_at, "available_at"),
|
||||
"max_attempts": max_attempts,
|
||||
}
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
INSERT INTO outbound_messages (
|
||||
id, user_uuid, channel, payload, idempotency_key,
|
||||
available_at, max_attempts
|
||||
) VALUES (
|
||||
%(id)s, %(user_uuid)s, %(channel)s, %(payload)s,
|
||||
%(idempotency_key)s, %(available_at)s, %(max_attempts)s
|
||||
)
|
||||
ON CONFLICT (idempotency_key)
|
||||
DO UPDATE SET idempotency_key = EXCLUDED.idempotency_key
|
||||
WHERE outbound_messages.user_uuid = EXCLUDED.user_uuid
|
||||
AND outbound_messages.channel = EXCLUDED.channel
|
||||
RETURNING *
|
||||
""",
|
||||
values,
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def get_message(message_id, cursor=None):
|
||||
"""Return one outbound message by UUID."""
|
||||
try:
|
||||
message_id = str(uuid.UUID(str(message_id)))
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
return None
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"SELECT * FROM outbound_messages WHERE id = %s", (message_id,)
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def list_messages(user_uuid=None, status=None, channel=None, limit=100, cursor=None):
|
||||
"""List outbound messages newest first with optional filters."""
|
||||
limit = _positive(limit, "limit", 500)
|
||||
clauses = []
|
||||
params = []
|
||||
if user_uuid is not None:
|
||||
clauses.append("user_uuid = %s")
|
||||
params.append(user_uuid)
|
||||
if status is not None:
|
||||
clauses.append("status = %s")
|
||||
params.append(status)
|
||||
if channel is not None:
|
||||
clauses.append("channel = %s")
|
||||
params.append(channel)
|
||||
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
|
||||
params.append(limit)
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
f"""
|
||||
SELECT * FROM outbound_messages
|
||||
{where}
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
""",
|
||||
params,
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
|
||||
|
||||
def claim_messages(
|
||||
worker_id,
|
||||
channel=None,
|
||||
limit=10,
|
||||
lease_seconds=DEFAULT_LEASE_SECONDS,
|
||||
cursor=None,
|
||||
):
|
||||
"""Atomically lease deliverable messages while skipping other workers."""
|
||||
if not worker_id:
|
||||
raise ValueError("worker_id is required")
|
||||
limit = _positive(limit, "limit", 100)
|
||||
lease_seconds = _positive(lease_seconds, "lease_seconds")
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
WITH exhausted AS (
|
||||
SELECT id
|
||||
FROM outbound_messages
|
||||
WHERE status = 'delivering'
|
||||
AND (lease_until IS NULL OR lease_until <= NOW())
|
||||
AND attempts >= max_attempts
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE outbound_messages AS message
|
||||
SET status = 'failed',
|
||||
lease_until = NULL,
|
||||
leased_by = NULL,
|
||||
last_error = COALESCE(last_error, 'lease expired'),
|
||||
updated_at = NOW()
|
||||
FROM exhausted
|
||||
WHERE message.id = exhausted.id
|
||||
"""
|
||||
)
|
||||
active_cursor.execute(
|
||||
"""
|
||||
WITH candidates AS (
|
||||
SELECT id
|
||||
FROM outbound_messages
|
||||
WHERE available_at <= NOW()
|
||||
AND attempts < max_attempts
|
||||
AND (
|
||||
status = 'pending'
|
||||
OR (
|
||||
status = 'delivering'
|
||||
AND (lease_until IS NULL OR lease_until <= NOW())
|
||||
)
|
||||
)
|
||||
AND (%(channel)s IS NULL OR channel = %(channel)s)
|
||||
ORDER BY available_at, created_at
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT %(limit)s
|
||||
)
|
||||
UPDATE outbound_messages AS message
|
||||
SET status = 'delivering',
|
||||
attempts = message.attempts + 1,
|
||||
leased_by = %(worker_id)s,
|
||||
lease_until = NOW() + (%(lease_seconds)s * INTERVAL '1 second'),
|
||||
updated_at = NOW()
|
||||
FROM candidates
|
||||
WHERE message.id = candidates.id
|
||||
RETURNING message.*
|
||||
""",
|
||||
{
|
||||
"channel": channel,
|
||||
"limit": limit,
|
||||
"worker_id": worker_id,
|
||||
"lease_seconds": lease_seconds,
|
||||
},
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
|
||||
|
||||
def renew_message_lease(
|
||||
message_id, worker_id, lease_seconds=DEFAULT_LEASE_SECONDS, cursor=None
|
||||
):
|
||||
"""Extend a delivery lease owned by the requesting worker."""
|
||||
lease_seconds = _positive(lease_seconds, "lease_seconds")
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE outbound_messages
|
||||
SET lease_until = NOW() + (%s * INTERVAL '1 second'), updated_at = NOW()
|
||||
WHERE id = %s AND status = 'delivering' AND leased_by = %s
|
||||
RETURNING *
|
||||
""",
|
||||
(lease_seconds, message_id, worker_id),
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def mark_delivered(
|
||||
message_id, worker_id, external_message_id=None, cursor=None
|
||||
):
|
||||
"""Mark a message delivered when its lease is owned by the worker."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE outbound_messages
|
||||
SET status = 'delivered', delivered_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL, last_error = NULL,
|
||||
external_message_id = %s
|
||||
WHERE id = %s AND status = 'delivering' AND leased_by = %s
|
||||
RETURNING *
|
||||
""",
|
||||
(external_message_id, message_id, worker_id),
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def retry_message(
|
||||
message_id,
|
||||
worker_id,
|
||||
error,
|
||||
retry_seconds=DEFAULT_RETRY_SECONDS,
|
||||
max_retry_seconds=MAX_RETRY_SECONDS,
|
||||
cursor=None,
|
||||
):
|
||||
"""Fail or reschedule an owned delivery with exponential backoff."""
|
||||
retry_seconds = _positive(retry_seconds, "retry_seconds")
|
||||
max_retry_seconds = _positive(max_retry_seconds, "max_retry_seconds")
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
SELECT attempts, max_attempts
|
||||
FROM outbound_messages
|
||||
WHERE id = %s AND status = 'delivering' AND leased_by = %s
|
||||
FOR UPDATE
|
||||
""",
|
||||
(message_id, worker_id),
|
||||
)
|
||||
current = active_cursor.fetchone()
|
||||
if not current:
|
||||
return None
|
||||
|
||||
exhausted = current["attempts"] >= current["max_attempts"]
|
||||
delay = min(
|
||||
max_retry_seconds,
|
||||
retry_seconds * (2 ** min(max(current["attempts"] - 1, 0), 30)),
|
||||
)
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE outbound_messages
|
||||
SET status = %(status)s,
|
||||
available_at = CASE
|
||||
WHEN %(exhausted)s THEN available_at
|
||||
ELSE NOW() + (%(delay)s * INTERVAL '1 second')
|
||||
END,
|
||||
leased_by = NULL,
|
||||
lease_until = NULL,
|
||||
last_error = %(error)s,
|
||||
updated_at = NOW()
|
||||
WHERE id = %(message_id)s
|
||||
RETURNING *
|
||||
""",
|
||||
{
|
||||
"status": "failed" if exhausted else "pending",
|
||||
"exhausted": exhausted,
|
||||
"delay": delay,
|
||||
"error": str(error)[:4000],
|
||||
"message_id": message_id,
|
||||
},
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def cancel_message(message_id, user_uuid=None, cursor=None):
|
||||
"""Cancel one undelivered message, optionally enforcing its owner."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE outbound_messages
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE id = %(message_id)s
|
||||
AND status IN ('pending', 'delivering')
|
||||
AND (%(user_uuid)s IS NULL OR user_uuid = %(user_uuid)s)
|
||||
RETURNING *
|
||||
""",
|
||||
{"message_id": message_id, "user_uuid": user_uuid},
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def cancel_messages(
|
||||
user_uuid=None, channel=None, idempotency_key=None, cursor=None
|
||||
):
|
||||
"""Cancel matching undelivered messages; at least one filter is required."""
|
||||
filters = {
|
||||
"user_uuid": user_uuid,
|
||||
"channel": channel,
|
||||
"idempotency_key": idempotency_key,
|
||||
}
|
||||
clauses = [f"{name} = %({name})s" for name, value in filters.items() if value is not None]
|
||||
if not clauses:
|
||||
raise ValueError("at least one cancellation filter is required")
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
f"""
|
||||
UPDATE outbound_messages
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE status IN ('pending', 'delivering')
|
||||
AND {' AND '.join(clauses)}
|
||||
RETURNING *
|
||||
""",
|
||||
filters,
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
289
core/postgres.py
Normal file
289
core/postgres.py
Normal file
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
postgres.py - Small parameterized PostgreSQL CRUD layer
|
||||
|
||||
Connection configuration is read from DB_HOST, DB_PORT, DB_NAME, DB_USER,
|
||||
and DB_PASS. Raw SQL remains available through execute() for domain services.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
from dotenv import load_dotenv
|
||||
|
||||
|
||||
load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
|
||||
|
||||
IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
|
||||
def _get_config():
|
||||
return {
|
||||
"host": os.environ.get("DB_HOST", "localhost"),
|
||||
"port": int(os.environ.get("DB_PORT", 5432)),
|
||||
"dbname": os.environ.get("DB_NAME", "app"),
|
||||
"user": os.environ.get("DB_USER", "app"),
|
||||
"password": os.environ.get("DB_PASS", ""),
|
||||
}
|
||||
|
||||
|
||||
def _safe_id(name):
|
||||
if not isinstance(name, str) or not IDENTIFIER.fullmatch(name):
|
||||
raise ValueError(f"Invalid SQL identifier: {name}")
|
||||
return f'"{name}"'
|
||||
|
||||
|
||||
def _build_where(where, prefix=""):
|
||||
if not isinstance(where, dict):
|
||||
raise ValueError("where must be a dictionary")
|
||||
clauses = []
|
||||
params = {}
|
||||
for index, (column, value) in enumerate(where.items()):
|
||||
paramName = f"{prefix}{column}_{index}"
|
||||
safeColumn = _safe_id(column)
|
||||
|
||||
if isinstance(value, tuple) and len(value) == 2:
|
||||
operator, operand = value
|
||||
operator = str(operator).upper()
|
||||
allowed = {"=", "!=", "<", ">", "<=", ">=", "LIKE", "ILIKE", "IN"}
|
||||
if operator not in allowed:
|
||||
raise ValueError(f"Unsupported operator: {operator}")
|
||||
if operator == "IN":
|
||||
values = list(operand)
|
||||
if not values:
|
||||
clauses.append("FALSE")
|
||||
continue
|
||||
placeholders = []
|
||||
for itemIndex, item in enumerate(values):
|
||||
itemName = f"{paramName}_{itemIndex}"
|
||||
placeholders.append(f"%({itemName})s")
|
||||
params[itemName] = item
|
||||
clauses.append(f"{safeColumn} IN ({', '.join(placeholders)})")
|
||||
else:
|
||||
clauses.append(f"{safeColumn} {operator} %({paramName})s")
|
||||
params[paramName] = operand
|
||||
elif value is None:
|
||||
clauses.append(f"{safeColumn} IS NULL")
|
||||
else:
|
||||
clauses.append(f"{safeColumn} = %({paramName})s")
|
||||
params[paramName] = value
|
||||
return " AND ".join(clauses), params
|
||||
|
||||
|
||||
def _order_clause(orderBy):
|
||||
if not orderBy:
|
||||
return ""
|
||||
items = orderBy if isinstance(orderBy, (list, tuple)) else str(orderBy).split(",")
|
||||
safeItems = []
|
||||
for item in items:
|
||||
if isinstance(item, (list, tuple)):
|
||||
if len(item) != 2:
|
||||
raise ValueError("order tuple must contain column and direction")
|
||||
column, direction = item
|
||||
else:
|
||||
parts = str(item).strip().split()
|
||||
if not parts or len(parts) > 2:
|
||||
raise ValueError(f"Invalid order expression: {item}")
|
||||
column = parts[0]
|
||||
direction = parts[1] if len(parts) == 2 else "ASC"
|
||||
direction = str(direction).upper()
|
||||
if direction not in {"ASC", "DESC"}:
|
||||
raise ValueError(f"Invalid order direction: {direction}")
|
||||
safeItems.append(f"{_safe_id(column)} {direction}")
|
||||
return ", ".join(safeItems)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_connection():
|
||||
connection = psycopg2.connect(**_get_config())
|
||||
try:
|
||||
yield connection
|
||||
connection.commit()
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_cursor(dict_cursor=True):
|
||||
with get_connection() as connection:
|
||||
factory = psycopg2.extras.RealDictCursor if dict_cursor else None
|
||||
cursor = connection.cursor(cursor_factory=factory)
|
||||
try:
|
||||
yield cursor
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
def insert(table, data):
|
||||
if not isinstance(data, dict) or not data:
|
||||
raise ValueError("insert data cannot be empty")
|
||||
columns = list(data.keys())
|
||||
query = f"""
|
||||
INSERT INTO {_safe_id(table)} ({', '.join(_safe_id(col) for col in columns)})
|
||||
VALUES ({', '.join(f'%({col})s' for col in columns)})
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, data)
|
||||
record = cursor.fetchone()
|
||||
return dict(record) if record else None
|
||||
|
||||
|
||||
def select(table, where=None, order_by=None, limit=None, offset=None):
|
||||
query = f"SELECT * FROM {_safe_id(table)}"
|
||||
params = {}
|
||||
if where:
|
||||
clauses, params = _build_where(where)
|
||||
query += f" WHERE {clauses}"
|
||||
orderClause = _order_clause(order_by)
|
||||
if orderClause:
|
||||
query += f" ORDER BY {orderClause}"
|
||||
if limit is not None:
|
||||
limit = int(limit)
|
||||
if limit < 0:
|
||||
raise ValueError("limit cannot be negative")
|
||||
query += " LIMIT %(query_limit)s"
|
||||
params["query_limit"] = limit
|
||||
if offset is not None:
|
||||
offset = int(offset)
|
||||
if offset < 0:
|
||||
raise ValueError("offset cannot be negative")
|
||||
query += " OFFSET %(query_offset)s"
|
||||
params["query_offset"] = offset
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, params)
|
||||
return [dict(record) for record in cursor.fetchall()]
|
||||
|
||||
|
||||
def select_one(table, where):
|
||||
records = select(table, where=where, limit=1)
|
||||
return records[0] if records else None
|
||||
|
||||
|
||||
def update(table, data, where):
|
||||
if not isinstance(data, dict) or not data:
|
||||
raise ValueError("update data cannot be empty")
|
||||
if not isinstance(where, dict) or not where:
|
||||
raise ValueError("update requires a non-empty where clause")
|
||||
setClause = ", ".join(f"{_safe_id(col)} = %(set_{col})s" for col in data)
|
||||
params = {f"set_{col}": value for col, value in data.items()}
|
||||
whereClause, whereParams = _build_where(where, prefix="where_")
|
||||
params.update(whereParams)
|
||||
query = f"""
|
||||
UPDATE {_safe_id(table)} SET {setClause}
|
||||
WHERE {whereClause}
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, params)
|
||||
return [dict(record) for record in cursor.fetchall()]
|
||||
|
||||
|
||||
def delete(table, where):
|
||||
if not isinstance(where, dict) or not where:
|
||||
raise ValueError("delete requires a non-empty where clause")
|
||||
whereClause, params = _build_where(where)
|
||||
query = f"DELETE FROM {_safe_id(table)} WHERE {whereClause} RETURNING *"
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, params)
|
||||
return [dict(record) for record in cursor.fetchall()]
|
||||
|
||||
|
||||
def count(table, where=None):
|
||||
query = f"SELECT COUNT(*) AS count FROM {_safe_id(table)}"
|
||||
params = {}
|
||||
if where:
|
||||
clauses, params = _build_where(where)
|
||||
query += f" WHERE {clauses}"
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, params)
|
||||
return cursor.fetchone()["count"]
|
||||
|
||||
|
||||
def exists(table, where):
|
||||
return count(table, where) > 0
|
||||
|
||||
|
||||
def upsert(table, data, conflict_columns):
|
||||
if not isinstance(data, dict) or not data:
|
||||
raise ValueError("upsert data cannot be empty")
|
||||
if not conflict_columns:
|
||||
raise ValueError("conflict columns are required")
|
||||
columns = list(data.keys())
|
||||
updates = [column for column in columns if column not in conflict_columns]
|
||||
action = "DO NOTHING"
|
||||
if updates:
|
||||
action = "DO UPDATE SET " + ", ".join(
|
||||
f"{_safe_id(column)} = EXCLUDED.{_safe_id(column)}" for column in updates
|
||||
)
|
||||
query = f"""
|
||||
INSERT INTO {_safe_id(table)} ({', '.join(_safe_id(col) for col in columns)})
|
||||
VALUES ({', '.join(f'%({col})s' for col in columns)})
|
||||
ON CONFLICT ({', '.join(_safe_id(col) for col in conflict_columns)})
|
||||
{action}
|
||||
RETURNING *
|
||||
"""
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, data)
|
||||
record = cursor.fetchone()
|
||||
return dict(record) if record else None
|
||||
|
||||
|
||||
def insert_many(table, rows):
|
||||
if not rows:
|
||||
return 0
|
||||
columns = list(rows[0].keys())
|
||||
if any(list(row.keys()) != columns for row in rows):
|
||||
raise ValueError("all rows must use the same columns in the same order")
|
||||
query = f"""
|
||||
INSERT INTO {_safe_id(table)} ({', '.join(_safe_id(col) for col in columns)})
|
||||
VALUES %s
|
||||
"""
|
||||
template = f"({', '.join(f'%({column})s' for column in columns)})"
|
||||
with get_cursor() as cursor:
|
||||
psycopg2.extras.execute_values(
|
||||
cursor, query, rows, template=template, page_size=100
|
||||
)
|
||||
return cursor.rowcount
|
||||
|
||||
|
||||
def execute(query, params=None):
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(query, params or {})
|
||||
if cursor.description:
|
||||
return [dict(record) for record in cursor.fetchall()]
|
||||
return cursor.rowcount
|
||||
|
||||
|
||||
def table_exists(table):
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT EXISTS (
|
||||
SELECT FROM information_schema.tables
|
||||
WHERE table_schema = 'public' AND table_name = %(table)s
|
||||
)
|
||||
""",
|
||||
{"table": table},
|
||||
)
|
||||
return cursor.fetchone()["exists"]
|
||||
|
||||
|
||||
def get_columns(table):
|
||||
with get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT column_name, data_type, is_nullable, column_default
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public' AND table_name = %(table)s
|
||||
ORDER BY ordinal_position
|
||||
""",
|
||||
{"table": table},
|
||||
)
|
||||
return [dict(record) for record in cursor.fetchall()]
|
||||
171
core/registry.py
Normal file
171
core/registry.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
registry.py - Discovery and registration for framework feature modules
|
||||
|
||||
Each package directly under ``modules`` may expose ``register(registry)``.
|
||||
The same registry is loaded by the API, Discord bot, and scheduler so a
|
||||
feature can keep its routes, commands, prompts, and jobs together.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import pkgutil
|
||||
|
||||
|
||||
class FrameworkRegistry:
|
||||
def __init__(self):
|
||||
self.modules = {}
|
||||
self.commands = {}
|
||||
self.route_registrars = []
|
||||
self.job_handlers = {}
|
||||
self._loading_module = None
|
||||
|
||||
def clear(self):
|
||||
self.modules.clear()
|
||||
self.commands.clear()
|
||||
self.route_registrars.clear()
|
||||
self.job_handlers.clear()
|
||||
self._loading_module = None
|
||||
|
||||
def begin_module(self, name, package):
|
||||
if name in self.modules:
|
||||
raise ValueError(f"Duplicate module name: {name}")
|
||||
self.modules[name] = {
|
||||
"name": name,
|
||||
"package": package,
|
||||
"description": "",
|
||||
}
|
||||
self._loading_module = name
|
||||
|
||||
def finish_module(self):
|
||||
self._loading_module = None
|
||||
|
||||
def describe(self, description):
|
||||
module_name = self._require_loading_module()
|
||||
self.modules[module_name]["description"] = description.strip()
|
||||
|
||||
def register_command(
|
||||
self,
|
||||
interaction_type,
|
||||
handler,
|
||||
prompt,
|
||||
validator=None,
|
||||
help_text=None,
|
||||
description="",
|
||||
):
|
||||
module_name = self._require_loading_module()
|
||||
if interaction_type in self.commands:
|
||||
raise ValueError(f"Duplicate command type: {interaction_type}")
|
||||
if not callable(handler):
|
||||
raise TypeError(f"Handler for {interaction_type} must be callable")
|
||||
if validator is not None and not callable(validator):
|
||||
raise TypeError(f"Validator for {interaction_type} must be callable")
|
||||
if not isinstance(prompt, dict) or not prompt.get("system") or not prompt.get(
|
||||
"user_template"
|
||||
):
|
||||
raise ValueError(
|
||||
f"Command {interaction_type} must provide system and user_template prompts"
|
||||
)
|
||||
|
||||
self.commands[interaction_type] = {
|
||||
"module": module_name,
|
||||
"handler": handler,
|
||||
"prompt": prompt,
|
||||
"validator": validator,
|
||||
"help_text": list(help_text or []),
|
||||
"description": description.strip(),
|
||||
}
|
||||
|
||||
def register_routes(self, registrar):
|
||||
module_name = self._require_loading_module()
|
||||
if not callable(registrar):
|
||||
raise TypeError(f"Route registrar for {module_name} must be callable")
|
||||
self.route_registrars.append((module_name, registrar))
|
||||
|
||||
def register_job(self, job_type, handler):
|
||||
module_name = self._require_loading_module()
|
||||
if job_type in self.job_handlers:
|
||||
raise ValueError(f"Duplicate job type: {job_type}")
|
||||
if not callable(handler):
|
||||
raise TypeError(f"Job handler for {job_type} must be callable")
|
||||
self.job_handlers[job_type] = {
|
||||
"module": module_name,
|
||||
"handler": handler,
|
||||
}
|
||||
|
||||
def get_command(self, interaction_type):
|
||||
return self.commands.get(interaction_type)
|
||||
|
||||
def get_job_handler(self, job_type):
|
||||
registration = self.job_handlers.get(job_type)
|
||||
return registration["handler"] if registration else None
|
||||
|
||||
def list_commands(self):
|
||||
return list(self.commands.keys())
|
||||
|
||||
def router_context(self):
|
||||
lines = []
|
||||
for name, command in sorted(self.commands.items()):
|
||||
description = command["description"] or "No description provided"
|
||||
lines.append(f"- {name}: {description}")
|
||||
return "\n".join(lines) if lines else "No modules are available"
|
||||
|
||||
def help_lines(self):
|
||||
lines = []
|
||||
for name, command in sorted(self.commands.items()):
|
||||
if command["help_text"]:
|
||||
lines.extend(f"- {item}" for item in command["help_text"])
|
||||
else:
|
||||
lines.append(f"- {name}: {command['description']}")
|
||||
return lines
|
||||
|
||||
def _require_loading_module(self):
|
||||
if not self._loading_module:
|
||||
raise RuntimeError("Registration must happen inside a module register() call")
|
||||
return self._loading_module
|
||||
|
||||
|
||||
registry = FrameworkRegistry()
|
||||
_loaded = False
|
||||
|
||||
|
||||
def discover_modules(force=False):
|
||||
"""Discover and register feature packages exactly once per process."""
|
||||
global _loaded
|
||||
if _loaded and not force:
|
||||
return registry
|
||||
|
||||
_loaded = False
|
||||
|
||||
package = importlib.import_module("modules")
|
||||
discovered = sorted(
|
||||
item.name
|
||||
for item in pkgutil.iter_modules(package.__path__)
|
||||
if item.ispkg and not item.name.startswith("_")
|
||||
)
|
||||
|
||||
registry.clear()
|
||||
try:
|
||||
for module_name in discovered:
|
||||
qualified_name = f"modules.{module_name}"
|
||||
feature_module = importlib.import_module(qualified_name)
|
||||
register_fn = getattr(feature_module, "register", None)
|
||||
if not callable(register_fn):
|
||||
raise RuntimeError(f"{qualified_name} must expose register(registry)")
|
||||
|
||||
registry.begin_module(module_name, qualified_name)
|
||||
try:
|
||||
register_fn(registry)
|
||||
finally:
|
||||
registry.finish_module()
|
||||
except Exception:
|
||||
registry.clear()
|
||||
raise
|
||||
|
||||
_loaded = True
|
||||
return registry
|
||||
|
||||
|
||||
def reset_registry():
|
||||
"""Reset discovery state for tests."""
|
||||
global _loaded
|
||||
registry.clear()
|
||||
_loaded = False
|
||||
205
core/users.py
Normal file
205
core/users.py
Normal file
@@ -0,0 +1,205 @@
|
||||
import os
|
||||
import uuid
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
import bcrypt
|
||||
from psycopg2.errors import UniqueViolation
|
||||
|
||||
import core.postgres as postgres
|
||||
|
||||
|
||||
MIN_PASSWORD_LENGTH = 8
|
||||
MAX_PASSWORD_BYTES = 72
|
||||
REGISTRATION_FIELDS = {"timezone"}
|
||||
USER_UPDATE_FIELDS = {"timezone"}
|
||||
|
||||
|
||||
def getUser(userUUID):
|
||||
return postgres.select_one("users", {"id": userUUID})
|
||||
|
||||
|
||||
def getUserUUID(username):
|
||||
if not username:
|
||||
return False
|
||||
userRecord = postgres.select_one("users", {"username": username})
|
||||
if userRecord:
|
||||
return userRecord["id"]
|
||||
return False
|
||||
|
||||
|
||||
def getUserFirstName(userUUID):
|
||||
userRecord = getUser(userUUID)
|
||||
if userRecord:
|
||||
return userRecord.get("username")
|
||||
return None
|
||||
|
||||
|
||||
def getUserTimezone(userUUID):
|
||||
userRecord = getUser(userUUID)
|
||||
if userRecord:
|
||||
return userRecord.get("timezone") or "UTC"
|
||||
return None
|
||||
|
||||
|
||||
def isUsernameAvailable(username):
|
||||
if not username:
|
||||
return False
|
||||
return not postgres.exists("users", {"username": username})
|
||||
|
||||
|
||||
def doesUserUUIDExist(userUUID):
|
||||
if not userUUID:
|
||||
return False
|
||||
return postgres.exists("users", {"id": userUUID})
|
||||
|
||||
|
||||
def isValidTimezone(timezoneName):
|
||||
if not isinstance(timezoneName, str) or not timezoneName.strip():
|
||||
return False
|
||||
try:
|
||||
ZoneInfo(timezoneName.strip())
|
||||
return True
|
||||
except (ZoneInfoNotFoundError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def normalizeTimezone(timezoneName):
|
||||
timezoneName = timezoneName.strip() if isinstance(timezoneName, str) else ""
|
||||
if not isValidTimezone(timezoneName):
|
||||
raise ValueError("timezone must be a valid IANA timezone name")
|
||||
return timezoneName
|
||||
|
||||
|
||||
def getDefaultTimezone():
|
||||
return normalizeTimezone(os.getenv("DEFAULT_TIMEZONE", "UTC"))
|
||||
|
||||
|
||||
def validatePassword(password):
|
||||
errors = []
|
||||
if not isinstance(password, str) or not password:
|
||||
return False, ["password"]
|
||||
|
||||
encodedPassword = password.encode("utf-8")
|
||||
if len(password) < MIN_PASSWORD_LENGTH:
|
||||
errors.append(f"password must be at least {MIN_PASSWORD_LENGTH} characters")
|
||||
if len(encodedPassword) > MAX_PASSWORD_BYTES:
|
||||
errors.append(f"password must be at most {MAX_PASSWORD_BYTES} UTF-8 bytes")
|
||||
return not errors, errors
|
||||
|
||||
|
||||
def registerUser(username, password, data=None):
|
||||
if not isinstance(username, str) or not username.strip():
|
||||
raise ValueError("username is required")
|
||||
username = username.strip()
|
||||
if len(username) > 255:
|
||||
raise ValueError("username must be at most 255 characters")
|
||||
|
||||
isValid, errors = validatePassword(password)
|
||||
if not isValid:
|
||||
raise ValueError(f"Invalid password: {', '.join(errors)}")
|
||||
if not isUsernameAvailable(username):
|
||||
return False
|
||||
|
||||
userData = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"username": username,
|
||||
"password_hashed": bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()),
|
||||
}
|
||||
if data is not None:
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("registration data must be an object")
|
||||
for key in REGISTRATION_FIELDS:
|
||||
if key in data:
|
||||
userData[key] = data[key]
|
||||
|
||||
try:
|
||||
createUser(userData)
|
||||
except UniqueViolation:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def updateUser(userUUID, data_dict):
|
||||
if not getUser(userUUID) or not isinstance(data_dict, dict):
|
||||
return False
|
||||
|
||||
updates = {key: data_dict[key] for key in USER_UPDATE_FIELDS if key in data_dict}
|
||||
if "timezone" in updates:
|
||||
updates["timezone"] = normalizeTimezone(updates["timezone"])
|
||||
if not updates:
|
||||
return False
|
||||
|
||||
postgres.update("users", updates, {"id": userUUID})
|
||||
return True
|
||||
|
||||
|
||||
def changePassword(userUUID, new_password):
|
||||
if not getUser(userUUID):
|
||||
return False
|
||||
isValid, errors = validatePassword(new_password)
|
||||
if not isValid:
|
||||
raise ValueError(f"Invalid password: {', '.join(errors)}")
|
||||
|
||||
hashedPassword = bcrypt.hashpw(new_password.encode("utf-8"), bcrypt.gensalt())
|
||||
postgres.update("users", {"password_hashed": hashedPassword}, {"id": userUUID})
|
||||
return True
|
||||
|
||||
|
||||
def deleteUser(userUUID):
|
||||
if not getUser(userUUID):
|
||||
return False
|
||||
postgres.delete("users", {"id": userUUID})
|
||||
return True
|
||||
|
||||
|
||||
def createUser(data_dict):
|
||||
if not isinstance(data_dict, dict):
|
||||
raise ValueError("user data must be an object")
|
||||
|
||||
allowedFields = {"id", "username", "password_hashed", "timezone", "created_at"}
|
||||
userData = {key: value for key, value in data_dict.items() if key in allowedFields}
|
||||
if "timezone" not in userData:
|
||||
userData["timezone"] = getDefaultTimezone()
|
||||
|
||||
isValid, errors = validateUser(userData)
|
||||
if not isValid:
|
||||
raise ValueError(f"Invalid user data: {', '.join(errors)}")
|
||||
|
||||
return postgres.insert("users", userData)
|
||||
|
||||
|
||||
def validateUser(user):
|
||||
errors = []
|
||||
if not isinstance(user, dict):
|
||||
return False, ["user"]
|
||||
|
||||
userUUID = user.get("id")
|
||||
if not userUUID:
|
||||
errors.append("id")
|
||||
else:
|
||||
try:
|
||||
uuid.UUID(str(userUUID))
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
errors.append("id must be a valid UUID")
|
||||
|
||||
username = user.get("username")
|
||||
passwordHash = user.get("password_hashed")
|
||||
if username is None and passwordHash is not None:
|
||||
errors.append("username")
|
||||
elif username is not None and passwordHash is None:
|
||||
errors.append("password_hashed")
|
||||
elif username is not None:
|
||||
if not isinstance(username, str) or not username.strip():
|
||||
errors.append("username")
|
||||
elif len(username) > 255:
|
||||
errors.append("username must be at most 255 characters")
|
||||
if isinstance(passwordHash, memoryview):
|
||||
passwordHash = bytes(passwordHash)
|
||||
if not isinstance(passwordHash, bytes) or not passwordHash:
|
||||
errors.append("password_hashed")
|
||||
|
||||
timezoneName = user.get("timezone", "UTC")
|
||||
if not isValidTimezone(timezoneName):
|
||||
errors.append("timezone must be a valid IANA timezone name")
|
||||
|
||||
return not errors, errors
|
||||
Reference in New Issue
Block a user