def test_no_pii_table(self):
schema = Schema('pii_db')
table = Table(schema, 'no_pii')
table.add(Column('a'))
table.add(Column('b'))
table.scan(self.explorer.get_connection().cursor())
self.assertFalse(table.has_pii())
def test_partial_pii_table(self):
schema = Schema('pii_db')
table = Table(schema, 'partial_pii')
table.add(Column('a'))
table.add(Column('b'))
table.scan(self.explorer.get_connection().cursor())
self.assertTrue(table.has_pii())
cols = table.get_columns()
self.assertTrue(cols[0].has_pii())
self.assertFalse(cols[1].has_pii())
def test_full_pii_table(self):
schema = Schema('pii_db')
table = Table(schema, 'full_pii')
table.add(Column('name'))
table.add(Column('location'))
table.scan(self.explorer.get_connection().cursor())
self.assertTrue(table.has_pii())
cols = table.get_columns()
self.assertTrue(cols[0].has_pii())
self.assertTrue(cols[1].has_pii())
def setUp(self):
self.explorer = SqliteExplorer("mock_connection")
col1 = Column('c1')
col2 = Column('c2')
col2._pii = [PiiTypes.LOCATION]
schema = Schema('s1')
table = Table(schema, 't1')
table.columns = [col1, col2]
schema = Schema('testSchema')
schema.tables = [table]
self.explorer._schemas = [schema]
def _load_catalog(self):
if self._cache_ts is None or self._cache_ts < datetime.now(
) - timedelta(minutes=10):
with self.get_connection().cursor() as cursor:
cursor.execute(self._catalog_query)
self._schemas = []
row = cursor.fetchone()
current_schema = None
current_table = None
if row is not None:
current_schema = Schema(row[0])
current_table = Table(current_schema, row[1])
while row is not None:
if current_schema.get_name() != row[0]:
current_schema.tables.append(current_table)
self._schemas.append(current_schema)
current_schema = Schema(row[0])
current_table = Table(current_schema, row[1])
elif current_table.get_name() != row[1]:
current_schema.tables.append(current_table)
current_table = Table(current_schema, row[1])
current_table.columns.append(Column(row[2]))
row = cursor.fetchone()
if current_schema is not None and current_table is not None:
current_schema.tables.append(current_table)
self._schemas.append(current_schema)
self._cache_ts = datetime.now()
def get_columns(self, schema_name, table_name):
query = self.pragma_query.format(where_clause="WHERE m.name = ?")
logging.debug(query)
logging.debug(table_name)
result_set = self.get_connection().execute(query, (table_name, ))
columns = []
row = result_set.fetchone()
while row is not None:
columns.append(Column(row[1]))
row = result_set.fetchone()
return columns
def get_tables(self, schema):
query = self.pragma_query.format(where_clause="")
logging.debug(query)
result_set = self.get_connection().execute(query)
tables = []
row = result_set.fetchone()
current_table = None
while row is not None:
if current_table is None:
current_table = Table(schema, row[0])
elif current_table.get_name() != row[0]:
tables.append(current_table)
current_table = Table(schema, row[0])
current_table.add(Column(row[1]))
row = result_set.fetchone()
if current_table is not None:
tables.append(current_table)
return tables
def test_positive_scan_column(self):
col = Column('col')
col.scan('Jonathan Smith')
self.assertTrue(col.has_pii())
def test_negative_scan_column(self):
col = Column('col')
col.scan('abc')
self.assertFalse(col.has_pii())