Merge pull request #58 from City-of-Helsinki/enhancement/improve_popu…

…late_script Enhancement/improve populate script
City-of-Helsinki · Nov 7, 2023 · b6e325a · b6e325a
2 parents c28025d + 4c88047
commit b6e325a
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 34 deletions.
diff --git a/mittaridatapumppu-deviceregistry/devices/scripts/populate_db_via_api.py b/mittaridatapumppu-deviceregistry/devices/scripts/populate_db_via_api.py
@@ -1,12 +1,11 @@
 import argparse
 import csv
-from pprint import pprint
+import logging
+from pprint import pprint, pformat
 
 import httpx
 import openpyxl
 
-from uirasmeta import META
-
 from django.utils.text import slugify
 
 
@@ -15,14 +14,18 @@ def parse_args():
     parser = argparse.ArgumentParser(description="Populate database via API")
     parser.add_argument("--api-url", type=str, help="API base URL", required=True)
     parser.add_argument("--api-token", type=str, help="API token", required=True)
+    parser.add_argument("--django-user", type=str, help="Owner of devices", required=True)
     # Create mutually exclusive group for csv and excel files
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument("--csv-file", type=str, help="Data file")
     group.add_argument("--excel-file", type=str, help="Excel data file")
     group.add_argument("--uirasmeta", action="store_true", help="Use Uirasmeta data")
     # add argument to choose which sheets will be used
     parser.add_argument("--sheets", nargs="+", type=str, help="Sheets to use in Excel file")
+    parser.add_argument("--log", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], default="ERROR")
     args = parser.parse_args()
+    logging.basicConfig(level=getattr(logging, args.log))
+    args.api_url = args.api_url.rstrip("/")
     return args
 
 
@@ -37,12 +40,10 @@ def parse_args():
 
 def get_device_type():
     """Check if one device type exists, if not create it and return URL."""
-    url = "{}/device-types/".format(args.api_url.rstrip("/"))
+    url = "{}/device-types/".format(args.api_url)
     r = httpx.get(url, headers=headers)
     if r.status_code != 200:
-        print("Failed to get device types")
-        print(r.status_code)
-        print(r.text)
+        logging.critical(f"Failed to get device types from {url}: {r.status_code} {r.text}")
         exit(1)
     r_data = r.json()
     if r_data["count"] > 0:
@@ -55,36 +56,34 @@ def get_device_type():
         }
         r = httpx.post(url, headers=headers, json=device_type_data)
         if r.status_code != 201:
-            print("Failed to create device type")
-            print(r.status_code)
-            print(r.text)
+            logging.critical(f"Failed to create device type: {r.status_code} {r.text}")
             exit(1)
         r_data = r.json()
         return r_data["url"]
 
 
 def post_device_data(args: argparse.Namespace, device_type: str, data: dict):
-    url = "{}/devices/".format(args.api_url.rstrip("/"))
+    url = "{}/devices/".format(args.api_url)
     device_data = {
         "device_id": data["device_id"],
         "name": data["name"],
-        "owner": "http://127.0.0.1:8000/api/v1/users/1/",
+        "owner": f"{args.api_url}/users/{args.django_user}/",
         "type": device_type,
         "parser_module": data["parser_module"],
         "logs": [],
         "images": [],
     }
+    logging.info(f"Posting device data to {url}: {device_data}")
     # Check first if device exists
     device_url = f"{url}{data['device_id']}/"
     r = httpx.get(device_url, headers=headers)
     if r.status_code == 200:
-        print("Device exists, updating...")
+        logging.debug("Device exists, updating...")
         r = httpx.put(device_url, headers=headers, json=device_data)
     else:  # POST new device
-        print("Device does not exist, creating...")
+        logging.debug("Device does not exist, creating...")
         r = httpx.post(url, headers=headers, json=device_data)
-
-    print(r.status_code, r.text)
+    logging.info(r.status_code, r.text)
 
 
 def parse_csv(args: argparse.Namespace):
@@ -94,7 +93,7 @@ def parse_csv(args: argparse.Namespace):
     with open(args.csv_file, newline="") as csvfile:
         lines = csv.DictReader(csvfile)
         for line in lines:
-            print(line)
+            logging.debug(line)
             post_device_data(args, device_type, line)
 
 
@@ -122,29 +121,36 @@ def read_excel_sheet(wb: openpyxl.Workbook, sheet_name: str) -> list:
 
 
 def post_data(args: argparse.Namespace, endpoint: str, lookup_value: str, data: dict, patch: bool = False):
-    url = "{}/{}/".format(args.api_url.rstrip("/"), endpoint)
+    url = "{}/{}/".format(args.api_url, endpoint)
     object_url = "{}{}/".format(url, lookup_value)
-    print(object_url, end=": ")
     r = httpx.get(object_url, headers=headers)
