Migrate to flask_sqlalchemy

apiserver/app.py

@@ -1,4 +1,5 @@
 from flask import Flask
+from dora.store import configure_db_with_app
 from env import load_app_env
 
 load_app_env()
@@ -21,6 +22,7 @@
 app.register_blueprint(deployment_analytics_api)
 app.register_blueprint(integrations_api)
 
+configure_db_with_app(app)
 initialize_database(app)
 
 

apiserver/dora/store/__init__.py

@@ -1,37 +1,22 @@
 from os import getenv
 
-from sqlalchemy import create_engine
-from sqlalchemy.orm import Session, declarative_base
+from flask_sqlalchemy import SQLAlchemy
 
-from dora.utils.log import LOG
+db = SQLAlchemy()
 
-DB_HOST = getenv("DB_HOST")
-DB_PORT = getenv("DB_PORT")
-DB_USER = getenv("DB_USER")
-DB_PASS = getenv("DB_PASS")
-DB_NAME = getenv("DB_NAME")
 
-ENVIRONMENT = getenv("ENVIRONMENT")
+def configure_db_with_app(app):
 
-engine = create_engine(
-    f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}",
-    connect_args={"application_name": f"dora--{ENVIRONMENT}"},
-)
-session = Session(engine)
+    DB_HOST = getenv("DB_HOST")
+    DB_PORT = getenv("DB_PORT")
+    DB_USER = getenv("DB_USER")
+    DB_PASS = getenv("DB_PASS")
+    DB_NAME = getenv("DB_NAME")
+    ENVIRONMENT = getenv("ENVIRONMENT", "local")
 
-Base = declarative_base()
+    connection_uri = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}?application_name=dora--{ENVIRONMENT}"
 
-
-def rollback_on_exc(func):
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except Exception as e:
-            session.rollback()
-            LOG.error(f"Error in {func.__name__} - {str(e)}")
-            raise
-        finally:
-            # session.close()
-            pass
-
-    return wrapper
+    app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
+    app.config["SQLALCHEMY_DATABASE_URI"] = connection_uri
+    app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {"pool_size": 20, "max_overflow": 5}
+    db.init_app(app)

apiserver/dora/store/initialise_db.py

@@ -1,14 +1,13 @@
-from dora.store import rollback_on_exc, session
+from dora.store import db
 from dora.store.models import Organization
 from dora.utils.string import uuid4_str
 from dora.utils.time import time_now
 
 
-@rollback_on_exc
 def initialize_database(app):
     with app.app_context():
         default_org = (
-            session.query(Organization)
+            db.session.query(Organization)
             .filter(Organization.name == "default")
             .one_or_none()
         )
@@ -20,8 +19,8 @@ def initialize_database(app):
             domain="default",
             created_at=time_now(),
         )
-        session.add(default_org)
-        session.commit()
+        db.session.add(default_org)
+        db.session.commit()
 
 
 if __name__ == "__main__":

apiserver/dora/store/models/code/pull_requests.py

@@ -1,49 +1,48 @@
 from datetime import datetime
 
-import pytz
-from sqlalchemy import Column, String, DateTime, ForeignKey, Integer, func
+from sqlalchemy import func
 from sqlalchemy.dialects.postgresql import UUID, JSONB, ARRAY, ENUM
 
-from dora.store import Base
+from dora.store import db
 from dora.store.models.code.enums import (
     PullRequestEventType,
     PullRequestState,
     PullRequestRevertPRMappingActorType,
 )
 
 
-class PullRequest(Base):
+class PullRequest(db.Model):
     __tablename__ = "PullRequest"
 
