github.com/erda-project/erda-infra@v1.0.10-0.20240327085753-f3a249292aeb/providers/mysqlxorm/examples/main.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"os"
    21  
    22  	"github.com/sirupsen/logrus"
    23  	"xorm.io/xorm"
    24  
    25  	"github.com/erda-project/erda-infra/base/servicehub"
    26  	"github.com/erda-project/erda-infra/providers/mysqlxorm"
    27  )
    28  
    29  type provider struct {
    30  	DB    *xorm.Engine        // autowired
    31  	MySQL mysqlxorm.Interface // autowired
    32  }
    33  
    34  func (p *provider) Init(ctx servicehub.Context) error {
    35  	fmt.Println(p.DB)
    36  	fmt.Println(p.MySQL)
    37  	// do something
    38  	return nil
    39  }
    40  
    41  func (p *provider) Run(ctx context.Context) error {
    42  	/* test db */
    43  	r, err := p.MySQL.DB().QueryString("show tables")
    44  	if err != nil {
    45  		panic(err)
    46  	}
    47  	for i, m := range r {
    48  		for k, v := range m {
    49  			fmt.Println(i, k, v)
    50  		}
    51  	}
    52  
    53  	/* test tx */
    54  	// create table for test
    55  	if err := p.DB.CreateTables(&Table{}); err != nil {
    56  		return err
    57  	}
    58  	defer func() {
    59  		if err := p.DB.DropTables(&Table{}); err != nil {
    60  			logrus.Fatalf("failed to cleanup table, err: %v", err)
    61  		}
    62  	}()
    63  
    64  	// tx
    65  	txSession := p.MySQL.NewSession()
    66  	defer txSession.Close()
    67  	// begin tx
    68  	err = txSession.Begin()
    69  	if err != nil {
    70  		return err
    71  	}
    72  	defer func() {
    73  		if err != nil {
    74  			logrus.Error(err)
    75  
    76  			// check table size before rollback, should be 0 (2 in tx)
    77  			count := p.checkTableRows(0)
    78  			logrus.Printf("count before rollback: %d", count)
    79  
    80  			// rollback
    81  			err = txSession.Rollback()
    82  			if err != nil {
    83  				logrus.Fatalf("failed to rollback, err: %v", err)
    84  			}
    85  		} else { // make insertFailed as success
    86  			p.checkTableRows(2)
    87  		}
    88  	}()
    89  
    90  	// insert success
    91  	err = p.insertSuccess(mysqlxorm.WithSession(txSession))
    92  	if err != nil {
    93  		return err
    94  	}
    95  
    96  	// insert failed
    97  	err = p.insertFailed(mysqlxorm.WithSession(txSession))
    98  	if err != nil {
    99  		return err
   100  	}
   101  
   102  	// tx commit
   103  	err = txSession.Commit()
   104  	if err != nil {
   105  		return err
   106  	}
   107  
   108  	return nil
   109  }
   110  
   111  // Table .
   112  type Table struct {
   113  	ID   uint64 `json:"id" xorm:"pk autoincr"`
   114  	Name string
   115  }
   116  
   117  // TableName .
   118  func (t *Table) TableName() string { return "table" }
   119  
   120  func (p *provider) insertSuccess(opts ...mysqlxorm.SessionOption) error {
   121  	s := p.MySQL.NewSession(opts...)
   122  	defer s.Close()
   123  	_, err := s.InsertOne(&Table{Name: "n1"})
   124  	return err
   125  }
   126  
   127  func (p *provider) insertFailed(opts ...mysqlxorm.SessionOption) error {
   128  	s := p.MySQL.NewSession(opts...)
   129  	defer s.Close()
   130  	_, err := s.InsertOne(&Table{Name: "n2"})
   131  	if err != nil {
   132  		return err
   133  	}
   134  	// force err
   135  	return fmt.Errorf("fake error by func: insertFailed")
   136  }
   137  
   138  func (p *provider) checkTableRows(expectRows int64) int64 {
   139  	count, _ := p.DB.Count(&Table{})
   140  	if count != expectRows {
   141  		logrus.Errorf("expectRows: %d, actualRows: %d", expectRows, count)
   142  	} else {
   143  		logrus.Infof("expectRows: %d, actualRows: %d", expectRows, count)
   144  	}
   145  	return count
   146  }
   147  
   148  func init() {
   149  	servicehub.Register("example", &servicehub.Spec{
   150  		Services:     []string{"example"},
   151  		Dependencies: []string{"mysql-xorm"},
   152  		Description:  "example",
   153  		Creator: func() servicehub.Provider {
   154  			return &provider{}
   155  		},
   156  	})
   157  }
   158  
   159  func main() {
   160  	hub := servicehub.New()
   161  	hub.Run("examples", "", os.Args...)
   162  }