-    print(r.status_code, r.text)
+    logging.debug(f"{object_url}: {r.status_code} {r.text[:100]}")
     if r.status_code == 200:
         if patch:
-            print(f"{lookup_value} exists, patching...")
+            logging.debug(f"{lookup_value} exists, patching...")
             r = httpx.patch(object_url, headers=headers, json=data)
         else:
-            print(f"{lookup_value} exists, updating...")
+            logging.debug(f"{lookup_value} exists, updating...")
             r = httpx.put(object_url, headers=headers, json=data)
+            # pretty_data = pformat(data)
+            # logging.debug(f"{r.status_code} {r.text} {pretty_data}")
     else:  # POST new device
-        print(f"{lookup_value} does not exist, creating...")
+        logging.debug(f"{lookup_value} does not exist, creating...")
         r = httpx.post(url, headers=headers, json=data)
-    print(r.status_code, r.text)
+        if r.status_code >= 400:
+            pretty_data = pformat(data)
+            logging.error(f"{r.status_code} {r.text} {pretty_data}")
+            exit(1)
+    logging.debug(f"{r.status_code} {r.text}")
     return r.json()
 
 
 def parse_excel(args: argparse.Namespace):
     wb = read_excel(args.excel_file)
 
     if args.sheets is None or "DeviceType" in args.sheets:
+        placeholder = {"name": "Placeholder", "slug": "placeholder", "description": "Placeholder device type"}
+        post_data(args, "device-types", "placeholder", placeholder)
         device_types = read_excel_sheet(wb, "DeviceType")
         for device_type in device_types:
             post_data(args, "device-types", device_type["slug"], device_type)
@@ -157,8 +163,8 @@ def parse_excel(args: argparse.Namespace):
     if args.sheets is None or "Device" in args.sheets:
         devices = read_excel_sheet(wb, "Device")
         for device in devices:
-            device["owner"] = "http://127.0.0.1:8000/api/v1/users/root/"
-            device["type"] = "http://127.0.0.1:8000/api/v1/device-types/{}/".format(device.get("type", "placeholder"))
+            device["owner"] = f"{args.api_url}/users/{args.django_user}/"
+            device["type"] = "{}/device-types/{}/".format(args.api_url, device.get("type", "placeholder"))
             # TODO: make these fields optional
             device["logs"] = []
             device["images"] = []
@@ -172,6 +178,8 @@ def parse_excel(args: argparse.Namespace):
 
 
 def parse_uirasmeta(args: argparse.Namespace):
+    from uirasmeta import META
+
     for device_id, data in META.items():
         location_slug = slugify(data["name"])[:50]
         print(device_id, location_slug, data)

diff --git a/pyproject.toml b/pyproject.toml
@@ -16,7 +16,7 @@ dependencies = [
   "Django >= 4.2.6, < 5.0",
   "django-auditlog",
   "django-environ",
-  "django-rest-framework",
+  "djangorestframework",
   "openpyxl",
   "Pillow",
   "psycopg2-binary",

diff --git a/requirements.txt b/requirements.txt
@@ -10,7 +10,7 @@ certifi==2023.7.22
     # via sentry-sdk
 click==8.1.7
     # via uvicorn
-django==4.2.6
+django==4.2.7
     # via
     #   django-auditlog
     #   djangorestframework
@@ -20,10 +20,8 @@ django-auditlog==2.3.0
     # via mittaridatapumppu-deviceregistry (pyproject.toml)
 django-environ==0.11.2
     # via mittaridatapumppu-deviceregistry (pyproject.toml)
-django-rest-framework==0.1.0
-    # via mittaridatapumppu-deviceregistry (pyproject.toml)
 djangorestframework==3.14.0
-    # via django-rest-framework
+    # via mittaridatapumppu-deviceregistry (pyproject.toml)
 et-xmlfile==1.1.0
     # via openpyxl
 h11==0.14.0
@@ -44,13 +42,13 @@ pytest==7.4.3
     # via
     #   mittaridatapumppu-deviceregistry (pyproject.toml)
     #   pytest-django
-pytest-django==4.5.2
+pytest-django==4.6.0
     # via mittaridatapumppu-deviceregistry (pyproject.toml)
 python-dateutil==2.8.2
     # via django-auditlog
 pytz==2023.3.post1
     # via djangorestframework
-sentry-sdk[django]==1.32.0
+sentry-sdk[django]==1.34.0
     # via mittaridatapumppu-deviceregistry (pyproject.toml)
 six==1.16.0
     # via python-dateutil
@@ -62,7 +60,7 @@ unittest-parametrize==1.4.0
     # via mittaridatapumppu-deviceregistry (pyproject.toml)
 urllib3==2.0.7
     # via sentry-sdk
-uvicorn==0.23.2
+uvicorn==0.24.0.post1
     # via mittaridatapumppu-deviceregistry (pyproject.toml)
 whitenoise==6.6.0
     # via mittaridatapumppu-deviceregistry (pyproject.toml)