github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/test/testcases/database_test.go (about)

     1  //go:build L0
     2  
     3  package testcases
     4  
     5  import (
     6  	"fmt"
     7  	"log"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/milvus-io/milvus-sdk-go/v2/entity"
    12  	"github.com/milvus-io/milvus-sdk-go/v2/test/base"
    13  	"github.com/stretchr/testify/assert"
    14  
    15  	"github.com/milvus-io/milvus-sdk-go/v2/client"
    16  	"github.com/milvus-io/milvus-sdk-go/v2/test/common"
    17  )
    18  
    19  // teardownTest
    20  func teardownTest(t *testing.T) func(t *testing.T) {
    21  	log.Println("setup test func")
    22  	return func(t *testing.T) {
    23  		log.Println("teardown func drop all non-default db")
    24  		// drop all db
    25  		ctx := createContext(t, time.Second*common.DefaultTimeout)
    26  		mc := createMilvusClient(ctx, t)
    27  		dbs, _ := mc.ListDatabases(ctx)
    28  		for _, db := range dbs {
    29  			if db.Name != common.DefaultDb {
    30  				_ = mc.UsingDatabase(ctx, db.Name)
    31  				collections, _ := mc.ListCollections(ctx)
    32  				for _, coll := range collections {
    33  					_ = mc.DropCollection(ctx, coll.Name)
    34  				}
    35  				_ = mc.DropDatabase(ctx, db.Name)
    36  			}
    37  		}
    38  	}
    39  }
    40  
    41  func TestDatabase(t *testing.T) {
    42  	teardownSuite := teardownTest(t)
    43  	defer teardownSuite(t)
    44  
    45  	ctx := createContext(t, time.Second*common.DefaultTimeout)
    46  	clientDefault := createMilvusClient(ctx, t)
    47  	tmpNb := 100
    48  	// create db1
    49  	dbName1 := common.GenRandomString(4)
    50  	err := clientDefault.CreateDatabase(ctx, dbName1)
    51  	common.CheckErr(t, err, true)
    52  
    53  	// list db and verify db1 in dbs
    54  	dbs, errList := clientDefault.ListDatabases(ctx)
    55  	common.CheckErr(t, errList, true)
    56  	common.CheckContainsDb(t, dbs, dbName1)
    57  
    58  	// new client with db1 -> create collections
    59  	clientDB1 := createMilvusClient(ctx, t, client.Config{
    60  		DBName: dbName1,
    61  	})
    62  	db1Col1, _ := createCollectionAllFields(ctx, t, clientDB1, tmpNb, 0)
    63  	db1Col2, _ := createCollectionAllFields(ctx, t, clientDB1, tmpNb, 0)
    64  	collections, errListCollections := clientDB1.ListCollections(ctx)
    65  	common.CheckErr(t, errListCollections, true)
    66  	common.CheckContainsCollection(t, collections, db1Col1)
    67  	common.CheckContainsCollection(t, collections, db1Col2)
    68  
    69  	// create db2
    70  	dbName2 := common.GenRandomString(4)
    71  	err = clientDefault.CreateDatabase(ctx, dbName2)
    72  	common.CheckErr(t, err, true)
    73  	dbs, err = clientDefault.ListDatabases(ctx)
    74  	common.CheckErr(t, err, true)
    75  	common.CheckContainsDb(t, dbs, dbName2)
    76  
    77  	// using db2 -> create collection -> drop collection
    78  	clientDefault.UsingDatabase(ctx, dbName2)
    79  	db2Col1, _ := createCollectionAllFields(ctx, t, clientDefault, tmpNb, 0)
    80  	err = clientDefault.DropCollection(ctx, db2Col1)
    81  	common.CheckErr(t, err, true)
    82  
    83  	// using empty db -> drop db2
    84  	clientDefault.UsingDatabase(ctx, "")
    85  	err = clientDefault.DropDatabase(ctx, dbName2)
    86  	common.CheckErr(t, err, true)
    87  
    88  	// list db and verify db drop success
    89  	dbs, err = clientDefault.ListDatabases(ctx)
    90  	common.CheckErr(t, err, true)
    91  	common.CheckNotContainsDb(t, dbs, dbName2)
    92  
    93  	// drop db1 which has some collections
    94  	err = clientDB1.DropDatabase(ctx, dbName1)
    95  	common.CheckErr(t, err, false, "must drop all collections before drop database")
    96  
    97  	// drop all db1's collections -> drop db1
    98  	err = clientDB1.DropCollection(ctx, db1Col1)
    99  	common.CheckErr(t, err, true)
   100  
   101  	err = clientDB1.DropCollection(ctx, db1Col2)
   102  	common.CheckErr(t, err, true)
   103  
   104  	err = clientDB1.DropDatabase(ctx, dbName1)
   105  	common.CheckErr(t, err, true)
   106  
   107  	// drop default db
   108  	err = clientDefault.DropDatabase(ctx, common.DefaultDb)
   109  	common.CheckErr(t, err, false, "can not drop default database")
   110  
   111  	dbs, err = clientDefault.ListDatabases(ctx)
   112  	common.CheckErr(t, err, true)
   113  	common.CheckContainsDb(t, dbs, common.DefaultDb)
   114  }
   115  
   116  // test create with invalid db name
   117  func TestCreateDb(t *testing.T) {
   118  	teardownSuite := teardownTest(t)
   119  	defer teardownSuite(t)
   120  
   121  	// create db
   122  	ctx := createContext(t, time.Second*common.DefaultTimeout)
   123  	mc := createMilvusClient(ctx, t)
   124  	dbName := common.GenRandomString(4)
   125  	err := mc.CreateDatabase(ctx, dbName)
   126  	common.CheckErr(t, err, true)
   127  
   128  	// create existed db
   129  	err = mc.CreateDatabase(ctx, dbName)
   130  	common.CheckErr(t, err, false, fmt.Sprintf("database already exist: %s", dbName))
   131  
   132  	// create default db
   133  	err = mc.CreateDatabase(ctx, common.DefaultDb)
   134  	common.CheckErr(t, err, false, fmt.Sprintf("database already exist: %s", common.DefaultDb))
   135  
   136  	emptyErr := mc.CreateDatabase(ctx, "")
   137  	common.CheckErr(t, emptyErr, false, "database name couldn't be empty")
   138  }
   139  
   140  // test drop db
   141  func TestDropDb(t *testing.T) {
   142  	teardownSuite := teardownTest(t)
   143  	defer teardownSuite(t)
   144  
   145  	// create collection in default db
   146  	ctx := createContext(t, time.Second*common.DefaultTimeout)
   147  	mc := createMilvusClient(ctx, t)
   148  	collName := createDefaultCollection(ctx, t, mc, true, common.DefaultShards)
   149  	collections, _ := mc.ListCollections(ctx)
   150  	common.CheckContainsCollection(t, collections, collName)
   151  
   152  	// create db
   153  	dbName := common.GenRandomString(4)
   154  	err := mc.CreateDatabase(ctx, dbName)
   155  	common.CheckErr(t, err, true)
   156  
   157  	// using db and drop the db
   158  	err = mc.UsingDatabase(ctx, dbName)
   159  	common.CheckErr(t, err, true)
   160  	err = mc.DropDatabase(ctx, dbName)
   161  	common.CheckErr(t, err, true)
   162  
   163  	// verify current db
   164  	_, err = mc.ListCollections(ctx)
   165  	common.CheckErr(t, err, false, fmt.Sprintf("database not found[database=%s]", dbName))
   166  
   167  	// using default db and verify collections
   168  	mc.UsingDatabase(ctx, common.DefaultDb)
   169  	collections, _ = mc.ListCollections(ctx)
   170  	common.CheckContainsCollection(t, collections, collName)
   171  
   172  	// drop not existed db
   173  	err = mc.DropDatabase(ctx, common.GenRandomString(4))
   174  	common.CheckErr(t, err, true)
   175  
   176  	// drop empty db
   177  	err = mc.DropDatabase(ctx, "")
   178  	common.CheckErr(t, err, false, "database name couldn't be empty")
   179  
   180  	// drop default db
   181  	err = mc.DropDatabase(ctx, common.DefaultDb)
   182  	common.CheckErr(t, err, false, "can not drop default database")
   183  }
   184  
   185  // test using db
   186  func TestUsingDb(t *testing.T) {
   187  	teardownSuite := teardownTest(t)
   188  	defer teardownSuite(t)
   189  
   190  	// create collection in default db
   191  	ctx := createContext(t, time.Second*common.DefaultTimeout)
   192  	mc := createMilvusClient(ctx, t)
   193  	collName := createDefaultCollection(ctx, t, mc, true, common.DefaultShards)
   194  	collections, _ := mc.ListCollections(ctx)
   195  	common.CheckContainsCollection(t, collections, collName)
   196  
   197  	// using not existed db
   198  	dbName := common.GenRandomString(4)
   199  	err := mc.UsingDatabase(ctx, dbName)
   200  	common.CheckErr(t, err, false, fmt.Sprintf("connect fail, database not found[database=%s]", dbName))
   201  
   202  	// using empty db
   203  	err = mc.UsingDatabase(ctx, "")
   204  	common.CheckErr(t, err, true)
   205  	collections, _ = mc.ListCollections(ctx)
   206  	common.CheckContainsCollection(t, collections, collName)
   207  
   208  	// using current db
   209  	err = mc.UsingDatabase(ctx, common.DefaultDb)
   210  	common.CheckErr(t, err, true)
   211  	collections, _ = mc.ListCollections(ctx)
   212  	common.CheckContainsCollection(t, collections, collName)
   213  }
   214  
   215  // test client with db
   216  func TestClientWithDb(t *testing.T) {
   217  	teardownSuite := teardownTest(t)
   218  	defer teardownSuite(t)
   219  	ctx := createContext(t, time.Second*common.DefaultTimeout)
   220  
   221  	// connect with not existed db
   222  	_, err := base.NewMilvusClient(ctx, client.Config{
   223  		Address: *addr,
   224  		DBName:  common.GenRandomString(4),
   225  	})
   226  	common.CheckErr(t, err, false, "database not found")
   227  
   228  	// connect default db -> create a collection in default db
   229  	mcDefault, errDefault := base.NewMilvusClient(ctx, client.Config{
   230  		Address: *addr,
   231  		DBName:  common.DefaultDb,
   232  	})
   233  	common.CheckErr(t, errDefault, true)
   234  	_, _ = mcDefault.ListCollections(ctx)
   235  	collDefault := createDefaultCollection(ctx, t, mcDefault, true, common.DefaultShards)
   236  
   237  	// create a db and create collection in db
   238  	dbName := common.GenRandomString(5)
   239  	err = mcDefault.CreateDatabase(ctx, dbName)
   240  	common.CheckErr(t, err, true)
   241  
   242  	// and connect with db
   243  	mcDb, err := base.NewMilvusClient(ctx, client.Config{
   244  		Address: *addr,
   245  		DBName:  dbName,
   246  	})
   247  	common.CheckErr(t, err, true)
   248  	collDb := createDefaultCollection(ctx, t, mcDb, true, common.DefaultShards)
   249  	collections, _ := mcDb.ListCollections(ctx)
   250  	common.CheckContainsCollection(t, collections, collDb)
   251  
   252  	// using default db and collection not in
   253  	_ = mcDb.UsingDatabase(ctx, common.DefaultDb)
   254  	collections, _ = mcDb.ListCollections(ctx)
   255  	common.CheckNotContainsCollection(t, collections, collDb)
   256  
   257  	// connect empty db (actually default db)
   258  	mcEmpty, err := base.NewMilvusClient(ctx, client.Config{
   259  		Address: *addr,
   260  		DBName:  "",
   261  	})
   262  	common.CheckErr(t, err, true)
   263  	collections, _ = mcEmpty.ListCollections(ctx)
   264  	common.CheckContainsCollection(t, collections, collDefault)
   265  }
   266  
   267  func TestAlterDatabase(t *testing.T) {
   268  	teardownSuite := teardownTest(t)
   269  	defer teardownSuite(t)
   270  
   271  	// create db
   272  	ctx := createContext(t, time.Second*common.DefaultTimeout)
   273  	mc := createMilvusClient(ctx, t)
   274  	dbName := common.GenRandomString(4)
   275  	err := mc.CreateDatabase(ctx, dbName)
   276  	common.CheckErr(t, err, true)
   277  
   278  	dbInfo, err := mc.DescribeDatabase(ctx, dbName)
   279  	common.CheckErr(t, err, true)
   280  	assert.Equal(t, dbInfo.Name, dbName)
   281  
   282  	err = mc.AlterDatabase(ctx, dbName, entity.DatabaseResourceGroups([]string{"rg1"}))
   283  	common.CheckErr(t, err, true)
   284  
   285  	dbInfo, err = mc.DescribeDatabase(ctx, dbName)
   286  	common.CheckErr(t, err, true)
   287  	assert.Equal(t, dbInfo.Name, dbName)
   288  	assert.Len(t, dbInfo.Properties, 1)
   289  
   290  	err = mc.AlterDatabase(ctx, dbName, entity.DatabaseReplica(1))
   291  	common.CheckErr(t, err, true)
   292  
   293  	dbInfo, err = mc.DescribeDatabase(ctx, dbName)
   294  	common.CheckErr(t, err, true)
   295  	assert.Equal(t, dbInfo.Name, dbName)
   296  	assert.Len(t, dbInfo.Properties, 2)
   297  }