Merge pull request #9 from nimh-dsst/add-work-in-doc

Generate work for document

.github/workflows/tests.yml

@@ -34,16 +34,24 @@ jobs:
       - name: Install dependencies
         run: |
           pip install uv
-          uv pip install --system -r <(uv pip compile --all-extras pyproject.toml)
+          pip install .[ci]
 
       - name: Start Docker stack
         run: |
           cp .mockenv .env
           docker compose -f .docker/postgres-compose.yaml up -d
 
-      - name: Run tests
-        run: pytest tests
+      - name: Run tests with coverage
+        run: |
+          coverage run -m pytest -s tests
+          coverage report
+          coverage xml
 
       - name: Stop Docker stack
         if: always()
-        run: docker compose -f compose.yaml -f compose.development.override.yaml down
+        run: docker compose -f .docker/postgres-compose.yaml down
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v3
+        with:
+          file: coverage.xml

.gitignore

@@ -12,4 +12,5 @@ _version.py
 pdfs/*
 test-pdf/*.pdf
 dsst_etl.egg-info/*
-uv.lock
+uv.lock
+.coverage

.mockenv

@@ -4,9 +4,15 @@ POSTGRES_USER=postgres
 POSTGRES_PASSWORD=postgres
 POSTGRES_HOST=localhost
 POSTGRES_PORT=5432
+POSTGRES_DB=pdx
+POSTGRES_DB_UNIT_TEST=pdx
 
 # requests will be rejected without user-agent
 USER_AGENT=
 
 # metapub; get key from https://pubmed.ncbi.nlm.nih.gov for faster requests
 NCBI_API_KEY=
+
+S3_BUCKET_NAME=osm-pdf-uploads
+HOSTNAME=localhost
+USERNAME=quang

README.md

@@ -124,7 +124,7 @@ docker compose -f .docker/postgres-compose.yaml down -v
 pre-commit install
 
 # run the pre-commit hooks on all files
-pre-commit run -all
+   pre-commit run --all-files
 
 # run the tests
 pytest

alembic/versions/4a908d10b459_update_non_null_of_document_in_works_.py

@@ -0,0 +1,40 @@
+"""update non-null of document in works table
+
+Revision ID: 4a908d10b459
+Revises: 360c65a62392
+Create Date: 2024-11-21 16:40:56.966690
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '4a908d10b459'
+down_revision: Union[str, None] = '360c65a62392'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('works', 'initial_document_id',
+               existing_type=sa.INTEGER(),
+               nullable=True)
+    op.alter_column('works', 'primary_document_id',
+               existing_type=sa.INTEGER(),
+               nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('works', 'primary_document_id',
+               existing_type=sa.INTEGER(),
+               nullable=False)
+    op.alter_column('works', 'initial_document_id',
+               existing_type=sa.INTEGER(),
+               nullable=False)
+    # ### end Alembic commands ###

dsst_etl/__init__.py

@@ -2,26 +2,45 @@
 DSST ETL Package
 """
 
+import logging
 import os
 
-from dotenv import load_dotenv
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
-load_dotenv()
+from .config import config
+
+logger = logging.getLogger(__name__)
+
 
 def get_db_url():
     database_url = (
         "postgresql://"
-        f"{os.environ['POSTGRES_USER']}"
-        f":{os.environ['POSTGRES_PASSWORD']}"
-        f"@{os.environ['POSTGRES_HOST']}:"
-        f"{os.environ['POSTGRES_PORT']}/{os.environ['POSTGRES_DB']}"
+        f"{config.POSTGRES_USER}"
+        f":{config.POSTGRES_PASSWORD}"
+        f"@{config.POSTGRES_HOST}:"
+        f"{config.POSTGRES_PORT}/{config.POSTGRES_DB}"
+    )
+    return database_url
+
+
+def get_db_url_test():
+    database_url = (
+        "postgresql://"
+        f"{config.POSTGRES_USER}"
+        f":{config.POSTGRES_PASSWORD}"
+        f"@{config.POSTGRES_HOST}:"
+        f"{config.POSTGRES_PORT}/{config.POSTGRES_DB_UNIT_TEST}"
     )
     return database_url
 
-def get_db_engine():
-    return create_engine(get_db_url())
+
+def get_db_engine(is_test=False):
+    if is_test:
+        return create_engine(get_db_url_test())
+    else:
+        return create_engine(get_db_url())
+
 
 engine = get_db_engine()
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

dsst_etl/_utils.py

@@ -1,12 +1,12 @@
-import os
+from .config import config
 
 
 def get_compute_context_id():
-    return hash(f"{os.environ.get('HOSTNAME')}_{os.environ.get('USERNAME')}")
+    return hash(f"{config.HOSTNAME}_{config.USERNAME}")
 
 
 def get_bucket_name():
-    bucket_name = os.getenv('S3_BUCKET_NAME')
+    bucket_name = config.S3_BUCKET_NAME
     if not bucket_name:
         raise ValueError("S3_BUCKET_NAME environment variable is not set")
-    return bucket_name
+    return bucket_name

dsst_etl/config.py

@@ -0,0 +1,11 @@
+from dotenv import dotenv_values
+
+
+class Config:
+    def __init__(self, config_dict):
+        for key, value in config_dict.items():
+            setattr(self, key, value)
+
+
+# Load the environment variables from the .env file
+config = Config(dotenv_values(".env"))

dsst_etl/db.py

@@ -1,19 +1,25 @@
-import os
+import logging
 
-from dsst_etl import get_db_engine
-from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
+from sqlalchemy_utils import create_database, database_exists
+
+from dsst_etl import get_db_engine
 
 from .models import Base
 
+logger = logging.getLogger(__name__)
 
 
-def get_db_session():
-    engine = get_db_engine()
+def get_db_session(is_test=False):
+    engine = get_db_engine(is_test)
     Session = sessionmaker(bind=engine)
     return Session()
 
 
-def init_db():
-    engine = get_db_engine()
+def init_db(is_test=False):
+    engine = get_db_engine(is_test)
+
+    if not database_exists(engine.url):
+        logger.info("Creating database.....")
+        create_database(engine.url)
     Base.metadata.create_all(engine)

dsst_etl/models.py

@@ -1,8 +1,5 @@
-from datetime import datetime
-
-from sqlalchemy import Column, DateTime, ForeignKey, Integer, LargeBinary, String, Text
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import relationship
+from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text
+from sqlalchemy.orm import declarative_base
 from sqlalchemy.sql import func
 
 Base = declarative_base()
@@ -14,12 +11,8 @@ class Works(Base):
     id = Column(Integer, primary_key=True)
     created_at = Column(DateTime, default=func.now())
     modified_at = Column(DateTime, default=func.now(), onupdate=func.now())
-    initial_document_id = Column(
-        Integer, ForeignKey("documents.id"), nullable=False
-    )
-    primary_document_id = Column(
-        Integer, ForeignKey("documents.id"), nullable=False
-    )
+    initial_document_id = Column(Integer, ForeignKey("documents.id"), nullable=True)
+    primary_document_id = Column(Integer, ForeignKey("documents.id"), nullable=True)
     provenance_id = Column(Integer, ForeignKey("provenance.id"))
 
     # Relationships
@@ -30,7 +23,7 @@ class Works(Base):
 
 class Documents(Base):
     __tablename__ = "documents"
-    
+
     id = Column(Integer, primary_key=True)
     hash_data = Column(String, nullable=False, unique=True)
     created_at = Column(DateTime, default=func.now())
@@ -42,7 +35,6 @@ class Documents(Base):
     # provenance = relationship("Provenance")
 
 
-
 class Provenance(Base):
     __tablename__ = "provenance"
 
@@ -52,4 +44,3 @@ class Provenance(Base):
     compute = Column(Text)
     personnel = Column(Text)
     comment = Column(Text)
-