Switch to postgresql backend. Add azure entrypoint.sh

Author: grvlbit

entrypoint.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+set -e
+python3 -m pip install --upgrade pip
+python3 -m pip install -r requirements.txt
+python3 ssh_keyservice_api/seed_database.py
+fastapi run ssh_keyservice_api/main.py

poetry.lock

@@ -11,18 +11,24 @@ python = "^3.11"
 fastapi = "^0.115.8"
 cryptography = "^44.0.1"
 fastapi-azure-auth = "^5.0.1"
-valkey = "^6.1.0"
 pydantic = "^2.10.6"
 pydantic-settings = "^2.7.1"
 cachetools = "^5.5.1"
 azure-identity = "^1.20.0"
 azure-keyvault-secrets = "^4.9.0"
-
+sqlmodel = "^0.0.22"
+psycopg2 = "^2.9.10"
+azure-monitor-opentelemetry = "^1.6.5"
+uvicorn = "^0.34.0"
 
 [tool.poetry.group.dev.dependencies]
-uvicorn = "^0.34.0"
 python-dotenv = "^1.0.1"
 
+[tool.poetry-auto-export]
+output = "requirements.txt"
+without_hashes = true
+without = ["dev"]
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

pyproject.toml

@@ -0,0 +1,74 @@
+# poetry.lock hash: 6aac512a4174ad31a98adbfa760fdd2b2372b242
+# This file is generated by poetry-auto-export
+# The SHA1 hash of the poetry.lock file is printed above
+annotated-types==0.7.0 ; python_version >= "3.11" and python_version < "4.0"
+anyio==4.8.0 ; python_version >= "3.11" and python_version < "4.0"
+asgiref==3.8.1 ; python_version >= "3.11" and python_version < "4.0"
+azure-core-tracing-opentelemetry==1.0.0b11 ; python_version >= "3.11" and python_version < "4.0"
+azure-core==1.32.0 ; python_version >= "3.11" and python_version < "4.0"
+azure-identity==1.20.0 ; python_version >= "3.11" and python_version < "4.0"
+azure-keyvault-secrets==4.9.0 ; python_version >= "3.11" and python_version < "4.0"
+azure-monitor-opentelemetry-exporter==1.0.0b33 ; python_version >= "3.11" and python_version < "4.0"
+azure-monitor-opentelemetry==1.6.5 ; python_version >= "3.11" and python_version < "4.0"
+cachetools==5.5.1 ; python_version >= "3.11" and python_version < "4.0"
+certifi==2025.1.31 ; python_version >= "3.11" and python_version < "4.0"
+cffi==1.17.1 ; python_version >= "3.11" and python_version < "4.0" and platform_python_implementation != "PyPy"
+charset-normalizer==3.4.1 ; python_version >= "3.11" and python_version < "4.0"
+click==8.1.8 ; python_version >= "3.11" and python_version < "4.0"
+colorama==0.4.6 ; python_version >= "3.11" and python_version < "4.0" and platform_system == "Windows"
+cryptography==44.0.1 ; python_version >= "3.11" and python_version < "4.0"
+deprecated==1.2.18 ; python_version >= "3.11" and python_version < "4.0"
+fastapi-azure-auth==5.0.1 ; python_version >= "3.11" and python_version < "4.0"
+fastapi==0.115.8 ; python_version >= "3.11" and python_version < "4.0"
+fixedint==0.1.6 ; python_version >= "3.11" and python_version < "4.0"
+greenlet==3.1.1 ; python_version < "3.14" and (platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "x86_64" or platform_machine == "amd64" or platform_machine == "AMD64" or platform_machine == "win32" or platform_machine == "WIN32") and python_version >= "3.11"
+h11==0.14.0 ; python_version >= "3.11" and python_version < "4.0"
+httpcore==1.0.7 ; python_version >= "3.11" and python_version < "4.0"
+httpx==0.28.1 ; python_version >= "3.11" and python_version < "4.0"
+idna==3.10 ; python_version >= "3.11" and python_version < "4.0"
+importlib-metadata==8.5.0 ; python_version >= "3.11" and python_version < "4.0"
+isodate==0.7.2 ; python_version >= "3.11" and python_version < "4.0"
+msal-extensions==1.2.0 ; python_version >= "3.11" and python_version < "4.0"
+msal==1.31.1 ; python_version >= "3.11" and python_version < "4.0"
+msrest==0.7.1 ; python_version >= "3.11" and python_version < "4.0"
+oauthlib==3.2.2 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-api==1.30.0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-asgi==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-dbapi==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-django==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-fastapi==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-flask==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-psycopg2==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-requests==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-urllib3==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-urllib==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation-wsgi==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-instrumentation==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-resource-detector-azure==0.1.5 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-sdk==1.30.0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-semantic-conventions==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+opentelemetry-util-http==0.51b0 ; python_version >= "3.11" and python_version < "4.0"
+packaging==24.2 ; python_version >= "3.11" and python_version < "4.0"
+portalocker==2.10.1 ; python_version >= "3.11" and python_version < "4.0"
+psutil==5.9.8 ; python_version >= "3.11" and python_version < "4.0"
+psycopg2==2.9.10 ; python_version >= "3.11" and python_version < "4.0"
+pycparser==2.22 ; python_version >= "3.11" and python_version < "4.0" and platform_python_implementation != "PyPy"
+pydantic-core==2.27.2 ; python_version >= "3.11" and python_version < "4.0"
+pydantic-settings==2.7.1 ; python_version >= "3.11" and python_version < "4.0"
+pydantic==2.10.6 ; python_version >= "3.11" and python_version < "4.0"
+pyjwt==2.10.1 ; python_version >= "3.11" and python_version < "4.0"
+pyjwt[crypto]==2.10.1 ; python_version >= "3.11" and python_version < "4.0"
+python-dotenv==1.0.1 ; python_version >= "3.11" and python_version < "4.0"
+pywin32==308 ; python_version >= "3.11" and python_version < "4.0" and platform_system == "Windows"
+requests-oauthlib==2.0.0 ; python_version >= "3.11" and python_version < "4.0"
+requests==2.32.3 ; python_version >= "3.11" and python_version < "4.0"
+six==1.17.0 ; python_version >= "3.11" and python_version < "4.0"
+sniffio==1.3.1 ; python_version >= "3.11" and python_version < "4.0"
+sqlalchemy==2.0.38 ; python_version >= "3.11" and python_version < "4.0"
+sqlmodel==0.0.22 ; python_version >= "3.11" and python_version < "4.0"
+starlette==0.45.3 ; python_version >= "3.11" and python_version < "4.0"
+typing-extensions==4.12.2 ; python_version >= "3.11" and python_version < "4.0"
+urllib3==2.3.0 ; python_version >= "3.11" and python_version < "4.0"
+uvicorn==0.34.0 ; python_version >= "3.11" and python_version < "4.0"
+wrapt==1.17.2 ; python_version >= "3.11" and python_version < "4.0"
+zipp==3.21.0 ; python_version >= "3.11" and python_version < "4.0"

