added external location to information_schema.tables and recovered from it

Author: dougspadotto

metastore_export_import/01_backup_catalog.py

@@ -11,19 +11,70 @@
 
 # COMMAND ----------
 
+from delta.tables import *
+
+# COMMAND ----------
+
 dbutils.widgets.removeAll()
 dbutils.widgets.text("storageLocation", "/mnt/externallocation", "Storage location for copy")
 dbutils.widgets.text("catalogName", "system", "information_schema catalog")
+dbutils.widgets.dropdown("getExternalLocations", "True", ["True", "False"])
 
 # COMMAND ----------
 
 storage_location = dbutils.widgets.get("storageLocation")
 catalog_name =  dbutils.widgets.get("catalogName")
+get_external_location =  dbutils.widgets.get("getExternalLocations")
 
 table_list = spark.catalog.listTables(f"{catalog_name}.information_schema")
 
 # COMMAND ----------
 
 for table in table_list:
-    df = spark.sql(f"SELECT * FROM {table.catalog}.information_schema.{table.name}")
-    df.write.format("delta").mode("overwrite").save(f"{storage_location}/{table.name}")
+    info_schema_table_df = spark.sql(f"SELECT * FROM {table.catalog}.information_schema.{table.name}")
+    info_schema_table_df.write.format("delta").mode("overwrite").save(f"{storage_location}/{table.name}")
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ##Optional step
+# MAGIC Get table locations from running DESCRIBE EXTENDED on each table on information_schema.tables
+
+# COMMAND ----------
+
+if get_external_location:
+    table_location_columns = ["table_catalog","table_schema","table_name","table_location"]
+    table_location_storage = "external_table_locations"
+    location_list = []
+
+    #Need to filter out Unity Catalog data source that counts as external
+    describe_table_list = spark.read.table(f"{catalog_name}.information_schema.tables").filter("table_type=='EXTERNAL' AND data_source_format <> 'UNITY_CATALOG'")
+    
+    for d_table in describe_table_list.collect():
+        d_location = spark.sql(f"DESCRIBE EXTENDED {d_table.table_catalog}.{d_table.table_schema}.{d_table.table_name}").filter("col_name = 'Location'").select("data_type").head()[0]
+        location_list.append([d_table.table_catalog, d_table.table_schema, d_table.table_name, d_location]) 
+
+    location_df = spark.createDataFrame(data=location_list, schema = table_location_columns)
+    
+    #merge with information_schema.tables and save external locations to storage_sub_directory column (that as of 03/09 only holds Managed table information)
+    table_df = DeltaTable.forPath(spark, f"{storage_location}/tables")
+    #table_df = spark.sql(f"SELECT * FROM {table.catalog}.information_schema.tables")
+    table_df.alias('tables') \
+    .merge(
+        location_df.alias('locations'),
+        'tables.table_catalog = locations.table_catalog and tables.table_schema = locations.table_schema and tables.table_name = locations.table_name'
+        ) \
+    .whenMatchedUpdate(set =
+        {
+          "storage_sub_directory": "locations.table_location"
+        }
+        ) \
+    .execute()
+    
+    display(table_df)
+    #or create a separate table only for this
+    #(location_df
+    #.write
+    #.mode("overwrite")
+    #.format("delta")
+    #.save(f"{storage_location}/{table_location_storage}"))

metastore_export_import/02_recreate_catalog.py

@@ -8,14 +8,14 @@
 # MAGIC 
 # MAGIC Assumptions:
 # MAGIC - The storage credential(s) and external location(s) of the parent external location needs to be created on the target UC beforehand.
-# MAGIC - The external location for all schemas is the same, formed of ```<storage location root>/<schema name>/<table name>```
+# MAGIC - The external location is taken from the overwritten storage_sub_directory column on information_schema.tables (previously it was formed by ```<storage location root>/<schema name>/<table name>```)
 # MAGIC - All tables are Delta
 
 # COMMAND ----------
 
 dbutils.widgets.removeAll()
 dbutils.widgets.text("storageLocation", "/mnt/externallocation", "Storage with source catalog info")
-dbutils.widgets.text("catalogName", "system", "information_schema catalog")
+dbutils.widgets.text("catalogName", "system_backup", "information_schema catalog")
 dbutils.widgets.text("rootExternalStorage", "abfss://[email protected]/root/", "Root of external tables' path")
 
 # COMMAND ----------
@@ -69,7 +69,6 @@ def return_schema(df):
 
 from pyspark.sql.functions import col, when, collect_list, upper
 
-
 #Get only user schemas
 schemas_df = spark.read.format("delta").load(f"{storage_location}/schemata").filter("schema_name<>'information_schema'")
 
@@ -97,7 +96,11 @@ def return_schema(df):
     columns = return_schema(columns_df)
 
     #Create Table
-    spark.sql(f"CREATE OR REPLACE TABLE {catalog_name}.{table.table_schema}.{table.table_name}({columns}) COMMENT '{table.comment}' LOCATION '{root_externalstorage}{table.table_schema}/{table.table_name}'")
+    #Hard-coded path
+    #spark.sql(f"CREATE OR REPLACE TABLE {catalog_name}.{table.table_schema}.{table.table_name}({columns}) COMMENT '{table.comment}' LOCATION '{root_externalstorage}{table.table_schema}/{table.table_name}'")
+    
+    #Extracted path
+    spark.sql(f"CREATE OR REPLACE TABLE {catalog_name}.{table.table_schema}.{table.table_name}({columns}) COMMENT '{table.comment}' LOCATION '{table.storage_sub_directory}'")
     spark.sql(f"ALTER TABLE {catalog_name}.{table.table_schema}.{table.table_name} SET OWNER to `{table.table_owner}`")
 
 # COMMAND ----------