Merge pull request #73 from 0x41424142/dt_fixes

docs/sql.md

@@ -44,7 +44,10 @@ And finally, you can use the following supported functions:
 
 >**Head's Up:** More upload functions are coming soon!
 
-Each upload function takes 2 parameters. The first is the ```BaseList``` of data, and the second is the ```sqlalchemy.Connection``` object you built above.
+Each upload function takes 2 positional parameters. The first is the ```BaseList``` of data, and the second is the ```sqlalchemy.Connection``` object you built above. 
+
+Functions also take an optional ```override_import_dt``` parameter that will set the resulting SQL table's ```import_datetime``` field to the value you specify. ```override_import_dt``` is a ```datetime.datetime``` object.
+
 
 | Function Name | Module  | ```qualyspy``` Function Data Source | Resulting Table Name |
 | -- | -- | -- | -- |
@@ -54,6 +57,18 @@ Each upload function takes 2 parameters. The first is the ```BaseList``` of data
 | ```upload_vmdr_hld``` | VMDR | ```vmdr.get_hld()```| ```vmdr_hld_hosts_list``` for hosts and ```vmdr_hld_detections``` for detections |
 | ```upload_vmdr_ips``` | VMDR | ```vmdr.get_ip_list()```| ```vmdr_ips``` |
 
+```py
+from qualyspy.sql import *
+
+# Get a connection to the DB
+cnxn = db_connect(host='10.0.0.1', db='qualysdata', trusted_cnxn=True)
+
+# Upload a previous vmdr.get_host_list() call to the DB, with override_import_dt set
+# to 10-25-2023 12:00:00
+dt = datetime.datetime(2023, 10, 25, 12, 0, 0)
+uploaded = upload_vmdr_hosts(vmdr_hosts, cnxn, override_import_dt=dt)
+>>>Uploaded 12345 records to vmdr_hosts_list
+
 ## A Friendly Recommendation For Getting Data
 
 When calling any of the data source functions to get the data to upload, it is recommended to make the call as verbose as possible via kwargs. 

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) -> 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.
 
@@ -84,13 +90,16 @@ def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict) -> int
         table (str): The name of the table to upload to.
         cnxn (Connection): The Connection object to the SQL database.
         dtype (dict): The data types of the columns in the table. Key is the column name, value is the data type as sqlalchemy.types.Something()
+        override_import_dt (datetime): If provided, will override the import_datetime column with this value.
 
     Returns:
     int: The number of rows uploaded.
     """
 
     # Add an import_datetime column:
-    df["import_datetime"] = datetime.now()
+    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

@@ -2,20 +2,25 @@
 supported_uploads.py - Contains the functions to upload supported API pulls to SQL Server.
 """
 
+from datetime import datetime
+
 from pandas import DataFrame
 from sqlalchemy import Connection, types
 
 from .base import upload_data, prepare_dataclass
 from ..vmdr.data_classes.lists import BaseList
 
 
-def upload_vmdr_ags(ags: BaseList, cnxn: Connection) -> 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.
 
     Returns:
     int: The number of rows uploaded.
@@ -63,16 +68,21 @@ def upload_vmdr_ags(ags: BaseList, cnxn: Connection) -> int:
     df = DataFrame([prepare_dataclass(ag) for ag in ags])
 
     # Upload the data:
-    return upload_data(df, "vmdr_assetgroups", cnxn, dtype=COLS)
+    return upload_data(
+        df, "vmdr_assetgroups", cnxn, dtype=COLS, override_import_dt=override_import_dt
+    )
 
 
-def upload_vmdr_kb(kbs: BaseList, cnxn: Connection) -> int:
+def upload_vmdr_kb(
+    kbs: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.query_kb() to SQL.
 
     Parameters:
     kbs (BaseList): The KB 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.
 
     Returns:
     int: The number of rows uploaded.
@@ -115,17 +125,29 @@ def upload_vmdr_kb(kbs: BaseList, cnxn: Connection) -> int:
     df = DataFrame([prepare_dataclass(kb) for kb in kbs])
 
     # Upload the data:
-    return upload_data(df, "vmdr_knowledgebase", cnxn, dtype=COLS)
+    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) -> 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.
 
     Parameters:
         hosts (BaseList): The Host List to upload.
         cnxn (Connection): The Connection object to the SQL database.
         is_hld (bool): If the data is from a Host List Detail pull. You can ignore this.
+        override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database.
 
     Returns:
         int: The number of rows uploaded.
@@ -192,17 +214,24 @@ def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False) -
 
     # 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
+        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) -> int:
+def upload_vmdr_ips(
+    ips: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.get_ip_list() to SQL.
 
     Parameters:
         ips (BaseList): The IP List to upload from vmdr.get_ip_list().
         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.
 
     Returns:
         int: The number of rows uploaded.
@@ -220,10 +249,14 @@ def upload_vmdr_ips(ips: BaseList, cnxn: Connection) -> int:
     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)
+    return upload_data(
+        df, "vmdr_ips", cnxn, dtype=COLS, override_import_dt=override_import_dt
+    )
 
 
-def upload_vmdr_hld(hld: BaseList, cnxn: Connection) -> int:
+def upload_vmdr_hld(
+    hld: BaseList, cnxn: Connection, override_import_dt: datetime = None
+) -> int:
     """
     Upload data from vmdr.get_hld() to SQL.
 
@@ -290,4 +323,10 @@ def upload_vmdr_hld(hld: BaseList, cnxn: Connection) -> int:
     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)
+    return upload_data(
+        df,
+        "vmdr_hld_detections",
+        cnxn,
+        dtype=COLS,
+        override_import_dt=override_import_dt,
+    )