retuve.trak.data

  1# Copyright 2024 Adam McArthur
  2#
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#     http://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""
 16Code to control inserting/reading present state
 17of results into the database.
 18"""
 19
 20import os
 21import sqlite3
 22from typing import List
 23
 24from retuve.app.classes import File, FileEnum
 25
 26# Get this files DIR
 27DIR = os.path.dirname(os.path.abspath(__file__))
 28
 29
 30def database_init(db_path: str):
 31    """
 32    Initialize the database.
 33
 34    :param db_path: The path to the database.
 35    """
 36    # delete the database if it exists
 37    if not os.path.exists(db_path):
 38        with sqlite3.connect(db_path) as conn:
 39            cursor = conn.cursor()
 40
 41            # Create the table if it doesn't exist
 42            cursor.execute(
 43                """
 44                CREATE TABLE IF NOT EXISTS files (
 45                    file_id TEXT PRIMARY KEY,
 46                    state TEXT,
 47                    metrics_url TEXT,
 48                    video_url TEXT,
 49                    img_url TEXT,
 50                    figure_url TEXT,
 51                    attempts INTEGER DEFAULT 0
 52                )
 53                """
 54            )
 55            conn.commit()
 56
 57
 58def insert_files(files: List[File], db_path: str):
 59    """
 60    Insert files into the database.
 61
 62    :param files: The list of files to insert.
 63    :param db_path: The path to the database.
 64    """
 65    with sqlite3.connect(db_path) as conn:
 66        cursor = conn.cursor()
 67
 68        for file in files:
 69            cursor.execute(
 70                """
 71                INSERT OR REPLACE INTO files (file_id, state, metrics_url, video_url, img_url, figure_url, attempts)
 72                VALUES (?, ?, ?, ?, ?, ?, ?)
 73                """,
 74                (
 75                    file.file_id,
 76                    file.state.value,
 77                    file.metrics_url,
 78                    file.video_url,
 79                    file.img_url,
 80                    file.figure_url,
 81                    file.attempts,
 82                ),
 83            )
 84            conn.commit()
 85
 86
 87# Extract function
 88def extract_files(db_path: str) -> List[File]:
 89    """
 90    Extract files from the database.
 91
 92    :param db_path: The path to the database.
 93    :return: The list of files.
 94    """
 95    with sqlite3.connect(db_path) as conn:
 96        cursor = conn.cursor()
 97        try:
 98            cursor.execute(
 99                "SELECT file_id, state, metrics_url, video_url, img_url, figure_url, attempts FROM files"
100            )
101        except sqlite3.OperationalError:
102            return []
103
104        rows = cursor.fetchall()
105        files = [
106            File(
107                file_id=row[0],
108                state=FileEnum(row[1]),
109                metrics_url=row[2],
110                video_url=row[3],
111                img_url=row[4],
112                figure_url=row[5],
113                attempts=row[6],
114            )
115            for row in rows
116        ]
117        conn.commit()
118
119        files.sort(key=lambda x: (x.state.value, x.file_id))
120
121        return files
122
123
124def get_file_data(file_id: str, db_path: str) -> File:
125    """
126    Get the file data from the database.
127
128    :param file_id: The file id.
129    :param db_path: The path to the database.
130    :return: The file data.
131    """
132    with sqlite3.connect(db_path) as conn:
133        cursor = conn.cursor()
134        cursor.execute(
135            "SELECT file_id, state, metrics_url, video_url, img_url, figure_url, attempts FROM files WHERE file_id = ?",
136            (file_id,),
137        )
138        row = cursor.fetchone()
139
140        conn.commit()
141
142        if not row:
143            return File(
144                file_id=file_id,
145                state=FileEnum.PENDING,
146                metrics_url="N/A",
147                video_url="N/A",
148                img_url="N/A",
149                figure_url="N/A",
150                attempts=0,
151            )
152
153        return File(
154            file_id=row[0],
155            state=FileEnum(row[1]),
156            metrics_url=row[2],
157            video_url=row[3],
158            img_url=row[4],
159            figure_url=row[5],
160            attempts=row[6],
161        )