database

Watts-Lab · Jul 12, 2024 · 6140e0b · 6140e0b
1 parent 9eaf88b
commit 6140e0b
Show file tree

Hide file tree

Showing 7 changed files with 151 additions and 60 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -28,5 +28,13 @@
     "editor.defaultFormatter": "ms-azuretools.vscode-docker"
   },
   "python.autoComplete.extraPaths": ["./server"],
-  "python.analysis.extraPaths": ["./server"]
+  "python.analysis.extraPaths": ["./server"],
+  "files.associations": {
+    ".env*": "dotenv"
+  },
+  "github.copilot.enable": {
+    "*": true,
+    // ...
+    "dotenv": false
+  }
 }
diff --git a/server/api.py b/server/api.py
@@ -33,6 +33,11 @@
 jwt = JWTManager(app)
 
 
+@app.before_first_request
+def my_func():
+    print("before first request")
+
+
 api.add_resource(GetFeatures, "/api/features")
 api.add_resource(
     RunAssistant, "/api/run_assistant", resource_class_kwargs={"socketio": socketio}

diff --git a/server/controllers/email_client.py b/server/controllers/email_client.py
@@ -0,0 +1,76 @@
+import os
+import boto3
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+from email.mime.application import MIMEApplication
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def send_ses_email(
+    subject, body_text, body_html, sender_email, recipient_emails, attachment=None
+):
+    # Create a multipart/mixed parent container
+    msg = MIMEMultipart("mixed")
+    msg["Subject"] = subject
+    msg["From"] = sender_email
+    msg["To"] = ", ".join(recipient_emails)
+
+    # Add body to email
+    msg_body = MIMEMultipart("alternative")
+    textpart = MIMEText(body_text.encode("utf-8"), "plain", "utf-8")
+    htmlpart = MIMEText(body_html.encode("utf-8"), "html", "utf-8")
+
+    msg_body.attach(textpart)
+    msg_body.attach(htmlpart)
+    msg.attach(msg_body)
+
+    # Attachment
+    if attachment:
+        with open(attachment, "rb") as f:
+            part = MIMEApplication(f.read())
+            part.add_header(
+                "Content-Disposition",
+                "attachment",
+                filename=os.path.basename(attachment),
+            )
+            msg.attach(part)
+
+    # Connect to AWS SES
+    client = boto3.client(
+        "ses",
+        aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
+        aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
+        region_name=os.getenv("AWS_REGION"),
+    )
+
+    # Try to send the email.
+    try:
+        response = client.send_raw_email(
+            Source=sender_email,
+            Destinations=recipient_emails,
+            RawMessage={"Data": msg.as_string()},
+        )
+    except Exception as e:
+        print(e)
+        return False
+    return True
+
+
+if __name__ == "__main__":
+
+    subject = "Your magic link to log in to Atlas"
+    body_text = "Content of your email."
+    body_html = """<html>
+    <head></head>
+    <body>
+    <h1>Welcome to Atlas!</h1>
+    <p>Click <a href='https://atlas.scaledhumanity.org'>here</a> to log in</p>
+    </body>
+    </html>"""
+    sender_email = "[email protected]"
+    recipient_emails = ["[email protected]"]
+
+    # Send the email
+    send_ses_email(subject, body_text, body_html, sender_email, recipient_emails)
diff --git a/server/db/database.py b/server/db/database.py
@@ -0,0 +1,35 @@
+import asyncio
+import os
+from dotenv import load_dotenv
+from beanie import init_beanie
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from models.users import User
+
+
+load_dotenv()
+
+
+# Call this from within your event loop to get beanie setup.
+async def init_db():
+    # Create Motor client
+    client = AsyncIOMotorClient(os.getenv("DB_URI"))
+
+    # Init beanie with the Product document class
+    await init_beanie(database=client.atlas_main, document_models=[User])
+
+
+async def main():
+    await init_db()
+    us = User(
+        email="[email protected]",
+        magic_link="1234",
+        magic_link_expiration="2021-01-01",
+        number_of_tokens=5,
+    )
+
+    await us.create()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/server/db/db.py b/server/db/db.py
diff --git a/server/db/models/users.py b/server/db/models/users.py
@@ -0,0 +1,16 @@
+from datetime import datetime
+from typing import Optional
+from beanie import Document
+from beanie.operators import *
+
+
+class User(Document):
+    email: str
+    magic_link: str
+    magic_link_expiration: datetime
+    number_of_tokens: Optional[int]
+    created_at: datetime = datetime.now()
+    updated_at: datetime = datetime.now()
+
+    class Settings:
+        name = "users"
diff --git a/server/requirements.txt b/server/requirements.txt
@@ -1,8 +1,11 @@
 aniso8601==9.0.1
 annotated-types==0.7.0
 anyio==4.4.0
+beanie==1.26.0
 bidict==0.23.1
 blinker==1.8.2
+boto3==1.34.143
+botocore==1.34.143
 certifi==2024.7.4
 click==8.1.7
 distro==1.9.0
@@ -18,20 +21,27 @@ httpx==0.27.0
 idna==3.7
 itsdangerous==2.2.0
 Jinja2==3.1.4
+jmespath==1.0.1
+lazy-model==0.2.0
 MarkupSafe==2.1.5
+motor==3.5.1
 openai==1.35.6
 pydantic==2.7.4
 pydantic_core==2.18.4
 PyJWT==2.8.0
 pymongo==4.8.0
+python-dateutil==2.9.0.post0
 python-dotenv==1.0.1
 python-engineio==4.9.1
 python-socketio==5.11.3
 pytz==2024.1
+s3transfer==0.10.2
 simple-websocket==1.0.0
 six==1.16.0
 sniffio==1.3.1
+toml==0.10.2
 tqdm==4.66.4
 typing_extensions==4.12.2
+urllib3==2.2.2
 Werkzeug==3.0.3
 wsproto==1.2.0