109 lines
2.6 KiB
Python
109 lines
2.6 KiB
Python
import os
|
|
from pprint import pprint
|
|
from dataclasses import dataclass
|
|
from litestar import Litestar, get, post
|
|
from litestar.logging import LoggingConfig
|
|
import psycopg2
|
|
import logging
|
|
|
|
logger = logging.getLogger()
|
|
|
|
@dataclass
|
|
class PostRequest:
|
|
hostname: str
|
|
source: str
|
|
uuid: str
|
|
file_path: str
|
|
file_size: int
|
|
shasum: str = None
|
|
|
|
@get("/")
|
|
async def hello_world() -> str:
|
|
return "Hello, world!"
|
|
|
|
@get("/files/uniq")
|
|
async def files_uniq() -> str:
|
|
db = psycopg2.connect('dbname=file_index_ian')
|
|
db.autocommit = True
|
|
cur = db.cursor()
|
|
sql = """
|
|
SELECT * FROM fileinfo
|
|
"""
|
|
cur.execute(sql)
|
|
filesnames = {}
|
|
for elt in cur.fetchall():
|
|
name = elt[0]
|
|
name = name[name.rindex('/')+1:]
|
|
if name not in filesnames:
|
|
filesnames[name] = 0
|
|
else:
|
|
filesnames[name] += 1
|
|
|
|
|
|
duplicates = [
|
|
(elt, n)
|
|
for elt, n in filesnames.items()
|
|
if n > 0
|
|
]
|
|
|
|
pprint(duplicates)
|
|
cur.close()
|
|
db.close()
|
|
|
|
return 'ok'
|
|
|
|
|
|
@post("/")
|
|
def register_file(data: PostRequest) -> str:
|
|
db = psycopg2.connect('dbname=file_index_ian')
|
|
db.autocommit = True
|
|
cur = db.cursor()
|
|
try:
|
|
sql = f"""
|
|
INSERT INTO source VALUES (%s, %s, %s);
|
|
"""
|
|
|
|
cur.execute(sql, (data.hostname, data.source, data.uuid))
|
|
except psycopg2.errors.UniqueViolation:
|
|
cur.close()
|
|
pass
|
|
|
|
if data.shasum:
|
|
try:
|
|
cur = db.cursor()
|
|
sql = f"""
|
|
INSERT INTO shasum VALUES (%s);
|
|
"""
|
|
cur.execute(sql, (data.shasum, ))
|
|
except psycopg2.errors.UniqueViolation:
|
|
cur.close()
|
|
pass
|
|
try:
|
|
cur = db.cursor()
|
|
if data.shasum:
|
|
sql = f"""
|
|
INSERT INTO fileinfo (file_path, file_size, source_uuid, shasum) VALUES (%s, %s, %s, %s);
|
|
"""
|
|
cur.execute(sql, (data.file_path, data.file_size, data.uuid, data.shasum))
|
|
else:
|
|
sql = f"""
|
|
INSERT INTO fileinfo (file_path, file_size, source_uuid) VALUES (%s, %s, %s);
|
|
"""
|
|
cur.execute(sql, (data.file_path, data.file_size, data.uuid))
|
|
cur.close()
|
|
except psycopg2.errors.UniqueViolation:
|
|
pass
|
|
db.close()
|
|
|
|
|
|
logger.info(f"Inserted {data.hostname}:{data.uuid}/{data.file_path} into database.")
|
|
return f"Inserted {data.hostname}:{data.uuid}/{data.file_path} into database."
|
|
|
|
app = Litestar([
|
|
register_file,
|
|
files_uniq
|
|
], logging_config=LoggingConfig(
|
|
root={'level': 'DEBUG', 'handlers': ['queue_listener']},
|
|
log_exceptions='always'
|
|
))
|