requirements.txt

@@ -6,8 +6,6 @@
 
 import uvicorn
 
-import valkey as redis
-
 from fastapi import FastAPI, Security, Depends, HTTPException
 from fastapi.responses import PlainTextResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -19,12 +17,22 @@
 from typing import AsyncGenerator
 from datetime import datetime
 
+from sqlmodel import Session, select
+
 from cachetools import TTLCache
 
-from models import UserModel, SSHKeyPutRequest, SSHKeyDeleteRequest
+from models import UserModel, SSHKeyPutRequest, SSHKeyDeleteRequest, engine, SSHKey
 
 from dotenv import load_dotenv
 
+from azure.monitor.opentelemetry import configure_azure_monitor
+
+# Setup logger and Azure Monitor:
+logger = logging.getLogger("ssh_keyservice_api")
+logger.setLevel(logging.INFO)
+if os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING"):
+    configure_azure_monitor()
+
 #from azure.identity import DefaultAzureCredential
 #from azure.keyvault.secrets import SecretClient
 # Configure Azure Key Vault
@@ -56,9 +64,13 @@ def get_secret(secret_name: str) -> str:
 DB_HOST = get_secret("DB_HOST")
 SCOPE = { f'api://{APP_CLIENT_ID}/user.read.profile' : 'user.read.profile' }
 
+# Dependency to get the database session
+def get_db_session():
+    with Session(engine) as session:
+        yield session
+
 # Cache API keys for 10 minutes (600 seconds)
 api_key_cache = TTLCache(maxsize=1, ttl=600)
