Add an external id to provider accounts so that they may be joined against external data sources

Author: gsoltis

introspector/aws/__init__.py

@@ -30,8 +30,9 @@ def get_boto_session() -> boto.Session:
   return session
 
 
-def _create_provider_and_credential(db: Session, proxy: Proxy,
-                                    identity) -> ProviderAccount:
+def _create_provider_and_credential(
+    db: Session, proxy: Proxy, identity,
+    external_id: Optional[int]) -> ProviderAccount:
   account_id = identity['Account']
   org = proxy.service('organizations')
   try:
@@ -43,7 +44,9 @@ def _create_provider_and_credential(db: Session, proxy: Proxy,
       org_id = f'OrgDummy:{account_id}'
     else:
       raise
-  provider = ProviderAccount(provider='aws', name=org_id)
+  provider = ProviderAccount(provider='aws',
+                             name=org_id,
+                             external_id=external_id)
   db.add(provider)
   db.flush()
   _require_credential(db, provider.id, identity)
@@ -87,18 +90,20 @@ def walk_graph(org, graph) -> Generator[Tuple[str, str, Dict], None, None]:
 
 
 def build_aws_import_job(db: Session, session: boto.Session,
-                         confirm: ConfirmAcct) -> ImportJob:
+                         confirm: ConfirmAcct,
+                         external_id: Optional[int]) -> ImportJob:
   proxy = Proxy.build(session)
   sts = session.create_client('sts')
   identity = sts.get_caller_identity()
-  provider = _get_or_create_provider(db, proxy, identity, confirm)
+  provider = _get_or_create_provider(db, proxy, identity, confirm, external_id)
   desc = _build_import_job_desc(proxy, identity)
   org_id = desc['aws_org']['Id']
   return ImportJob.create(provider, desc, org_id)
 
 
 def _get_or_create_provider(db: Session, proxy: Proxy, identity: Dict,
-                            confirm: ConfirmAcct) -> ProviderAccount:
+                            confirm: ConfirmAcct,
+                            external_id: Optional[int]) -> ProviderAccount:
   org = proxy.service('organizations')
   try:
     org_resp = org.get('describe_organization')['Organization']
@@ -118,7 +123,7 @@ def _get_or_create_provider(db: Session, proxy: Proxy, identity: Dict,
   add = confirm(identity)
   if not add:
     raise GFError('User cancelled')
-  return _create_provider_and_credential(db, proxy, identity)
+  return _create_provider_and_credential(db, proxy, identity, external_id)
 
 
 def _build_import_job_desc(proxy: Proxy, identity: Dict) -> Dict:

introspector/cli/account/aws.py

@@ -58,9 +58,18 @@ def cmd():
               is_flag=True,
               required=False,
               help='Set this flag to import a govcloud account')
+@click.option(
+    '-e',
+    '--external_id',
+    default=None,
+    required=False,
+    type=int,
+    help=
+    '''Optional identifier to include in the provider_account row to facilitate joining against other data sources'''
+)
 @click.option('--dry-run', 'dry_run', default=False, hidden=True, is_flag=True)
-def import_aws_cmd(debug: bool, force: bool, dry_run: bool,
-                   service: Optional[str], gov_cloud: bool):
+def import_aws_cmd(debug: bool, force: bool, external_id: Optional[int],
+                   dry_run: bool, service: Optional[str], gov_cloud: bool):
   partition = 'aws-us-gov' if gov_cloud else 'aws'
   os.environ[
       'AWS_DEFAULT_REGION'] = 'us-gov-east-1' if gov_cloud else 'us-east-2'
@@ -77,7 +86,7 @@ def _confirm(identity: Dict) -> bool:
           default='yes')
 
     confirm = _confirm
-  import_job = build_aws_import_job(db, boto, confirm)
+  import_job = build_aws_import_job(db, boto, confirm, external_id)
   db.add(import_job)
   db.flush()
   region_cache = RegionCache(boto, partition)

introspector/models/provider_account.py

@@ -14,18 +14,9 @@ class ProviderAccount(Base):
   name = Column(String(256), comment='Provider name.')
   provider = Column(Enum('aws', 'gcp', 'azure', name='provider'),
                     comment='Provider enum.')
-
-  @classmethod
-  def get_or_create(cls, session: Session, account_info):
-    provider = account_info['provider']
-    account_id = account_info['account_id']
-    provider_account = session.query(ProviderAccount).filter(
-        ProviderAccount.provider == provider,
-        ProviderAccount.name == account_id).one_or_none()
-    if provider_account is None:
-      provider_account = ProviderAccount(provider=provider, name=account_id)
-      session.add(provider_account)
-    return provider_account
+  external_id = Column(
+      Integer,
+      comment='Optional external id for joining against external data sources')
 
   @classmethod
   def all(cls,

migrations/introspector/20210425202304_add_external_id.sql

@@ -0,0 +1,5 @@
+-- migrate:up
+ALTER TABLE provider_account ADD COLUMN external_id INTEGER;
+
+-- migrate:down
+ALTER TABLE provider_account DROP COLUMN external_id;