def test_connect_args():
"""
Tests connect string
Snowflake connect string supports accout name as a replacement of
host:port
"""
from sqlalchemy import create_engine
engine = create_engine(
'snowflake://{user}:{password}@{account}/{database}/{schema}'.format(
user=CONNECTION_PARAMETERS2['user'],
password=CONNECTION_PARAMETERS2['password'],
account=CONNECTION_PARAMETERS2['account'],
database=CONNECTION_PARAMETERS2['database'],
schema=CONNECTION_PARAMETERS2['schema'],
)
)
try:
results = engine.execute('select current_version()').fetchone()
assert results is not None
finally:
engine.dispose()
engine = create_engine(
'snowflake://{user}:{password}@{account}/'.format(
user=CONNECTION_PARAMETERS2['user'],
password=CONNECTION_PARAMETERS2['password'],
account=CONNECTION_PARAMETERS2['account'],
)
)
try:
results = engine.execute('select current_version()').fetchone()
assert results is not None
finally:
engine.dispose()
engine = create_engine(URL(
user=CONNECTION_PARAMETERS2['user'],
password=CONNECTION_PARAMETERS2['password'],
account=CONNECTION_PARAMETERS2['account'],
)
)
try:
results = engine.execute('select current_version()').fetchone()
assert results is not None
finally:
engine.dispose()
engine = create_engine(URL(
user=CONNECTION_PARAMETERS2['user'],
password=CONNECTION_PARAMETERS2['password'],
account=CONNECTION_PARAMETERS2['account'],
warehouse='testwh'
)
)
try:
results = engine.execute('select current_version()').fetchone()
assert results is not None
finally:
engine.dispose()
def __init__(self, config: Dict = None) -> None:
if not config:
config = {
"user": os.getenv("PERMISSION_BOT_USER"),
"password": os.getenv("PERMISSION_BOT_PASSWORD"),
"account": os.getenv("PERMISSION_BOT_ACCOUNT"),
"database": os.getenv("PERMISSION_BOT_DATABASE"),
"role": os.getenv("PERMISSION_BOT_ROLE"),
"warehouse": os.getenv("PERMISSION_BOT_WAREHOUSE"),
"oauth_token": os.getenv("PERMISSION_BOT_OAUTH_TOKEN"),
}
if config["oauth_token"] is not None:
self.engine = sqlalchemy.create_engine(
URL(
user=config["user"],
account=config["account"],
authenticator="oauth",
token=config["oauth_token"],
warehouse=config["warehouse"],
)
)
else:
self.engine = sqlalchemy.create_engine(
URL(
user=config["user"],
password=config["password"],
account=config["account"],
database=config["database"],
role=config["role"],
warehouse=config["warehouse"],
# Enable the insecure_mode if you get OCSP errors while testing
# insecure_mode=True,
)
)
def test_url():
assert URL(account='testaccount', user='******',
password='******', warehouse='testwh') == \
"snowflake://*****:*****@testaccount/?warehouse=testwh"
assert URL(account='testaccount', user='******',
password='******') == "snowflake://*****:*****@testaccount/"
assert URL(account='testaccount', user='******',
password='******', database='testdb') == \
"snowflake://*****:*****@testaccount/testdb"
assert URL(account='testaccount', user='******',
password='******', database='testdb', schema='testschema') == \
"snowflake://*****:*****@testaccount/testdb/testschema"
assert URL(account='testaccount', user='******',
password='******', database='testdb', schema='testschema',
warehouse='testwh') == \
"snowflake://*****:*****@testaccount/testdb/testschema?warehouse" \
"=testwh"
assert URL(host='snowflake.reg.local', account='testaccount', user='******',
password='******', database='testdb', schema='testschema') == \
"snowflake://*****:*****@snowflake.reg.local:443/testdb" \
"/testschema?account=testaccount"
assert URL(
host='testaccount.snowflakecomputing.com',
user='******',
account='testaccount',
password='******',
region='eu-central-1') == (
'snowflake://*****:*****@testaccount.snowflakecomputing.com:443'
'/?account=testaccount®ion=eu-central-1')
assert URL(host='testaccount.eu-central-1.snowflakecomputing.com',
user='******',
account='testaccount',
password='******') == (
'snowflake://*****:*****@testaccount.eu-central-1'
'.snowflakecomputing.com:443/?account=testaccount')
# empty password should be acceptable in URL utility. The validation will
# happen in Python connector anyway.
assert URL(host='testaccount.eu-central-1.snowflakecomputing.com',
user='******',
account='testaccount') == (
'snowflake://admin:@testaccount.eu-central-1'
'.snowflakecomputing.com:443/?account=testaccount')
# authenticator=externalbrowser doesn't require a password.
assert URL(host='testaccount.eu-central-1.snowflakecomputing.com',
user='******',
account='testaccount',
authenticator='externalbrowser') == (
'snowflake://admin:@testaccount.eu-central-1'
'.snowflakecomputing.com:443/?account=testaccount'
'&authenticator=externalbrowser')
def profile_to_db(dp_df, user_id):
db_config = Config()
engine = create_engine(
URL(account=db_config.SNOWFLAKE_ACCOUNT,
user=db_config.SNOWFLAKE_USER,
password=db_config.SNOWFLAKE_PASSWORD,
database='ADQ',
schema='PUBLIC',
warehouse='COMPUTE_WH',
role='SYSADMIN'))
session = sessionmaker(bind=engine)()
try:
# Creating a new data profile row in the data_profile table
newProfile = DataProfile(user_id=user_id)
# Note, newProfile dp_id isn't created until commited since SF creates the value on insert
session.add(newProfile)
session.commit()
# Adding data profile rows to column_profile table with a new column of dp_id
# so rows are related by their new data profile id
dp_df["dp_id"] = newProfile.dp_id
dp_df.to_sql('column_profile', engine, if_exists='append', index=False)
print("Data profile uploaded sucessfully.")
finally:
# connection.close()
engine.dispose()
def _authenticate_with_key_pair(self, rsa_key_path: str,
private_passphrase: str,
**kwargs) -> tuple:
with open(rsa_key_path, "rb") as key:
if private_passphrase:
private_passphrase = private_passphrase.encode()
p_key = serialization.load_pem_private_key(
key.read(),
password=private_passphrase,
backend=default_backend())
pkb = p_key.private_bytes(
encoding=serialization.Encoding.DER,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption())
client = snowflake.connector.connect(account=self.account_name,
application=APP_NAME,
validate_default_parameters=True,
protocol='https',
private_key=pkb,
**kwargs)
engine = create_engine(
URL(account=self.account_name,
application=APP_NAME,
validate_default_parameters=True,
protocol='https',
private_key=pkb,
**kwargs))
return client, engine
def readData():
snow_user = os.getenv("snow_user")
snow_pwd = os.getenv("snow_pwd")
print(snow_user)
engine = create_engine(
URL(
account='xb33235.eu-west-1',
user=snow_user,
password=snow_pwd,
database='ANALYTICS',
schema='ANALYTICS',
warehouse='reporting',
role='ANALYST',
))
cs = engine.connect()
# Get the total ID we have
sql = f"SELECT * from ANALYTICS_BASE.EXPERIAN_PD_RAW"
df1 = pd.read_sql_query(sql, cs)
mapping_csv = os.path.dirname(
os.path.dirname(
os.path.dirname(
os.path.dirname(os.path.dirname(
os.path.realpath(__file__)))))) + '/name_mapping.csv'
df2 = FinancialSheet(df1, mapping_csv, max_period=1).data
df2.fillna(value=pd.np.nan, inplace=True)
for col in df2:
try:
df2[col] = df2[col].astype(float)
except:
continue
return df2
def get_sf_data(query: str,
role: str = SF_ROLE,
warehouse: str = SF_WAREHOUSE,
db: str = SF_DB,
schema: str = SF_SCHEMA) -> pd.DataFrame:
"""
Use this to run a select statment and export that data to a pandas data frame.
"""
url = URL(
account=SF_ACCT,
user=SF_USER,
password=SF_PASS,
database=db,
schema=schema,
warehouse=warehouse,
role=role,
)
engine = create_engine(url)
connection = engine.connect()
df = pd.read_sql_query(query, connection)
connection.close()
engine.dispose()
return df
def get_engine_with_numpy(db_parameters,
user=None,
password=None,
account=None):
"""
Creates a connection using the parameters defined in JDBC connect string
"""
from sqlalchemy import create_engine
from snowflake.sqlalchemy import URL
if user is not None:
db_parameters['user'] = user
if password is not None:
db_parameters['password'] = password
if account is not None:
db_parameters['account'] = account
from sqlalchemy.pool import NullPool
engine = create_engine(URL(
user=db_parameters['user'],
password=db_parameters['password'],
host=db_parameters['host'],
port=db_parameters['port'],
database=db_parameters['database'],
schema=db_parameters['schema'],
account=db_parameters['account'],
protocol=db_parameters['protocol'],
numpy=True,
),
poolclass=NullPool)
return engine
def multiple_snowflake_queries(queries: list,
role: str = SF_ROLE,
warehouse: str = SF_WAREHOUSE,
db: str = SF_DB,
schema: str = SF_SCHEMA) -> pd.DataFrame:
"""
Use this to run a select statment and export that data to a pandas data frame.
"""
url = URL(
account=SF_ACCT,
user=SF_USER,
password=SF_PASS,
database=db,
schema=schema,
warehouse=warehouse,
role=role,
)
engine = create_engine(url)
connection = engine.connect()
out = {}
for q in queries:
df = pd.read_sql_query(q, connection)
out.update({q: df})
connection.close()
engine.dispose()
return out
def work():
engine1 = create_engine(
URL(
account='fh25363.us-east-1.privatelink',
user='******',
password='******',
database='DEVELOPMENT_POC',
schema='MAIN',
warehouse='DEVELOPMENT_COMPUTE_WH',
role='DEVELOPMENT_RW',
))
connection = engine1.connect()
try:
#connection = engine1.connect()
results = connection.execute('select current_version()').fetchone()
results = connection.execute('select * from cars')
returnArray = []
for result in results:
colname = ['reg', 'make', 'model', 'price', 'totalvotes']
item = {}
for i, colName in enumerate(colname):
value = result[i]
item[colName] = value
returnArray.append(item)
print(returnArray)
#print(results[0])
finally:
connection.close()
engine1.dispose()
def get_engine(user=None, password=None, account=None):
"""
Creates a connection using the parameters defined in JDBC connect string
"""
ret = get_db_parameters()
if user is not None:
ret['user'] = user
if password is not None:
ret['password'] = password
if account is not None:
ret['account'] = account
from sqlalchemy.pool import NullPool
engine = create_engine(URL(user=ret['user'],
password=ret['password'],
host=ret['host'],
port=ret['port'],
database=ret['database'],
schema=TEST_SCHEMA,
account=ret['account'],
protocol=ret['protocol']),
poolclass=NullPool)
return engine, ret
def manage_database(database: str,
action: str,
target_role: Optional[str] = None):
database = database.replace("-", "_")
if action == "create":
stmts = [
f'CREATE OR REPLACE DATABASE "{database}" CLONE "ANALYTICS_PRODUCTION"',
f'GRANT OWNERSHIP ON DATABASE "{database}" TO ROLE {target_role} REVOKE CURRENT GRANTS',
f'GRANT OWNERSHIP ON ALL SCHEMAS IN DATABASE "{database}" TO ROLE {target_role} REVOKE CURRENT GRANTS',
f'GRANT USAGE ON DATABASE "{database}" TO ROLE PUBLIC',
]
elif action == "drop":
stmts = [f'DROP DATABASE IF EXISTS "{database}"']
else:
stmts = [] # do nothing
engine = create_engine(
URL(
account=SNOWFLAKE_ACCOUNT,
user=SNOWFLAKE_USERNAME,
password=SNOWFLAKE_PASSWORD,
role=SNOWFLAKE_ROLE,
database=database,
))
with engine.begin() as tx:
for stmt in stmts:
tx.execute(stmt)
if action == "create":
change_objects_ownership(engine, database, target_role)
change_functions_ownership(engine, database, target_role)
def test_special_schema_character(db_parameters):
"""Make sure we decode special characters correctly"""
# Constants
database = "a/b/c" # "'/'.join([choice(ascii_lowercase) for _ in range(3)])
schema = "d/e/f" # '/'.join([choice(ascii_lowercase) for _ in range(3)])
# Setup
options = dict(**db_parameters)
conn = connect(**options)
conn.cursor().execute(
"CREATE OR REPLACE DATABASE \"{0}\"".format(database))
conn.cursor().execute("CREATE OR REPLACE SCHEMA \"{0}\"".format(schema))
conn.close()
# Test
options.update({
'database': '"' + database + '"',
'schema': '"' + schema + '"'
})
sf_conn = connect(**options)
sf_connection = [
res for res in sf_conn.cursor().execute("select current_database(), "
"current_schema();")
]
sa_conn = create_engine(URL(**options)).connect()
sa_connection = [
res for res in sa_conn.execute("select current_database(), "
"current_schema();")
]
sa_conn.close()
sf_conn.close()
# Teardown
conn = connect(**options)
conn.cursor().execute("DROP DATABASE IF EXISTS \"{0}\"".format(database))
conn.close()
assert [(database, schema)] == sf_connection == sa_connection
def _conn_params_to_sqlalchemy_uri(self, conn_params: Dict) -> str:
return URL(
**{
k: v
for k, v in conn_params.items() if v and k not in
['session_parameters', 'insecure_mode', 'private_key']
})
def _get_engine_with_columm_metadata_cache(db_parameters,
user=None,
password=None,
account=None):
"""
Creates a connection with column metadata cache
"""
if user is not None:
db_parameters['user'] = user
if password is not None:
db_parameters['password'] = password
if account is not None:
db_parameters['account'] = account
from sqlalchemy.pool import NullPool
from sqlalchemy import create_engine
from snowflake.sqlalchemy import URL
engine = create_engine(URL(
user=db_parameters['user'],
password=db_parameters['password'],
host=db_parameters['host'],
port=db_parameters['port'],
database=db_parameters['database'],
schema=db_parameters['schema'],
account=db_parameters['account'],
protocol=db_parameters['protocol'],
cache_column_metadata=True,
),
poolclass=NullPool)
return engine
def get_engine(self):
self.engine = create_engine(
URL(account=self.account,
user=self.user,
warehouse=self.warehouse,
database=self.database),
connect_args={'authenticator': 'externalbrowser'})
return self.engine
def create_engine(self, settings):
from snowflake.sqlalchemy import URL
url_params = dict()
for param in db_parameters:
if param in settings:
url_params[param] = settings.pop(param)
return create_engine(URL(**url_params), **settings)
def _get_snowflake_engine(self):
return create_engine(URL(
user=Config.snowflake_account['user'],
password=Config.snowflake_account['password'],
account=Config.snowflake_account['account'],
database=Config.snowflake_account['database'],
schema=Config.snowflake_account['schema'],
warehouse=Config.snowflake_account['warehouse']
), echo=True
)
def _authenticate_with_user_credentials(self, **kwargs) -> tuple:
client = snowflake.connector.connect(account=self.account_name,
application=APP_NAME,
validate_default_parameters=True,
protocol='https',
**kwargs)
engine = create_engine(
URL(account=self.account_name, application=APP_NAME, **kwargs))
return client, engine
def __init__(self, user, password, schema, account, warehouse, database,
role):
self.engine = create_engine(
URL(account=account,
user=user,
password=password,
role=role,
warehouse=warehouse,
database=database,
schema=schema))
def createEngine(SNOW_USER, SNOW_PASS, SNOW_ACCOUNT, SNOW_DB, SNOW_SCHEMA):
engine = create_engine(
URL(
user=SNOW_USER,
password=SNOW_PASS,
account=SNOW_ACCOUNT,
database=SNOW_DB,
schema=SNOW_SCHEMA,
role='sysadmin',
))
return engine
def create_connection(engine_name):
engine = create_engine(
URL(user=engine_name['user'],
password=engine_name['password'],
account=engine_name['account'],
warehouse=engine_name['warehouse'],
database=engine_name['database'],
schema=engine_name['schema'],
role=engine_name['role'],
numpy=engine_name['numpy']))
return engine.connect()
def create_snowflake_engine(): # -> Connection:
return create_engine(
URL(
account=config.snowflakeAccount,
user=config.snowflakeUsername,
password=config.snowflakePassword,
role=config.snowflakeRole,
warehouse=config.snowflakeWarehouse,
database=config.snowflakeDatabase,
schema=config.snowflakeSchema,
))
def generate_engine(self):
self.engine = create_engine(
URL(
account=self.credentials.credentials.get("account"),
user=self.credentials.credentials.get("user"),
password=self.credentials.credentials.get("password"),
role=self.credentials.credentials.get("role"),
warehouse=self.credentials.credentials.get("warehouse"),
database=self.credentials.credentials.get("database"),
schema=self.credentials.credentials.get("schema"),
))
def executeScriptsFromFile(filename):
# Open and read the file as a single buffer
f = open(filename, 'r')
sqlFile = f.read()
f.close()
# all SQL commands (split on ';')
#sqlCommands = sqlFile.split(';')
print("Connecting to DB - Snowflake")
try:
engine = create_engine(
URL(account='',
region='',
user=user,
database='',
schema='',
warehouse='',
authenticator='externalbrowser',
role=''))
except Exception as e:
print(e)
engine = create_engine(
URL(account='',
region='',
user=user,
database='',
schema='',
warehouse='',
authenticator='externalbrowser',
role=''))
try:
connection = engine.connect()
print("Executing SQL")
results = connection.execute(sqlFile).fetchall()
print("Got data")
print(results)
return results
except Exception, e:
print("Can't query: {0}").format(e)
def snowflake_connection ( account, username, password, database, schema, warehouse ) :
engine = create_engine(URL(
account = account,
user = username,
password = password,
database = database,
schema = schema,
warehouse = warehouse,
numpy=True
))
return engine.connect()
def create_engine(cls):
return db.create_engine(
URL(
account=cls._account,
user=cls._user,
password=cls._password,
database=cls._database,
schema=cls._schema,
warehouse=cls._warehouse,
role=cls._role,
)
)
def _set_connection(self):
if self.auth_type == 'basic' and self.connector_type == 'snowflake_connector':
self.conn = snowflake.connector.connect(
user=self.credentials['user'],
password=self.credentials['password'],
account=self.credentials['account'])
if self.auth_type == 'basic' and self.connector_type == 'snowflake_sqlalchemy':
engine = create_engine(
URL(account=self.credentials['account'],
user=self.credentials['user'],
password=self.credentials['password']))
self.conn = engine.connect()
def get_engine(self):
json_data = self.instance.extra_json
self.engine = create_engine(
URL(
user=self.instance.login,
password=self.instance.password,
account=json_data["account"],
region=json_data["region"],
database=self.instance.db,
warehouse=json_data["warehouse"],
role=json_data["role"],
))
def test_url():
assert URL(account='testaccount', user='******',
password='******', warehouse='testwh') == \
"snowflake://*****:*****@testaccount/?warehouse=testwh"
assert URL(account='testaccount', user='******',
password='******') == "snowflake://*****:*****@testaccount/"
assert URL(account='testaccount', user='******',
password='******', database='testdb') == \
"snowflake://*****:*****@testaccount/testdb"
assert URL(account='testaccount', user='******',
password='******', database='testdb', schema='testschema') == \
"snowflake://*****:*****@testaccount/testdb/testschema"
assert URL(account='testaccount', user='******',
password='******', database='testdb', schema='testschema',
warehouse='testwh') == \
"snowflake://*****:*****@testaccount/testdb/testschema?warehouse" \
"=testwh"
assert URL(host='snowflake.reg.local', account='testaccount', user='******',
password='******', database='testdb', schema='testschema') == \
"snowflake://*****:*****@snowflake.reg.local:443/testdb" \
"/testschema?account=testaccount"