hipom_data_mapping/data_import/select_db.py

60 lines
2.2 KiB
Python

import psycopg2
import pandas as pd
# Function to read the db connection info
def read_db_connection_info(filename="db_connection_info.txt"):
connection_info = {}
try:
with open(filename, 'r') as file:
for line in file:
key, value = line.strip().split('=')
connection_info[key] = value
except Exception as e:
print(f"Failed to read database connection info: {e}")
raise
return connection_info
# Load the connection info
connection_info = read_db_connection_info()
try:
# Connect to the database
conn = psycopg2.connect(
host=connection_info["host"],
user=connection_info["user"],
password=connection_info["password"],
dbname=connection_info["database"],
port=connection_info["port"]
)
# This ensures that resources are cleaned up properly
with conn:
with conn.cursor() as cursor:
# Export data_mapping table
query_mapping = """
SELECT * FROM data_mapping
WHERE ships_idx BETWEEN 1000 AND 1999
"""
cursor.execute(query_mapping)
results_mapping = cursor.fetchall()
columns_mapping = [desc[0] for desc in cursor.description]
df_mapping = pd.DataFrame(results_mapping, columns=columns_mapping)
df_mapping.to_csv('exports/data_mapping.csv', index=False, encoding='utf-8-sig')
# Export data_master_model table
query_master = """
SELECT * FROM data_model_master
"""
cursor.execute(query_master)
results_master = cursor.fetchall()
columns_master = [desc[0] for desc in cursor.description]
df_master = pd.DataFrame(results_master, columns=columns_master)
df_master.to_csv('exports/data_model_master_export.csv', index=False, encoding='utf-8-sig')
print("Data exported successfully to 'exports/data_mapping.csv' and 'exports/data_model_master_export.csv'")
except (Exception, psycopg2.DatabaseError) as error:
print(f"An error occurred: {error}")
finally:
if conn is not None:
conn.close()