Added check for 2 formats

Author: jprakash-db

src/databricks/sql/utils.py

@@ -628,6 +628,20 @@ def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table":
     return table
 
 
+def datetime_parser(datetime_string):
+    formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"]
+
+    for fmt in formats:
+        try:
+            return datetime.datetime.strptime(datetime_string, fmt)
+        except ValueError:
+            continue
+
+    raise ValueError(
+        f"Datetime string '{datetime_string}' does not match any expected formats"
+    )
+
+
 def convert_to_assigned_datatypes_in_column_table(column_table, description):
 
     converted_column_table = []
@@ -642,16 +656,7 @@ def convert_to_assigned_datatypes_in_column_table(column_table, description):
             )
         elif description[i][1] == "timestamp":
             converted_column_table.append(
-                tuple(
-                    (
-                        v
-                        if v is None
-                        else datetime.datetime.strptime(
-                            v, "%Y-%m-%d %H:%M:%S.%f"
-                        ).replace(tzinfo=pytz.UTC)
-                    )
-                    for v in col
-                )
+                tuple((v if v is None else datetime_parser(v)) for v in col)
             )
         else:
             converted_column_table.append(col)