Write owner_id on new dataset creation

app/authentication/token.py

@@ -4,16 +4,17 @@
 from fastapi.logger import logger
 from fastapi.security import OAuth2PasswordBearer
 from httpx import Response
-from ..routes import dataset_dependency
 
-from ..utils.rw_api import who_am_i
+from ..routes import dataset_dependency
 from ..settings.globals import PROTECTED_QUERY_DATASETS
+from ..utils.rw_api import who_am_i
 
 # token dependency where we immediately cause an exception if there is no auth token
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/token")
 # token dependency where we don't cause exception if there is no auth token
 oauth2_scheme_no_auto = OAuth2PasswordBearer(tokenUrl="/token", auto_error=False)
 
+
 async def is_service_account(token: str = Depends(oauth2_scheme)) -> bool:
     """Calls GFW API to authorize user.
 
@@ -40,27 +41,44 @@ async def is_admin(token: str = Depends(oauth2_scheme)) -> bool:
 
     return await is_app_admin(token, "gfw", "Unauthorized")
 
-async def is_gfwpro_admin_for_query(dataset: str = Depends(dataset_dependency),
-                                    token: str | None = Depends(oauth2_scheme_no_auto)) -> bool:
+
+async def rw_user_id(token: str = Depends(oauth2_scheme)) -> str:
+    """Gets user ID from token."""
+
+    return await who_am_i(token).json()["id"]
+
+
+async def is_gfwpro_admin_for_query(
+    dataset: str = Depends(dataset_dependency),
+    token: str | None = Depends(oauth2_scheme_no_auto),
+) -> bool:
     """If the dataset is protected dataset, calls GFW API to authorize user by
-    requiring the user must be ADMIN for gfw-pro app.  If the dataset is not
-    protected, just returns True without any required token or authorization.
+    requiring the user must be ADMIN for gfw-pro app.
+
+    If the dataset is not protected, just returns True without any
+    required token or authorization.
     """
-    
+
     if dataset in PROTECTED_QUERY_DATASETS:
-        if token == None:
-            raise HTTPException(status_code=401, detail="Unauthorized query on a restricted dataset")
+        if token is None:
+            raise HTTPException(
+                status_code=401, detail="Unauthorized query on a restricted dataset"
+            )
         else:
-            return await is_app_admin(cast(str, token), "gfw-pro",
-                               error_str="Unauthorized query on a restricted dataset")
+            return await is_app_admin(
+                cast(str, token),
+                "gfw-pro",
+                error_str="Unauthorized query on a restricted dataset",
+            )
 
     return True
 
+
 async def is_app_admin(token: str, app: str, error_str: str) -> bool:
     """Calls GFW API to authorize user.
 
-    User must be an ADMIN for the specified app, else it will throw
-    an exception with the specified error string.
+    User must be an ADMIN for the specified app, else it will throw an
+    exception with the specified error string.
     """
 
     response: Response = await who_am_i(token)

app/models/orm/datasets.py

@@ -5,5 +5,5 @@ class Dataset(Base):
     __tablename__ = "datasets"
     dataset = db.Column(db.String, primary_key=True)
     is_downloadable = db.Column(db.Boolean, nullable=False, default=True)
-    owner_id = db.Column(db.UUID, nullable=True, default=None)
+    owner_id = db.Column(db.String, nullable=True, default=None)
     # metadata = db.Column(db.JSONB, default=dict())

app/models/orm/migrations/versions/d767b6dd2c4c_.py

@@ -6,7 +6,6 @@
 """
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
 revision = "d767b6dd2c4c"
@@ -17,7 +16,7 @@
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
-    op.add_column("datasets", sa.Column("owner_id", postgresql.UUID(), nullable=True))
+    op.add_column("datasets", sa.Column("owner_id", sa.String(), nullable=True))
     # ### end Alembic commands ###
 
 

app/routes/datasets/dataset.py

@@ -8,7 +8,7 @@
 from sqlalchemy.schema import CreateSchema, DropSchema
 
 from ...application import db
-from ...authentication.token import is_admin
+from ...authentication.token import is_admin, rw_user_id
 from ...crud import datasets, versions
 from ...errors import RecordAlreadyExistsError, RecordNotFoundError
 from ...models.orm.datasets import Dataset as ORMDataset
@@ -53,11 +53,13 @@ async def create_dataset(
     dataset: str = Depends(dataset_dependency),
     request: DatasetCreateIn,
     is_authorized: bool = Depends(is_admin),
+    owner_id: str = Depends(rw_user_id),
     response: Response,
 ) -> DatasetResponse:
     """Create or update a dataset."""
 
     input_data: Dict = request.dict(exclude_none=True, by_alias=True)
+    input_data["owner_id"] = owner_id
 
     try:
         new_dataset: ORMDataset = await datasets.create_dataset(dataset, **input_data)

tests/__init__.py

@@ -313,6 +313,10 @@ async def get_api_key_mocked() -> Tuple[Optional[str], Optional[str]]:
     return str(uuid.uuid4()), "localhost"
 
 
+async def get_rw_user_id() -> str:
+    return "5874bfcca049b7a56ad42771"  # pragma: allowlist secret
+
+
 def setup_clients(ec2_client, iam_client):
     """
     Do prerequisite setup

tests/conftest.py

@@ -16,7 +16,7 @@
 from httpx import AsyncClient
 
 from app.authentication.api_keys import get_api_key
-from app.authentication.token import is_admin, is_service_account
+from app.authentication.token import is_admin, is_service_account, rw_user_id
 from app.settings.globals import (
     AURORA_JOB_QUEUE,
     AURORA_JOB_QUEUE_FAST,
@@ -52,6 +52,7 @@
     AWSMock,
     MemoryServer,
     get_api_key_mocked,
+    get_rw_user_id,
     is_admin_mocked,
     is_service_account_mocked,
     setup_clients,
@@ -251,6 +252,7 @@ async def async_client():
     app.dependency_overrides[is_admin] = is_admin_mocked
     app.dependency_overrides[is_service_account] = is_service_account_mocked
     app.dependency_overrides[get_api_key] = get_api_key_mocked
+    app.dependency_overrides[rw_user_id] = get_rw_user_id
 
     async with AsyncClient(app=app, base_url="http://test", trust_env=False) as client:
         yield client