-
 api_key_header_auth = APIKeyHeader(name="x-api-key", auto_error=True)
 
 def get_api_keys():
@@ -72,13 +84,6 @@ async def api_key_auth(api_key_header: str = Security(api_key_header_auth)):
     if not any(secrets.compare_digest(api_key_header, key.strip()) for key in valid_api_keys):
         raise HTTPException(status_code=401, detail="Invalid API Key")
 
-# Logging setup
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("SSHKeyAPI")
-
-# Redis connection
-redis_client = redis.Redis(host=DB_HOST, port=6379, decode_responses=True)
-
 @asynccontextmanager
 async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     """
@@ -110,46 +115,67 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     scopes=SCOPE,
 )
 
-def generate_user_key(email: str) -> str:
-    return f"user:{hashlib.sha256(email.encode()).hexdigest()}"
+def generate_user_hash(email: str) -> str:
+    return f"{hashlib.sha256(email.encode()).hexdigest()}"
 
 @app.get("/api/v1/users/me", dependencies=[Security(azure_scheme)])
-async def get_user_info(user: User = Depends(azure_scheme)) -> UserModel:
+async def get_user_info(user: User = Depends(azure_scheme), session: Session = Depends(get_db_session)) -> UserModel:
     """Retrieve SSH keys for authenticated user."""
     email = user.claims.get("preferred_username") or user.claims.get("email")
-    user_key = generate_user_key(email)
-    stored_keys = redis_client.hgetall(f"{user_key}:keys")
-    # Parse stored data into a structured format
-    ssh_keys = {
-        key: {"comment": value.split("|")[0], "timestamp": value.split("|")[1]} 
-        for key, value in stored_keys.items()
-    }
+    user_hash = generate_user_hash(email)
+
+    # Fetch SSH keys from the database
+    results = session.exec(select(SSHKey).filter(SSHKey.user_hash == user_hash)).all()
+
+    ssh_keys = {}
+    for key in results:
+        ssh_keys.update({key.ssh_key: {"comment": key.comment, "timestamp": key.timestamp}})
+
     return {"email": email, "ssh_keys": ssh_keys}
 
 @app.put("/api/v1/users/me/keys", dependencies=[Security(azure_scheme)])
-async def add_ssh_key(request: SSHKeyPutRequest, user: User = Depends(azure_scheme)) -> dict[str, str]:
+async def add_ssh_key( request: SSHKeyPutRequest, user: User = Depends(azure_scheme), session: Session = Depends(get_db_session)) -> dict[str, str]:
     """Add an SSH key for the authenticated user."""
     email = user.claims.get("preferred_username") or user.claims.get("email")
-    user_key = generate_user_key(email)
+    user_hash = generate_user_hash(email)
     timestamp = datetime.utcnow().isoformat()
-    redis_client.hset(f"{user_key}:keys", request.ssh_key, f"{request.comment}|{timestamp}")
+
+    ssh_key = SSHKey()
+    ssh_key.ssh_key = request.ssh_key
+    ssh_key.user_hash = user_hash
+    ssh_key.comment = request.comment
+    ssh_key.timestamp = timestamp
+    session.add(ssh_key)
+    session.commit()
+    session.refresh(ssh_key)
+
     return {"message": "SSH key added."}
 
 @app.delete("/api/v1/users/me/keys", dependencies=[Security(azure_scheme)])
-async def delete_ssh_key(request: SSHKeyDeleteRequest, user: User = Depends(azure_scheme)) -> dict[str, str]:
+async def delete_ssh_key(request: SSHKeyDeleteRequest, user: User = Depends(azure_scheme), session: Session = Depends(get_db_session)) -> dict[str, str]:
     """Delete a specific SSH key for the authenticated user."""
     email = user.claims.get("preferred_username") or user.claims.get("email")
-    user_key = generate_user_key(email)
-    if not redis_client.hexists(f"{user_key}:keys", request.ssh_key):
+    user_hash = generate_user_hash(email)
+
+    ssh_key = session.exec(select(SSHKey).filter(SSHKey.user_hash == user_hash, SSHKey.ssh_key == request.ssh_key)).first()
+    if not ssh_key:
         raise HTTPException(status_code=404, detail="SSH key not found.")
-    redis_client.hdel(f"{user_key}:keys", request.ssh_key)
+    session.delete(ssh_key)
+    session.commit()
+
     return {"message": "SSH key deleted."}
 
 @app.get("/api/v1/users/{email}/keys", response_class=PlainTextResponse, dependencies=[Security(api_key_auth)])
-async def get_ssh_keys_by_mail(email: str) -> str:
+async def get_ssh_keys_by_mail(email: str, session: Session = Depends(get_db_session)) -> str:
     """Get registered SSH keys for a given email."""
-    user_key = generate_user_key(email)
-    ssh_keys = redis_client.hgetall(f"{user_key}:keys")
+    user_hash = generate_user_hash(email)
+
+    results = session.exec(select(SSHKey).filter(SSHKey.user_hash == user_hash)).all()
+
+    ssh_keys = {}
+    for key in results:
+        ssh_keys.update({key.ssh_key: {"comment": key.comment, "timestamp": key.timestamp}})
+
     return "\n".join(ssh_keys.keys()) if ssh_keys else ""
 
 if __name__ == '__main__':

ssh_keyservice_api/main.py

@@ -1,12 +1,64 @@
+import logging
+import os
+import typing
+from urllib.parse import quote_plus
+
+from dotenv import load_dotenv
+from sqlmodel import Field, SQLModel, create_engine
+
 from pydantic import BaseModel
 
+logger = logging.getLogger("app")
+logger.setLevel(logging.INFO)
+
+sql_url = ""
+if os.getenv("WEBSITE_HOSTNAME"):
+    logger.info("Connecting to Azure PostgreSQL Flexible server based on AZURE_POSTGRESQL_CONNECTIONSTRING...")
+    env_connection_string = os.getenv("AZURE_POSTGRESQL_CONNECTIONSTRING")
+    if env_connection_string is None:
+        logger.info("Missing environment variable AZURE_POSTGRESQL_CONNECTIONSTRING")
+    else:
+        # Parse the connection string
+        details = dict(item.split('=') for item in env_connection_string.split())
+
+        # Properly format the URL for SQLAlchemy
+        sql_url = (
+            f"postgresql://{quote_plus(details['user'])}:{quote_plus(details['password'])}"
+            f"@{details['host']}:{details['port']}/{details['dbname']}?sslmode={details['sslmode']}"
+        )
+
+else:
+    logger.info("Connecting to local PostgreSQL server based on .env file...")
+    load_dotenv()
+    POSTGRES_USERNAME = os.environ.get("DBUSER")
+    POSTGRES_PASSWORD = os.environ.get("DBPASS")
+    POSTGRES_HOST = os.environ.get("DBHOST")
+    POSTGRES_DATABASE = os.environ.get("DBNAME")
+    POSTGRES_PORT = os.environ.get("DBPORT", 5432)
+
+    sql_url = f"postgresql://{POSTGRES_USERNAME}:{POSTGRES_PASSWORD}@{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DATABASE}"
+
+engine = create_engine(sql_url)
+
+def create_db_and_tables():
+    return SQLModel.metadata.create_all(engine)
+
 class UserModel(BaseModel):
     email: str
-    ssh_keys: dict[str, dict[str, str]]
+    ssh_keys: dict[str, dict[str, str]] = {}
 
 class SSHKeyPutRequest(BaseModel):
     ssh_key: str
     comment: str
 
 class SSHKeyDeleteRequest(BaseModel):
     ssh_key: str
+
+class SSHKey(SQLModel, table=True):
+    ssh_key: typing.Optional[str] = Field(default=None, primary_key=True)
+    user_hash: str = Field(max_length=250)
+    comment: str = Field(max_length=50)
+    timestamp: str = Field(max_length=50)
+
+    def __str__(self):
+        return f"{self.ssh_key}"

ssh_keyservice_api/models.py

@@ -0,0 +1,11 @@
+from sqlmodel import SQLModel
+
+from models import SSHKey, create_db_and_tables, engine
+
+def drop_all():
+    # Explicitly remove these tables first to avoid cascade errors
+    SQLModel.metadata.remove(SSHKey.__table__)
+    SQLModel.metadata.drop_all(engine)
+
+if __name__ == "__main__":
+    create_db_and_tables()