Compare commits
2 Commits
ca82885fb7
...
47bb6585d1
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
47bb6585d1 | ||
|
|
c271656f51 |
4
DB.py
4
DB.py
@@ -9,6 +9,7 @@ class Database:
|
||||
|
||||
|
||||
def _setup_tables(self):
|
||||
# create the shifts table.
|
||||
cur = self._con.cursor()
|
||||
cur.execute('''
|
||||
CREATE TABLE IF NOT EXISTS shifts (
|
||||
@@ -27,10 +28,12 @@ class Database:
|
||||
|
||||
|
||||
def insert_shifts(self, shifts):
|
||||
# start db connection
|
||||
cur = self._con.cursor()
|
||||
inserted = 0
|
||||
|
||||
for shift in shifts:
|
||||
# parse all information from the shift
|
||||
start = shift.get('start_datetime', 'N/A')
|
||||
end = shift.get('end_datetime', 'N/A')
|
||||
|
||||
@@ -38,6 +41,7 @@ class Database:
|
||||
duration = shift.get('duration', 'N/A')
|
||||
description = shift.get('shift_remark', shift.get('description', 'N/A'))
|
||||
|
||||
# insert it into the db.
|
||||
cur.execute("""
|
||||
INSERT OR IGNORE INTO shifts
|
||||
(shift_start, shift_end, department, duration, description)
|
||||
|
||||
21
PMT.py
21
PMT.py
@@ -15,7 +15,6 @@ class PMT:
|
||||
self.base_url = os.environ.get("URL_PMT")
|
||||
|
||||
def get_week_range(self, week_offset):
|
||||
"""Get Monday date for week_offset weeks from now"""
|
||||
today = datetime.now()
|
||||
days_to_monday = 0 if today.weekday() == 0 else 7 - today.weekday()
|
||||
monday = today + timedelta(days=days_to_monday)
|
||||
@@ -23,12 +22,13 @@ class PMT:
|
||||
return target_monday.strftime("%Y-%m-%d")
|
||||
|
||||
def login(self):
|
||||
|
||||
self.session = requests.Session()
|
||||
|
||||
login_url = f"{self.base_url}/login"
|
||||
resp = self.session.get(login_url)
|
||||
|
||||
# Step 2: SSO login POST
|
||||
# SSO login POST
|
||||
sso_data = {
|
||||
"username": self._username,
|
||||
"password": self._password,
|
||||
@@ -44,23 +44,29 @@ class PMT:
|
||||
json=sso_data,
|
||||
headers={'Content-Type': 'application/json'})
|
||||
|
||||
|
||||
# Check if it succeeded.
|
||||
if sso_resp.status_code != 200 or not sso_resp.json().get('result', {}).get('authenticated'):
|
||||
print("Login failed!")
|
||||
print(sso_resp.text)
|
||||
sys.exit(1)
|
||||
|
||||
# Store the useful information in some class fields
|
||||
login_data = sso_resp.json()['result']
|
||||
self.context_token = login_data['context_token']
|
||||
self.user_token = login_data['user_token']
|
||||
|
||||
print(f"✅ Logged in as {self._username} (store AH 8541 Rhoon)")
|
||||
# return session info
|
||||
print(f"✅ Logged in as {self._username}")
|
||||
return self.session, self.context_token, self.user_token
|
||||
|
||||
def get_shifts(self, days_ahead=14):
|
||||
"""Fetch shifts from now to days_ahead using the REAL endpoint"""
|
||||
# raise if not logged in.
|
||||
if not self.session or not self.context_token or not self.user_token:
|
||||
raise ValueError("Must call login() first!")
|
||||
|
||||
|
||||
# define date ranges
|
||||
shifts = []
|
||||
|
||||
from_date = datetime.now().strftime("%Y-%m-%d")
|
||||
@@ -68,6 +74,7 @@ class PMT:
|
||||
|
||||
print(f"Fetching shifts from {from_date} to {to_date}...")
|
||||
|
||||
# send request for shifts
|
||||
headers = {
|
||||
'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:148.0) Gecko/20100101 Firefox/148.0',
|
||||
'Accept': 'application/json, text/plain, */*',
|
||||
@@ -91,6 +98,8 @@ class PMT:
|
||||
|
||||
resp = self.session.get(shifts_url, headers=headers, params=params)
|
||||
|
||||
|
||||
# check if request succeeded
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
week_shifts = data.get('result', {}).get('shift_instances', [])
|
||||
@@ -100,16 +109,16 @@ class PMT:
|
||||
print(f"❌ Failed: {resp.status_code}")
|
||||
print(resp.text[:500])
|
||||
|
||||
# return all shifts fetched
|
||||
return shifts
|
||||
|
||||
@classmethod
|
||||
def print_shifts(cls, shifts):
|
||||
"""Pretty print shifts (class method since it doesn't need instance)"""
|
||||
if not shifts:
|
||||
print("No upcoming shifts found.")
|
||||
return
|
||||
|
||||
print("\n📅 UPCOMING SHIFTS:\n")
|
||||
print("\nUPCOMING SHIFTS:\n")
|
||||
print("Date | Time | Dept | Duration | Description")
|
||||
print("-" * 80)
|
||||
|
||||
|
||||
7
main.py
7
main.py
@@ -4,21 +4,26 @@ from dotenv import load_dotenv
|
||||
import os
|
||||
|
||||
def main():
|
||||
# Load the environment
|
||||
load_dotenv(override=True)
|
||||
|
||||
|
||||
# Set up database, webserver and PMT connection
|
||||
db = Database()
|
||||
pmt = PMT()
|
||||
|
||||
|
||||
# fetch PMT shifts
|
||||
pmt.login()
|
||||
|
||||
days_ahead = int(os.environ.get("DAYS_TO_FETCH"))
|
||||
shifts = pmt.get_shifts(days_ahead)
|
||||
PMT.print_shifts(shifts)
|
||||
|
||||
|
||||
# Insert those shifts into the database
|
||||
db.insert_shifts(shifts)
|
||||
|
||||
print(f"\nFound {len(shifts)} shifts total.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user