Document source

migrations/versions/a68c6bb2972c_add_document_source.py

@@ -0,0 +1,24 @@
+"""add_document_source
+
+Revision ID: a68c6bb2972c
+Revises: d63e9f7661dd
+Create Date: 2023-07-26 22:56:57.273888
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = 'a68c6bb2972c'
+down_revision = 'd63e9f7661dd'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column('file', sa.Column('source', sa.String(), nullable=True))
+
+
+def downgrade():
+    op.drop_column('file', 'source')

print_service/models/__init__.py

@@ -42,6 +42,7 @@ class File(Model):
         DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow
     )
     number_of_pages: Mapped[int] = Column(Integer)
+    source: Mapped[str] = Column(String, default='unknown', nullable=False)
 
     owner: Mapped[UnionMember] = relationship('UnionMember', back_populates='files')
     print_facts: Mapped[list[PrintFact]] = relationship('PrintFact', back_populates='file')
@@ -88,5 +89,4 @@ class PrintFact(Model):
 
     owner: Mapped[UnionMember] = relationship('UnionMember', back_populates='print_facts')
     file: Mapped[File] = relationship('File', back_populates='print_facts')
-
     sheets_used: Mapped[int] = Column(Integer)

print_service/routes/file.py

@@ -71,6 +71,7 @@ class SendInput(BaseModel):
         description='Название файла',
         example='filename.pdf',
     )
+    source: str | None = None
     options: PrintOptions = PrintOptions()
 
 
@@ -128,7 +129,7 @@ async def send(inp: SendInput, settings: Settings = Depends(get_settings)):
     except RuntimeError:
         raise PINGenerateError()
     filename = generate_filename(inp.filename)
-    file_model = FileModel(pin=pin, file=filename)
+    file_model = FileModel(pin=pin, file=filename, source=inp.source)
     file_model.owner = user
     file_model.option_copies = inp.options.copies
     file_model.option_pages = inp.options.pages

tests/test_routes/test_file.py

@@ -20,13 +20,27 @@ def test_post_success(union_member_user, client, dbsession):
         "surname": union_member_user['surname'],
         "number": union_member_user['union_number'],
         "filename": "filename.pdf",
+        "source": "webapp",
         "options": {"pages": "", "copies": 1, "two_sided": False},
     }
     res = client.post(url, data=json.dumps(body))
     assert res.status_code == status.HTTP_200_OK
     db_file = dbsession.query(File).filter(File.pin == res.json()['pin']).one_or_none()
     assert db_file is not None
+    assert db_file.source == 'webapp'
+    body2 = {
+        "surname": union_member_user['surname'],
+        "number": union_member_user['union_number'],
+        "filename": "filename2.pdf",
+        "options": {"pages": "", "copies": 1, "two_sided": False},
+    }
+    res2 = client.post(url, data=json.dumps(body2))
+    assert res2.status_code == status.HTTP_200_OK
+    db_file2 = dbsession.query(File).filter(File.pin == res2.json()['pin']).one_or_none()
+    assert db_file2 is not None
+    assert db_file2.source == 'unknown'
     dbsession.delete(db_file)
+    dbsession.delete(db_file2)
     dbsession.commit()
 
 
@@ -35,6 +49,7 @@ def test_post_unauthorized_user(client):
         "surname": 'surname',
         "number": 'union_number',
         "filename": "filename.pdf",
+        "source": "vkbot",
         "options": {"pages": "", "copies": 1, "two_sided": False},
     }
     res = client.post(url, data=json.dumps(body))