commit

f/CCR_ETL/ccr_etl_db_init.py

@@ -0,0 +1,89 @@
+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"}