Update pbs_data.py

pbs_data.py

@@ -210,6 +210,8 @@ class PBSPublicDataAPIClient:
             schedules[0]  # fallback to the most recent schedule if no match
         )
         latest_schedule = current_schedule['schedule_code']
+        schedule_year = current_schedule['effective_year']
+        schedule_month = current_schedule['effective_month']
 
         print(f"Selected schedule: {latest_schedule} (Effective: {current_schedule['effective_date']})")
 
@@ -338,6 +340,9 @@ class PBSPublicDataAPIClient:
                 pbs_code = item['pbs_code']
                 if pbs_code not in data:
                     data[pbs_code] = {
+                        "schedule_code": latest_schedule,
+                        "schedule_year": schedule_year,
+                        "schedule_month": schedule_month,
                         "name": item['drug_name'],
                         "brands": [],  # Change this to a list
                         "formulation": classify_formulation(item['li_form']),
@@ -414,7 +419,10 @@ class PBSPublicDataAPIClient:
                         'streamlined_code': restriction['streamlined_code'],
                         'online_application': restriction['online_application'],
                         'authority_method': restriction['authority_method'],
-                        'hospital_type': item['hospital_type']
+                        'hospital_type': item['hospital_type'],
+                        'schedule_code': item['schedule_code'],
+                        'schedule_year': item['schedule_year'],
+                        'schedule_month': item['schedule_month']
                     })
         
         return {k: sorted(v) if isinstance(v, set) else v for k, v in processed.items()}
@@ -446,7 +454,7 @@ class PBSPublicDataAPIClient:
                           (id INTEGER PRIMARY KEY, pbs_code TEXT, drug_id INTEGER, brand_id INTEGER,
                            formulation_id INTEGER, indication_id INTEGER, treatment_phase_id INTEGER,
                            streamlined_code TEXT, online_application BOOLEAN, authority_method TEXT,
-                           hospital_type_id INTEGER,
+                           hospital_type_id INTEGER, schedule_code TEXT, schedule_year INTEGER, schedule_month TEXT,
                            FOREIGN KEY (drug_id) REFERENCES drugs(id),
                            FOREIGN KEY (brand_id) REFERENCES brands(id),
                            FOREIGN KEY (formulation_id) REFERENCES formulations(id),
@@ -464,7 +472,7 @@ class PBSPublicDataAPIClient:
             cursor.execute('''INSERT INTO combinations
                               (pbs_code, drug_id, brand_id, formulation_id, indication_id,
                                treatment_phase_id, streamlined_code, online_application, authority_method,
-                               hospital_type_id)
+                               hospital_type_id, schedule_code, schedule_year, schedule_month)
                               VALUES (?, 
                                       (SELECT id FROM drugs WHERE name = ?),
                                       (SELECT id FROM brands WHERE name = ?),
@@ -472,10 +480,12 @@ class PBSPublicDataAPIClient:
                                       (SELECT id FROM indications WHERE name = ?),
                                       (SELECT id FROM treatment_phases WHERE name = ?),
                                       ?, ?, ?,
-                                      (SELECT id FROM hospital_types WHERE name = ?))''',
+                                      (SELECT id FROM hospital_types WHERE name = ?),
+                                      ?, ?, ?)''',
                            (combo['pbs_code'], combo['drug'], combo['brand'], combo['formulation'],
                             combo['indication'], combo['treatment_phase'], combo['streamlined_code'],
-                            combo['online_application'], combo['authority_method'], combo['hospital_type']))
+                            combo['online_application'], combo['authority_method'], combo['hospital_type'],
+                            combo['schedule_code'], combo['schedule_year'], combo['schedule_month']))
 
         # Add last_updated column and insert timestamp
         cursor.execute('''CREATE TABLE IF NOT EXISTS metadata