def test_encrypt_decrypt_large_file(tmpdir, test_files): """ Encrypt and Decrypt a large file """ encryption_material = SnowflakeFileEncryptionMaterial( query_stage_master_key='ztke8tIdVt1zmlQIZm0BMA==', query_id='123873c7-3a66-40c4-ab89-e3722fbccce1', smk_id=3112) # generates N files number_of_files = 1 number_of_lines = 10000 tmp_dir = test_files(tmpdir, number_of_lines, number_of_files) files = glob.glob(os.path.join(tmp_dir, 'file*')) input_file = files[0] encrypted_file = None decrypted_file = None try: (metadata, encrypted_file) = SnowflakeEncryptionUtil.encrypt_file( encryption_material, input_file) decrypted_file = SnowflakeEncryptionUtil.decrypt_file( metadata, encryption_material, encrypted_file) contents = '' cnt = 0 fd = codecs.open(decrypted_file, 'r', encoding=UTF8) for line in fd: contents += line cnt += 1 assert cnt == number_of_lines, "number of lines" finally: os.remove(input_file) if encrypted_file: os.remove(encrypted_file) if decrypted_file: os.remove(decrypted_file)
def put_to_stage(self, file, stream, count): logger.info(f"Uploading {count} rows to stage from {stream} on S3") # Generating key in S3 bucket bucket = self.connection_config["s3_bucket"] s3_key_prefix = self.connection_config.get("s3_key_prefix", "") s3_key = "{}pipelinewise_{}_{}.csv".format( s3_key_prefix, stream, datetime.datetime.now().strftime("%Y%m%d-%H%M%S-%f") ) logger.info("Target S3 bucket: {}, local file: {}, S3 key: {}".format(bucket, file, s3_key)) # Encrypt csv if client side encryption enabled master_key = self.connection_config.get("client_side_encryption_master_key", "") if master_key != "": # Encrypt the file encryption_material = SnowflakeFileEncryptionMaterial( query_stage_master_key=master_key, query_id="", smk_id=0 ) encryption_metadata, encrypted_file = SnowflakeEncryptionUtil.encrypt_file( encryption_material, file ) # Upload to s3 # Send key and iv in the metadata, that will be required to decrypt and upload the encrypted file metadata = {"x-amz-key": encryption_metadata.key, "x-amz-iv": encryption_metadata.iv} self.s3.upload_file(encrypted_file, bucket, s3_key, ExtraArgs={"Metadata": metadata}) # Remove the uploaded encrypted file os.remove(encrypted_file) # Upload to S3 without encrypting else: self.s3.upload_file(file, bucket, s3_key) return s3_key
def put_to_stage(self, file, stream, count): logger.info( "Uploading {} rows to external snowflake stage on S3".format( count)) s3_key_prefix = self.connection_config.get('s3_key_prefix', '') s3_key = "{}pipelinewise_{}_{}.csv".format( s3_key_prefix, stream, datetime.datetime.now().strftime("%Y%m%d-%H%M%S-%f")) # internal staging if not 'aws_access_key_id' in self.connection_config: stage = self.connection_config['stage'] logger.info("Target internal staging: {} ".format(stage)) with self.open_connection() as connection: with connection.cursor(snowflake.connector.DictCursor) as cur: cur.execute("USE SCHEMA {}".format( self.connection_config['default_target_schema'])) logger.info("file: {}, stage: {}, s3_key: {}".format( file, stage, s3_key)) put_sql = "PUT file:///{} @{}/{}".format( file, stage, s3_key) logger.info("SNOWFLAKE - {}".format(put_sql)) cur.execute(put_sql) logger.info("PUT complete") else: # Generating key in S3 bucket bucket = self.connection_config['s3_bucket'] endpoint = self.connection_config['s3_endpoint'] region = self.connection_config['s3_region'] logger.info( "Target S3 bucket: {}, local file: {}, S3 key: {}, endpoint: {}, region: {} " .format(bucket, file, s3_key, endpoint, region)) # Encrypt csv if client side encryption enabled master_key = self.connection_config.get( 'client_side_encryption_master_key', '') if master_key != '': # Encrypt the file encryption_material = SnowflakeFileEncryptionMaterial( query_stage_master_key=master_key, query_id='', smk_id=0) encryption_metadata, encrypted_file = SnowflakeEncryptionUtil.encrypt_file( encryption_material, file) # Upload to s3 # Send key and iv in the metadata, that will be required to decrypt and upload the encrypted file metadata = { 'x-amz-key': encryption_metadata.key, 'x-amz-iv': encryption_metadata.iv } self.s3.upload_file(encrypted_file, bucket, s3_key, ExtraArgs={'Metadata': metadata}) # Remove the uploaded encrypted file os.remove(encrypted_file) # Upload to S3 without encrypting else: self.s3.upload_file(file, bucket, s3_key) return s3_key