89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
import os
|
|
import wmill
|
|
import yaml
|
|
from sqlalchemy import create_engine, text, MetaData, Table, Column, Integer, String, inspect
|
|
from sqlalchemy.engine import Engine
|
|
import psycopg2
|
|
|
|
# You can import any PyPi package.
|
|
# See here for more info: https://www.windmill.dev/docs/advanced/dependencies_in_python
|
|
|
|
# you can use typed resources by doing a type alias to dict
|
|
#postgresql = dict
|
|
|
|
DB_RESOURCE_PATH = 'u/joshuakrzemien/slick_postgresql'
|
|
DB_CONFIG_PATH = 'f/CCR_ETL/ccr_db_config'
|
|
|
|
def create_db_engine(db: dict):
|
|
db_url = f"postgresql+psycopg2://postgres:{db['password']}@{db['host']}:{db['port']}/{db['dbname']}"
|
|
engine = create_engine(db_url)
|
|
engine.connect()
|
|
return engine
|
|
|
|
def table_exists(engine: Engine, table_name: str) -> bool:
|
|
"""Check if a table exists in the database."""
|
|
inspector = inspect(engine)
|
|
return table_name in inspector.get_table_names()
|
|
|
|
def create_table(engine: Engine, table: dict, strategy: str = "create_if_not_exists"):
|
|
try:
|
|
table_name = table['name']
|
|
columns = table['columns']
|
|
|
|
# Handle different table strategies
|
|
if strategy == "drop_and_recreate":
|
|
if table_exists(engine, table_name):
|
|
print(f"Dropping existing table: {table_name}")
|
|
with engine.connect() as conn:
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {table_name} CASCADE"))
|
|
conn.commit()
|
|
|
|
elif strategy == "create_if_not_exists":
|
|
if table_exists(engine, table_name):
|
|
print(f"Table {table_name} already exists, skipping creation")
|
|
return
|
|
|
|
else:
|
|
raise ValueError(f"Unknown table strategy: {strategy}")
|
|
|
|
# Map config types to SQLAlchemy types
|
|
type_mapping = {
|
|
'integer': Integer,
|
|
'string': String
|
|
}
|
|
|
|
# Build SQLAlchemy columns
|
|
sqlalchemy_columns = []
|
|
for column in columns:
|
|
col_type = type_mapping.get(column['type'], String)
|
|
sqlalchemy_columns.append(Column(column['name'], col_type, primary_key=column.get('primary_key', False), nullable=column.get('nullable', True), index=column.get('index', False), autoincrement=column.get('autoincrement', False)))
|
|
|
|
# Create table using SQLAlchemy Core
|
|
metadata = MetaData()
|
|
new_table = Table(table_name, metadata, *sqlalchemy_columns)
|
|
|
|
# Create the table
|
|
metadata.create_all(engine)
|
|
|
|
print(f"Successfully created table: {table_name}")
|
|
|
|
except Exception as e:
|
|
print(f"Error creating table {table_name}: {str(e)}")
|
|
raise
|
|
|
|
def main():
|
|
db = wmill.client.get_resource(DB_RESOURCE_PATH)
|
|
config_yaml = wmill.get_variable(DB_CONFIG_PATH)
|
|
config = yaml.safe_load(config_yaml)
|
|
engine = create_db_engine(db)
|
|
|
|
# Get table strategy from config (default to drop_and_recreate)
|
|
table_strategy = config.get('table_strategy', 'drop_and_recreate')
|
|
print(f"Using table strategy: {table_strategy}")
|
|
|
|
for table in config['schema']['tables']:
|
|
# Allow per-table strategy override
|
|
table_specific_strategy = table.get('strategy', table_strategy)
|
|
create_table(engine, table, table_specific_strategy)
|
|
|
|
return {"status": "success"} |