-    id = Column(UUID(as_uuid=True), primary_key=True)
-    repo_id = Column(UUID(as_uuid=True), ForeignKey("OrgRepo.id"))
-    title = Column(String)
-    url = Column(String)
-    number = Column(String)
-    author = Column(String)
-    state = Column(ENUM(PullRequestState))
-    requested_reviews = Column(ARRAY(String))
-    base_branch = Column(String)
-    head_branch = Column(String)
-    data = Column(JSONB)
-    created_at = Column(DateTime(timezone=True))
-    updated_at = Column(DateTime(timezone=True))
-    state_changed_at = Column(DateTime(timezone=True))
-    first_response_time = Column(Integer)
-    rework_time = Column(Integer)
-    merge_time = Column(Integer)
-    cycle_time = Column(Integer)
-    reviewers = Column(ARRAY(String))
-    meta = Column(JSONB)
-    provider = Column(String)
-    rework_cycles = Column(Integer, default=0)
-    first_commit_to_open = Column(Integer)
-    merge_to_deploy = Column(Integer)
-    lead_time = Column(Integer)
-    merge_commit_sha = Column(String)
-    created_in_db_at = Column(DateTime(timezone=True), server_default=func.now())
-    updated_in_db_at = Column(
-        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
+    id = db.Column(UUID(as_uuid=True), primary_key=True)
+    repo_id = db.Column(UUID(as_uuid=True), db.ForeignKey("OrgRepo.id"))
+    title = db.Column(db.String)
+    url = db.Column(db.String)
+    number = db.Column(db.String)
+    author = db.Column(db.String)
+    state = db.Column(ENUM(PullRequestState))
+    requested_reviews = db.Column(ARRAY(db.String))
+    base_branch = db.Column(db.String)
+    head_branch = db.Column(db.String)
+    data = db.Column(JSONB)
+    created_at = db.Column(db.DateTime(timezone=True))
+    updated_at = db.Column(db.DateTime(timezone=True))
+    state_changed_at = db.Column(db.DateTime(timezone=True))
+    first_response_time = db.Column(db.Integer)
+    rework_time = db.Column(db.Integer)
+    merge_time = db.Column(db.Integer)
+    cycle_time = db.Column(db.Integer)
+    reviewers = db.Column(ARRAY(db.String))
+    meta = db.Column(JSONB)
+    provider = db.Column(db.String)
+    rework_cycles = db.Column(db.Integer, default=0)
+    first_commit_to_open = db.Column(db.Integer)
+    merge_to_deploy = db.Column(db.Integer)
+    lead_time = db.Column(db.Integer)
+    merge_commit_sha = db.Column(db.String)
+    created_in_db_at = db.Column(db.DateTime(timezone=True), server_default=func.now())
+    updated_in_db_at = db.Column(
+        db.DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
     )
 
     def __eq__(self, other):
@@ -80,20 +79,20 @@ def username(self) -> str:
         return self.meta.get("user_profile", {}).get("username", "")
 
 
-class PullRequestEvent(Base):
+class PullRequestEvent(db.Model):
     __tablename__ = "PullRequestEvent"
 
-    id = Column(UUID(as_uuid=True), primary_key=True)
-    pull_request_id = Column(UUID(as_uuid=True), ForeignKey("PullRequest.id"))
-    type = Column(ENUM(PullRequestEventType))
-    data = Column(JSONB)
-    created_at = Column(DateTime(timezone=True))
-    idempotency_key = Column(String)
-    org_repo_id = Column(UUID(as_uuid=True), ForeignKey("OrgRepo.id"))
-    actor_username = Column(String)
-    created_in_db_at = Column(DateTime(timezone=True), server_default=func.now())
-    updated_in_db_at = Column(
-        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
+    id = db.Column(UUID(as_uuid=True), primary_key=True)
+    pull_request_id = db.Column(UUID(as_uuid=True), db.ForeignKey("PullRequest.id"))
+    type = db.Column(ENUM(PullRequestEventType))
+    data = db.Column(JSONB)
+    created_at = db.Column(db.DateTime(timezone=True))
+    idempotency_key = db.Column(db.String)
+    org_repo_id = db.Column(UUID(as_uuid=True), db.ForeignKey("OrgRepo.id"))
+    actor_username = db.Column(db.String)
+    created_in_db_at = db.Column(db.DateTime(timezone=True), server_default=func.now())
+    updated_in_db_at = db.Column(
+        db.DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
     )
 
     @property
@@ -107,42 +106,42 @@ def state(self):
         return ""
 
 
-class PullRequestCommit(Base):
+class PullRequestCommit(db.Model):
     __tablename__ = "PullRequestCommit"
 
-    hash = Column(String, primary_key=True)
-    pull_request_id = Column(UUID(as_uuid=True), ForeignKey("PullRequest.id"))
-    message = Column(String)
-    url = Column(String)
-    data = Column(JSONB)
-    author = Column(String)
-    created_at = Column(DateTime(timezone=True))
-    org_repo_id = Column(UUID(as_uuid=True), ForeignKey("OrgRepo.id"))
-    created_in_db_at = Column(DateTime(timezone=True), server_default=func.now())
-    updated_in_db_at = Column(
-        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
+    hash = db.Column(db.String, primary_key=True)
+    pull_request_id = db.Column(UUID(as_uuid=True), db.ForeignKey("PullRequest.id"))
+    message = db.Column(db.String)
+    url = db.Column(db.String)
+    data = db.Column(JSONB)
+    author = db.Column(db.String)
+    created_at = db.Column(db.DateTime(timezone=True))
+    org_repo_id = db.Column(UUID(as_uuid=True), db.ForeignKey("OrgRepo.id"))
+    created_in_db_at = db.Column(db.DateTime(timezone=True), server_default=func.now())
+    updated_in_db_at = db.Column(
+        db.DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
     )
 
 
-class PullRequestRevertPRMapping(Base):
+class PullRequestRevertPRMapping(db.Model):
     __tablename__ = "PullRequestRevertPRMapping"
 
-    pr_id = Column(
+    pr_id = db.Column(
         UUID(as_uuid=True),
-        ForeignKey("PullRequest.id"),
+        db.ForeignKey("PullRequest.id"),
         primary_key=True,
         nullable=False,
     )
-    actor_type = Column(
+    actor_type = db.Column(
         ENUM(PullRequestRevertPRMappingActorType), primary_key=True, nullable=False
     )
-    actor = Column(UUID(as_uuid=True), ForeignKey("Users.id"))
-    reverted_pr = Column(
-        UUID(as_uuid=True), ForeignKey("PullRequest.id"), nullable=False
+    actor = db.Column(UUID(as_uuid=True), db.ForeignKey("Users.id"))
+    reverted_pr = db.Column(
+        UUID(as_uuid=True), db.ForeignKey("PullRequest.id"), nullable=False
     )
-    created_at = Column(DateTime(timezone=True), server_default=func.now())
-    updated_at = Column(
-        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
+    created_at = db.Column(db.DateTime(timezone=True), server_default=func.now())
+    updated_at = db.Column(
+        db.DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
     )
 
     def __hash__(self):