/ src / revolve / source_template / postgres / db_utils.py
db_utils.py
 1  import uuid
 2  import os
 3  import psycopg2
 4  from datetime import datetime, date
 5  import json
 6  from uuid import UUID
 7  
 8  def get_db_connection(test_mode):
 9      try:
10          _db_name = os.getenv("DB_NAME") if not test_mode else os.getenv("DB_NAME_TEST")
11          return psycopg2.connect(
12              dbname=_db_name,
13              user=os.getenv("DB_USER"),
14              password=os.getenv("DB_PASSWORD"),
15              host=os.getenv("DB_HOST"),
16              port=os.getenv("DB_PORT"),
17          )
18      except psycopg2.Error as e:
19          raise Exception(f"Database connection error: {e}")
20  
21  def json_serial(obj):
22      if isinstance(obj, (datetime, date)):
23          return obj.isoformat()
24      raise TypeError(f"Type {type(obj)} not serializable")
25  
26  def is_valid_uuid(val):
27      try:
28          UUID(str(val))
29          return True
30      except Exception:
31          return False
32  
33  def sanitize_str(val):
34      if val is None:
35          return None
36      return str(val)
37  
38  def json_serial(obj):
39      if isinstance(obj, (datetime, date)):
40          return obj.isoformat()
41      if isinstance(obj, UUID):
42          return str(obj)
43      raise TypeError(f"Type {type(obj)} not serializable")
44  
45  def sanitize_str(val):
46      if not isinstance(val, str):
47          raise ValueError("Expected string value")
48      return val
49  
50  def sanitize_uuid(val):
51      try:
52          return str(UUID(val))
53      except Exception:
54          raise ValueError("Invalid UUID format")
55  
56  def sanitize_bool(val):
57      if isinstance(val, bool):
58          return val
59      if isinstance(val, str):
60          if val.lower() in ['true', '1', 't', 'yes']:
61              return True
62          elif val.lower() in ['false', '0', 'f', 'no']:
63              return False
64      raise ValueError("Invalid boolean value")
65  
66  def sanitize_json(val):
67      if val is None:
68          return None
69      if isinstance(val, dict):
70          return val
71      try:
72          return json.loads(val)
73      except Exception:
74          raise ValueError("Invalid JSON value")
75  
76  def sanitize_array(val):
77      if isinstance(val, list):
78          return val
79      if isinstance(val, str):
80          try:
81              return json.loads(val)
82          except Exception:
83              raise ValueError("Invalid array value")
84      raise ValueError("Invalid array value")
85  
86  def get_unique_id():
87      """Generate a unique UUID."""
88      return str(uuid.uuid4())