From b6593a69e6a79212df2dbea988b604a94f964c39 Mon Sep 17 00:00:00 2001
From: Tilly Woodfield <22456167+tillywoodfield@users.noreply.github.com>
Date: Wed, 10 Jul 2024 14:33:15 +0300
Subject: [PATCH] Move create schema to its own function

---
 iatidata/__init__.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/iatidata/__init__.py b/iatidata/__init__.py
index 5fc96a0..430f426 100644
--- a/iatidata/__init__.py
+++ b/iatidata/__init__.py
@@ -89,6 +89,20 @@ def create_table(table, sql, **params):
         _create_table(table.lower(), con, sql, **params)
 
 
+def create_schema():
+    if schema:
+        engine = get_engine()
+        with engine.begin() as connection:
+            connection.execute(
+                text(
+                    f"""
+                    DROP schema IF EXISTS {schema} CASCADE;
+                    CREATE schema {schema};
+                    """
+                )
+            )
+
+
 def create_activities_table():
     logger.debug("Creating table: _all_activities")
     engine = get_engine()
@@ -304,6 +318,7 @@ def load_datasets(datasets: Iterable[iatikit.Dataset]) -> None:
 
 def load(processes: int, sample: Optional[int] = None) -> None:
     logger.info("Loading registry data into database")
+    create_schema()
     create_activities_table()
     datasets_sample = islice(iatikit.data().datasets, sample)
     chunked_datasets = distribute(processes, list(datasets_sample))
@@ -314,18 +329,6 @@ def load(processes: int, sample: Optional[int] = None) -> None:
 
 def process_registry() -> None:
     logger.info("Starting process step")
-    if schema:
-        engine = get_engine()
-        with engine.begin() as connection:
-            connection.execute(
-                text(
-                    f"""
-                    DROP schema IF EXISTS {schema} CASCADE;
-                    CREATE schema {schema};
-                    """
-                )
-            )
-
     activity_objects()
     schema_analysis()
     postgres_tables()