Merge pull request #5 from NigelHambly/main

Fix setup to switch DB context before checking table sets

gaiadmpsetup/gaiadmpsetup.py

@@ -1,5 +1,6 @@
 from pyspark.sql.types import *
 from pyspark.sql.session import SparkSession
+from pyspark.sql.utils import AnalysisException
 
 from . import gaiaedr3_pyspark_schema_structures as edr3
 from . import gaiadr3_pyspark_schema_structures as dr3
@@ -18,17 +19,24 @@ def __init__(self):
     @staticmethod
     def setup():
 
-        def tablesExist(expected_tables):
-            actual_tables = [i.name for i in spark.catalog.listTables()]
-            check =  all(item in actual_tables for item in expected_tables)
+        def tablesExist(expected_tables, database):
+
+            check = False
+
+            try:
+
+                spark.sql("use " + database)
+                actual_tables = [i.name for i in spark.catalog.listTables()]
+                check =  all(item in actual_tables for item in expected_tables)
+
+            except AnalysisException: pass
+
             return check
 
         # check EDR3
-        if not tablesExist(edr3.table_dict.keys()):
+        database = "gaiaedr3"
+        if not tablesExist(edr3.table_dict.keys(), database):
 
-            # database name to create
-            database = "gaiaedr3"
-
             # create the database and switch the current SQL database context to it (from default)
             spark.sql("create database if not exists " + database)
             spark.sql("use " + database)
@@ -40,10 +48,10 @@ def tablesExist(expected_tables):
                 reattachParquetFileResourceToSparkContext(table_key, data_store + folder_path, schemas)
 
         # check DR3
-        if not tablesExist(dr3.table_dict.keys()):
+        database = "gaiadr3"
+        if not tablesExist(dr3.table_dict.keys(), database):
 
             # ... similarly for Gaia DR3
-            database = "gaiadr3"
             spark.sql("create database if not exists " + database)
             spark.sql("use " + database)