Merge pull request #72 from 0x41424142/actions/black

qualyspy/sql/base.py

@@ -75,7 +75,13 @@ def db_connect(
     return engine.connect()
 
 
-def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict, override_import_dt: datetime = None) -> int:
+def upload_data(
+    df: DataFrame,
+    table: str,
+    cnxn: Connection,
+    dtype: dict,
+    override_import_dt: datetime = None,
+) -> int:
     """
     Upload a DataFrame to a SQL table. Appends 'import_datetime' column to the DataFrame.
 
@@ -91,7 +97,9 @@ def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict, overri
     """
 
     # Add an import_datetime column:
-    df["import_datetime"] = datetime.now() if not override_import_dt else override_import_dt
+    df["import_datetime"] = (
+        datetime.now() if not override_import_dt else override_import_dt
+    )
     dtype["import_datetime"] = types.DateTime()
 
     # For any string values in the DataFrame, make sure it doesn't

qualyspy/sql/supported_uploads.py

@@ -11,14 +11,16 @@
 from ..vmdr.data_classes.lists import BaseList
 
 
-def upload_vmdr_ags(ags: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int:
+def upload_vmdr_ags(
+    ags: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.get_ag_list() to SQL.
 
     Parameters:
     ags (BaseList): The Asset Group List to upload.
     cnxn (Connection): The Connection object to the SQL database.
-    override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. 
+    override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database.
 
     Returns:
     int: The number of rows uploaded.
@@ -66,10 +68,14 @@ def upload_vmdr_ags(ags: BaseList, cnxn: Connection, override_import_dt: datetim
     df = DataFrame([prepare_dataclass(ag) for ag in ags])
 
     # Upload the data:
-    return upload_data(df, "vmdr_assetgroups", cnxn, dtype=COLS, override_import_dt=override_import_dt)
+    return upload_data(
+        df, "vmdr_assetgroups", cnxn, dtype=COLS, override_import_dt=override_import_dt
+    )
 
 
-def upload_vmdr_kb(kbs: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int:
+def upload_vmdr_kb(
+    kbs: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.query_kb() to SQL.
 
@@ -119,10 +125,21 @@ def upload_vmdr_kb(kbs: BaseList, cnxn: Connection, override_import_dt: datetime
     df = DataFrame([prepare_dataclass(kb) for kb in kbs])
 
     # Upload the data:
-    return upload_data(df, "vmdr_knowledgebase", cnxn, dtype=COLS, override_import_dt=override_import_dt)
+    return upload_data(
+        df,
+        "vmdr_knowledgebase",
+        cnxn,
+        dtype=COLS,
+        override_import_dt=override_import_dt,
+    )
 
 
-def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False, override_import_dt: datetime = None) -> int:
+def upload_vmdr_hosts(
+    hosts: BaseList,
+    cnxn: Connection,
+    is_hld: bool = False,
+    override_import_dt: datetime = None,
+) -> int:
     """
     Upload data from vmdr.get_host_list() to SQL.
 
@@ -197,11 +214,17 @@ def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False, o
 
     # Upload the data, with table depdening on if it is a Host List Detail pull or not:
     return upload_data(
-        df, "vmdr_hosts_list" if not is_hld else "vmdr_hld_hosts_list", cnxn, dtype=COLS, override_import_dt=override_import_dt
+        df,
+        "vmdr_hosts_list" if not is_hld else "vmdr_hld_hosts_list",
+        cnxn,
+        dtype=COLS,
+        override_import_dt=override_import_dt,
     )
 
 
-def upload_vmdr_ips(ips: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int:
+def upload_vmdr_ips(
+    ips: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.get_ip_list() to SQL.
 
@@ -226,10 +249,14 @@ def upload_vmdr_ips(ips: BaseList, cnxn: Connection, override_import_dt: datetim
     df["TYPE"] = df["IP"].apply(lambda x: "Single IP" if "/" not in x else "IP Range")
 
     # Upload the data:
-    return upload_data(df, "vmdr_ips", cnxn, dtype=COLS, override_import_dt=override_import_dt)
+    return upload_data(
+        df, "vmdr_ips", cnxn, dtype=COLS, override_import_dt=override_import_dt
+    )
 
 
-def upload_vmdr_hld(hld: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int:
+def upload_vmdr_hld(
+    hld: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.get_hld() to SQL.
 
@@ -296,4 +323,10 @@ def upload_vmdr_hld(hld: BaseList, cnxn: Connection, override_import_dt: datetim
     df["QDS"] = df["QDS"].apply(lambda x: int(x) if x else None)
 
     # Upload the data:
-    return upload_data(df, "vmdr_hld_detections", cnxn, dtype=COLS, override_import_dt=override_import_dt)
+    return upload_data(
+        df,
+        "vmdr_hld_detections",
+        cnxn,
+        dtype=COLS,
+        override_import_dt=override_import_dt,
+    )