Ejemplo n.º 1
0
 def execute(self, context: 'Context') -> Optional[bool]:
     hook = SpannerHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     if not hook.get_database(
         project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id
     ):
         self.log.info(
             "Creating Cloud Spanner database '%s' in project '%s' and instance '%s'",
             self.database_id,
             self.project_id,
             self.instance_id,
         )
         return hook.create_database(
             project_id=self.project_id,
             instance_id=self.instance_id,
             database_id=self.database_id,
             ddl_statements=self.ddl_statements,
         )
     else:
         self.log.info(
             "The database '%s' in project '%s' and instance '%s'"
             " already exists. Nothing to do. Exiting.",
             self.database_id,
             self.project_id,
             self.instance_id,
         )
     return True
Ejemplo n.º 2
0
 def setUp(self):
     with mock.patch(
             'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__',
             new=mock_base_gcp_hook_no_default_project_id,
     ):
         self.spanner_hook_no_default_project_id = SpannerHook(
             gcp_conn_id='test')
Ejemplo n.º 3
0
 def execute(self, context: 'Context') -> None:
     hook = SpannerHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     if not hook.get_database(project_id=self.project_id,
                              instance_id=self.instance_id,
                              database_id=self.database_id):
         raise AirflowException(
             f"The Cloud Spanner database '{self.database_id}' in project '{self.project_id}' "
             f"and instance '{self.instance_id}' is missing. "
             f"Create the database first before you can update it.")
     else:
         SpannerDatabaseLink.persist(
             context=context,
             task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,
         )
         return hook.update_database(
             project_id=self.project_id,
             instance_id=self.instance_id,
             database_id=self.database_id,
             ddl_statements=self.ddl_statements,
             operation_id=self.operation_id,
         )
Ejemplo n.º 4
0
 def execute(self, context: 'Context'):
     hook = SpannerHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     if isinstance(self.query, str):
         queries = [x.strip() for x in self.query.split(';')]
         self.sanitize_queries(queries)
     else:
         queries = self.query
     self.log.info(
         "Executing DML query(-ies) on projects/%s/instances/%s/databases/%s",
         self.project_id,
         self.instance_id,
         self.database_id,
     )
     self.log.info(queries)
     hook.execute_dml(
         project_id=self.project_id,
         instance_id=self.instance_id,
         database_id=self.database_id,
         queries=queries,
     )
     SpannerDatabaseLink.persist(
         context=context,
         task_instance=self,
         instance_id=self.instance_id,
         database_id=self.database_id,
         project_id=self.project_id or hook.project_id,
     )
Ejemplo n.º 5
0
 def execute(self, context: 'Context') -> None:
     hook = SpannerHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     if not hook.get_instance(project_id=self.project_id,
                              instance_id=self.instance_id):
         self.log.info("Creating Cloud Spanner instance '%s'",
                       self.instance_id)
         func = hook.create_instance
     else:
         self.log.info("Updating Cloud Spanner instance '%s'",
                       self.instance_id)
         func = hook.update_instance
     func(
         project_id=self.project_id,
         instance_id=self.instance_id,
         configuration_name=self.configuration_name,
         node_count=self.node_count,
         display_name=self.display_name,
     )
     SpannerInstanceLink.persist(
         context=context,
         task_instance=self,
         instance_id=self.instance_id,
         project_id=self.project_id or hook.project_id,
     )
Ejemplo n.º 6
0
 def execute(self, context):
     hook = SpannerHook(gcp_conn_id=self.gcp_conn_id)
     if hook.get_instance(project_id=self.project_id, instance_id=self.instance_id):
         return hook.delete_instance(project_id=self.project_id,
                                     instance_id=self.instance_id)
     else:
         self.log.info("Instance '%s' does not exist in project '%s'. "
                       "Aborting delete.", self.instance_id, self.project_id)
         return True
Ejemplo n.º 7
0
 def execute(self, context):
     hook = SpannerHook(gcp_conn_id=self.gcp_conn_id)
     if not hook.get_instance(project_id=self.project_id, instance_id=self.instance_id):
         self.log.info("Creating Cloud Spanner instance '%s'", self.instance_id)
         func = hook.create_instance
     else:
         self.log.info("Updating Cloud Spanner instance '%s'", self.instance_id)
         func = hook.update_instance
     func(project_id=self.project_id,
          instance_id=self.instance_id,
          configuration_name=self.configuration_name,
          node_count=self.node_count,
          display_name=self.display_name)
Ejemplo n.º 8
0
 def execute(self, context: 'Context') -> Optional[bool]:
     hook = SpannerHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     if hook.get_instance(project_id=self.project_id, instance_id=self.instance_id):
         return hook.delete_instance(project_id=self.project_id, instance_id=self.instance_id)
     else:
         self.log.info(
             "Instance '%s' does not exist in project '%s'. Aborting delete.",
             self.instance_id,
             self.project_id,
         )
         return True
Ejemplo n.º 9
0
 def execute(self, context):
     hook = SpannerHook(gcp_conn_id=self.gcp_conn_id)
     database = hook.get_database(project_id=self.project_id,
                                  instance_id=self.instance_id,
                                  database_id=self.database_id)
     if not database:
         self.log.info("The Cloud Spanner database was missing: "
                       "'%s' in project '%s' and instance '%s'. Assuming success.",
                       self.database_id, self.project_id, self.instance_id)
         return True
     else:
         return hook.delete_database(project_id=self.project_id,
                                     instance_id=self.instance_id,
                                     database_id=self.database_id)
Ejemplo n.º 10
0
 def execute(self, context):
     hook = SpannerHook(gcp_conn_id=self.gcp_conn_id)
     queries = self.query
     if isinstance(self.query, str):
         queries = [x.strip() for x in self.query.split(';')]
         self.sanitize_queries(queries)
     self.log.info("Executing DML query(-ies) on "
                   "projects/%s/instances/%s/databases/%s",
                   self.project_id, self.instance_id, self.database_id)
     self.log.info(queries)
     hook.execute_dml(project_id=self.project_id,
                      instance_id=self.instance_id,
                      database_id=self.database_id,
                      queries=queries)
Ejemplo n.º 11
0
 def execute(self, context):
     hook = SpannerHook(gcp_conn_id=self.gcp_conn_id)
     if not hook.get_database(project_id=self.project_id,
                              instance_id=self.instance_id,
                              database_id=self.database_id):
         raise AirflowException("The Cloud Spanner database '{}' in project '{}' and "
                                "instance '{}' is missing. Create the database first "
                                "before you can update it.".format(self.database_id,
                                                                   self.project_id,
                                                                   self.instance_id))
     else:
         return hook.update_database(project_id=self.project_id,
                                     instance_id=self.instance_id,
                                     database_id=self.database_id,
                                     ddl_statements=self.ddl_statements,
                                     operation_id=self.operation_id)
Ejemplo n.º 12
0
 def execute(self, context):
     hook = SpannerHook(gcp_conn_id=self.gcp_conn_id)
     if not hook.get_database(project_id=self.project_id,
                              instance_id=self.instance_id,
                              database_id=self.database_id):
         self.log.info("Creating Cloud Spanner database "
                       "'%s' in project '%s' and instance '%s'",
                       self.database_id, self.project_id, self.instance_id)
         return hook.create_database(project_id=self.project_id,
                                     instance_id=self.instance_id,
                                     database_id=self.database_id,
                                     ddl_statements=self.ddl_statements)
     else:
         self.log.info("The database '%s' in project '%s' and instance '%s'"
                       " already exists. Nothing to do. Exiting.",
                       self.database_id, self.project_id, self.instance_id)
     return True