How to Restart & Run All code if there is a Key Error during a ! pip install in Google Colab?

Question

In Google Colab, when I install & import this code on each run:

! pip install --upgrade snowflake-connector-python
! pip install pyspark
! pip install snowflake-connector-python asn1crypto==0.24.0
! pip install azure-storage==0.34.3
! pip install azure.storage.common

from pyspark.sql import SparkSession,SQLContext
import pyspark
import snowflake.connector as con
from pyspark.sql import SQLContext
sc = pyspark.SparkContext.getOrCreate()
sqlContext = SQLContext(sc)

I get this error:

KeyError: 'snowflake-connector-python'

I've spoken to Snowflake and they say this is normal and that I just need to click Restart & Run All again.

I'm wondering if there is a way I can automatically Restart & Run All within the code so that I don't need to do this manually? I'm thinking a while loop, but I'm not familiar with error handling during installation.

I'm aware that I can save pip installs to my Drive, but because these files will be acting as shared dashboards, I don't want colleagues to have to do this.

I've read a few things about how to restart the kernel, but not running the whole file again, e.g:

import os
os._exit(00)

Answer 1

This is the best I could do:

• Try to import a package
• If it fails (KeyError, ImportError, etc), pip install it and force a kernel restart.

try:
  import snowflake.connector as con
except (ImportError, KeyError, ModuleNotFoundError):
  !pip install snowflake-connector-python 
  print('Stopping RUNTIME. Colaboratory will restart automatically. Please run again.')
  exit()

Based on multiple answers to Google Colab - How to 'restart runtime' using python code or command line interface?.